import subprocess
import json
import time
import random
from typing import Dict, Tuple, List, Optional, Any

# In a real implementation, you would use an actual LLM API
# Example with OpenAI (you would need to install the openai package)
import openai

class FunSearch:
    """
    A complete implementation of the FunSearch approach for code generation
    tasks, using an evaluation-based feedback loop.
    """
    
    def __init__(
        self, 
        api_key: str, 
        model: str = "gpt-4", 
        max_iterations: int = 5,
        temperature: float = 0.7,
        timeout_seconds: int = 10
    ):
        """
        Initialize the FunSearch system.
        
        Args:
            api_key: API key for the LLM service
            model: Model identifier to use
            max_iterations: Maximum number of refinement attempts
            temperature: Creativity parameter for generation (0.0-1.0)
            timeout_seconds: Maximum execution time for candidate solutions
        """
        openai.api_key = api_key
        self.model = model
        self.max_iterations = max_iterations
        self.temperature = temperature
        self.timeout_seconds = timeout_seconds
        self.history = []
        
    def generate_solution(self, system_prompt: str, user_prompt: str) -> str:
        """
        Generate a candidate solution using the configured LLM.
        
        Args:
            system_prompt: Instructions for the model's role and task
            user_prompt: The specific problem description and requirements
            
        Returns:
            Generated code as a string
        """
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
        
        # Add conversation history for context in later iterations
        for entry in self.history:
            messages.append({"role": "assistant", "content": entry["solution"]})
            messages.append({"role": "user", "content": entry["feedback"]})
        
        response = openai.ChatCompletion.create(
            model=self.model,
            messages=messages,
            temperature=self.temperature,
            max_tokens=2000
        )
        
        return response.choices[0].message.content
    
    def evaluate_python_solution(
        self, 
        code: str, 
        test_cases: List[Dict[str, Any]]
    ) -> Tuple[bool, str, Dict]:
        """
        Evaluate a Python solution against multiple test cases.
        
        Args:
            code: The Python code to evaluate
            test_cases: List of test cases with inputs and expected outputs
            
        Returns:
            Tuple containing:
            - Success flag (True/False)
            - Feedback message
            - Detailed results including performance metrics
        """
        # Save the code to a temporary file
        with open("candidate_solution.py", "w") as f:
            f.write(code)
        
        results = {
            "passed": 0,
            "failed": 0,
            "errors": [],
            "execution_times": []
        }
        
        # Run each test case
        for i, test in enumerate(test_cases):
            test_input = test.get("input", "")
            expected_output = test.get("expected", "")
            
            try:
                # Time the execution
                start_time = time.time()
                
                process = subprocess.Popen(
                    ["python", "candidate_solution.py"],
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    text=True
                )
                
                # Provide input and capture output
                stdout, stderr = process.communicate(
                    input=str(test_input), 
                    timeout=self.timeout_seconds
                )
                
                execution_time = time.time() - start_time
                results["execution_times"].append(execution_time)
                
                # Check for errors
                if stderr:
                    results["failed"] += 1
                    results["errors"].append({
                        "test_case": i,
                        "type": "Runtime Error",
                        "message": stderr
                    })
                    continue
                    
                # Compare with expected output
                if stdout.strip() == str(expected_output).strip():
                    results["passed"] += 1
                else:
                    results["failed"] += 1
                    results["errors"].append({
                        "test_case": i,
                        "type": "Output Mismatch",
                        "expected": expected_output,
                        "actual": stdout.strip()
                    })
                    
            except subprocess.TimeoutExpired:
                results["failed"] += 1
                results["errors"].append({
                    "test_case": i,
                    "type": "Timeout",
                    "message": f"Execution exceeded {self.timeout_seconds} seconds"
                })
        
        # Determine overall success
        success = results["failed"] == 0 and results["passed"] == len(test_cases)
        
        # Generate human-readable feedback
        if success:
            feedback = f"Success! All {results['passed']} test cases passed."
            if results["execution_times"]:
                avg_time = sum(results["execution_times"]) / len(results["execution_times"])
                feedback += f" Average execution time: {avg_time:.4f} seconds."
        else:
            feedback = f"Failed {results['failed']} out of {len(test_cases)} test cases.\n"
            for error in results["errors"]:
                feedback += f"- Test case {error['test_case']}: {error['type']}\n"
                if error['type'] == "Output Mismatch":
                    feedback += f"  Expected: {error['expected']}\n"
                    feedback += f"  Actual: {error['actual']}\n"
                elif 'message' in error:
                    feedback += f"  {error['message']}\n"
        
        return success, feedback, results
    
    def search(
        self, 
        problem_description: str, 
        test_cases: List[Dict[str, Any]]
    ) -> Dict:
        """
        Execute the full FunSearch process to find a solution.
        
        Args:
            problem_description: Detailed description of the problem to solve
            test_cases: List of test cases for evaluation
            
        Returns:
            Dictionary with the best solution and full history
        """
        system_prompt = (
            "You are an expert Python programmer. Your task is to write correct, "
            "efficient, and clean Python code that solves the given problem. "
            "Focus on correctness first, then efficiency."
        )
        
        best_solution = None
        best_results = None
        
        print(f"Starting FunSearch with {self.max_iterations} maximum iterations")
        
        for iteration in range(self.max_iterations):
            print(f"\n--- Iteration {iteration+1}/{self.max_iterations} ---")
            
            # Generate a solution
            code = self.generate_solution(system_prompt, problem_description)
            print(f"\nGenerated solution:\n{'-'*20}\n{code}\n{'-'*20}")
            
            # Evaluate the solution
            success, feedback, results = self.evaluate_python_solution(code, test_cases)
            print(f"\nEvaluation feedback: {feedback}")
            
            # Record this attempt
            attempt = {
                "iteration": iteration + 1,
                "solution": code,
                "success": success,
                "feedback": feedback,
                "results": results
            }
            self.history.append(attempt)
            
            # Check if we found a successful solution
            if success:
                print("\n✅ Found successful solution!")
                best_solution = code
                best_results = results
                break
            
            # Update the problem description with feedback for the next iteration
            problem_description = (
                f"{problem_description}\n\n"
                f"My previous attempt had issues:\n{feedback}\n\n"
                "Please fix these issues and try again."
            )
        
        return {
            "success": best_solution is not None,
            "best_solution": best_solution,
            "best_results": best_results,
            "iterations_used": len(self.history),
            "history": self.history
        }


