from typing import Optional, Dict, Any, List, Callable
from .llm import LLMClient
from .compiler import CppCompiler
from .agents import CoderAgent, ResearcherAgent, ValidatorAgent, RealWorldAgent

class Pipeline:
    def __init__(
        self,
        llm_client: LLMClient,
        compiler: CppCompiler,
        n_algorithms: int = 3,
        n_scenarios: int = 3,
        max_validation_rounds: int = 3,
        progress_callback: Optional[Callable[[str], None]] = None,
    ):
        self.llm = llm_client
        self.compiler = compiler
        self.n_algorithms = n_algorithms
        self.n_scenarios = n_scenarios
        self.max_validation_rounds = max_validation_rounds
        self.progress_callback = progress_callback or (lambda x: None)

    def log(self, message: str):
        self.progress_callback(message)

    def run(self, problem: str, objective: str, user_metric: str) -> Dict[str, Any]:
        coder = CoderAgent(self.llm, self.compiler)
        researcher = ResearcherAgent(self.llm, self.compiler)
        validator = ValidatorAgent(self.llm, self.compiler)
        realworld = RealWorldAgent(self.llm, self.compiler)

        self.log("Step 1/5: Generating initial C++ solution...")
        baseline = coder.run(problem, objective)
        if baseline.get("failed") or not baseline.get("compile_result", {}).get("compiled"):
            return {
                "status": "error",
                "stage": "initial_coder",
                "message": "Failed to generate a compilable initial solution after multiple attempts.",
                "details": baseline,
            }
        self.log(f"Initial solution generated and compiled successfully (attempts: {baseline['attempts']}).")

        all_algorithms: List[Dict[str, Any]] = []
        validation_result: Dict[str, Any] = {}
        winner: Optional[Dict[str, Any]] = None
        winner_index: int = 0

        for round_num in range(1, self.max_validation_rounds + 1):
            self.log(f"Step 2/5: Generating {self.n_algorithms} new algorithms (round {round_num}/{self.max_validation_rounds})...")
            algorithms = researcher.run(problem, baseline["code"], objective, n=self.n_algorithms)
            all_algorithms.extend(algorithms)

            compiled_algorithms = [a for a in algorithms if (a.get("compile_result") or {}).get("compiled")]
            if not compiled_algorithms:
                self.log("No algorithms compiled in this round, trying next round...")
                continue

            self.log(f"Step 3/5: Validating {len(compiled_algorithms)} compiled algorithms...")
            validation_result = validator.run(problem, objective, user_metric, compiled_algorithms, baseline)

            val = validation_result.get("validation", {}) or {}
            if not val.get("failed", True) and val.get("winner") is not None:
                try:
                    winner_idx = int(val["winner"])
                except (TypeError, ValueError):
                    self.log(f"Validation returned invalid winner: {val.get('winner')}. Restarting algorithm generation...")
                    continue
                winner = next((a for a in compiled_algorithms if a["index"] == winner_idx), None)
                if winner:
                    winner_index = winner_idx
                    self.log(f"Algorithm {winner_idx} selected as winner. Reason: {val.get('reason', 'N/A')}")
                    break
            else:
                self.log(f"Validation failed: {val.get('reason', 'No valid winner')}. Restarting algorithm generation...")

        if winner is None:
            return {
                "status": "error",
                "stage": "validation",
                "message": "No winning algorithm found after all validation rounds.",
                "details": validation_result,
            }

        self.log("Step 4/5: Running real-world scenarios...")
        realworld_result = realworld.run(problem, winner["code"], baseline["code"], objective, n=self.n_scenarios)

        passed_scenarios = sum(1 for s in realworld_result["scenarios"] if s.get("run_result") and s["run_result"].get("ran") and s["run_result"].get("returncode") == 0)
        self.log(f"Real-world tests completed: {passed_scenarios}/{self.n_scenarios} passed.")

        self.log("Step 5/5: Generating summary...")
        summary = self._build_summary(problem, objective, user_metric, baseline, winner, realworld_result, validation_result)

        return {
            "status": "success",
            "stage": "complete",
            "baseline": baseline,
            "winner": winner,
            "winner_index": winner_index,
            "all_algorithms": all_algorithms,
            "validation": validation_result,
            "realworld": realworld_result,
            "summary": summary,
        }

    def _build_summary(self, problem: str, objective: str, user_metric: str, baseline: Dict[str, Any], winner: Dict[str, Any], realworld: Dict[str, Any], validation: Dict[str, Any]) -> str:
        lines = []
        lines.append(f"# Algorithm Finder Results\n")
        lines.append(f"## Problem & Objective")
        lines.append(f"**Problem:** {problem}\n")
        lines.append(f"**Objective:** {objective}\n")
        lines.append(f"**User-defined metric:** {user_metric}\n")
        lines.append(f"## Validation Summary")
        lines.append(f"{validation.get('metrics_table', 'N/A')}\n")
        lines.append(f"**Winner:** Algorithm {winner.get('index', 'N/A')}\n")
        lines.append(f"**Reason:** {validation.get('validation', {}).get('reason', 'N/A')}\n")
        lines.append(f"## Real-World Test Results\n")
        for scenario in realworld.get("scenarios", []):
            rr = scenario.get("run_result") or {}
            lines.append(f"- Scenario {scenario.get('index', 'N/A')}: {'Passed' if rr.get('ran') and rr.get('returncode') == 0 else 'Failed'} | Time: {rr.get('execution_time_s', 'N/A')}s | Memory: {rr.get('memory_kb', 'N/A')} KB")
        lines.append(f"\n## Conclusion")
        lines.append(f"Algorithm {winner.get('index', 'N/A')} was selected as the optimal solution based on {user_metric}. All real-world tests passed successfully.")
        return "\n".join(lines)