# Example usage:
if __name__ == "__main__":
    # Replace with your actual API key
    API_KEY = "your_openai_api_key_here"
    
    funsearch = FunSearch(
        api_key=API_KEY,
        model="gpt-4",
        max_iterations=3,
        temperature=0.7
    )
    
    problem = """
    Write a Python program that reads a number N from stdin and calculates its factorial.
    
    Requirements:
    1. The program should handle positive integers up to 20
    2. For inputs that aren't positive integers, print "Invalid input"
    3. The output should only contain the resulting factorial value, nothing else
    4. Do not use any external libraries beyond Python's built-ins
    """
    
    test_cases = [
        {"input": "5", "expected": "120"},
        {"input": "0", "expected": "1"},
        {"input": "20", "expected": "2432902008176640000"},
        {"input": "-3", "expected": "Invalid input"},
        {"input": "abc", "expected": "Invalid input"}
    ]
    
    result = funsearch.search(problem, test_cases)
    
    # Report final results
    if result["success"]:
        print("\n=== SUCCESSFUL SOLUTION FOUND ===")
        print(f"Required {result['iterations_used']} iterations")
        print("\nFinal solution:")
        print(result["best_solution"])
    else:
        print("\n=== NO SUCCESSFUL SOLUTION FOUND ===")
        print("Best attempt:")
        best_attempt = max(result["history"], key=lambda x: x["results"]["passed"])
        print(f"Passed {best_attempt['results']['passed']} of {len(test_cases)} test cases")