import asyncio
import importlib.util
import json
import logging
import os
import re
import sys
import tempfile
import time
import traceback
from typing import Any, Dict, List, Optional, Tuple, Union
import errno
from skydiscover.config import EvaluatorConfig
from skydiscover.search.base_database import ProgramDatabase
from skydiscover.evaluation.evaluation_result import EvaluationResult
from skydiscover.llm.llm_pool import LLMPool
from skydiscover.utils.async_utils import TaskPool
from skydiscover.prompt.default import DefaultContextBuilder
from skydiscover.utils.format_utils import format_metrics_safe

logger = logging.getLogger(__name__)


class Evaluator:
    """
    Evaluates programs by executing them and assigning scores.
    """

    def __init__(
        self,
        config: EvaluatorConfig,
        evaluation_file: str,
        llm_pool: Optional[LLMPool] = None,
        prompt_sampler: Optional[DefaultContextBuilder] = None,
        database: Optional[ProgramDatabase] = None,
        suffix: Optional[str] = ".py",
        is_image_mode: bool = False,
    ):
        self.config = config
        self.evaluation_file = evaluation_file
        self.program_suffix = suffix
        self.is_image_mode = is_image_mode

        # LLM pool if using llm-as-a-judge
        self.llm_pool = llm_pool
        self.prompt_sampler = prompt_sampler
        self.database = database

        # Allow multiple evaluations to run in parallel
        self.task_pool = TaskPool(max_concurrency=config.parallel_evaluations)
        self._load_evaluation_function()

        logger.info(f"Initialized evaluator with {evaluation_file}")
        if config.use_llm_feedback and self.llm_pool:
            model_names = ", ".join(m.name for m in self.llm_pool.models_cfg)
            logger.info(f"LLM feedback enabled (weight={config.llm_feedback_weight}) using models: {model_names}")

    def _load_evaluation_function(self) -> None:
        if not os.path.exists(self.evaluation_file):
            raise ValueError(f"Evaluation file {self.evaluation_file} not found")

        try:
            eval_dir = os.path.dirname(os.path.abspath(self.evaluation_file))
            if eval_dir not in sys.path:
                sys.path.insert(0, eval_dir)

            spec = importlib.util.spec_from_file_location("evaluation_module", self.evaluation_file)
            if spec is None or spec.loader is None:
                raise ImportError(f"Failed to load spec from {self.evaluation_file}")

            module = importlib.util.module_from_spec(spec)
            sys.modules["evaluation_module"] = module
            spec.loader.exec_module(module)

            if not hasattr(module, "evaluate"):
                raise AttributeError(
                    f"Evaluation file {self.evaluation_file} does not contain an 'evaluate' function"
                )

            self.evaluate_function = module.evaluate
            self._eval_module = module
            self._validate_cascade_configuration(module)
        except Exception as e:
            logger.error(f"Error loading evaluation function: {str(e)}")
            raise

    def _validate_cascade_configuration(self, module) -> None:
        """Warn if cascade_evaluation is enabled but the module lacks stage functions."""
        if self.config.cascade_evaluation:
            has_stage1 = hasattr(module, "evaluate_stage1")
            has_stage2 = hasattr(module, "evaluate_stage2")

            if not has_stage1:
                logger.warning(
                    f"Configuration has 'cascade_evaluation: true' but evaluator "
                    f"'{self.evaluation_file}' does not define 'evaluate_stage1' function. "
                    f"This will fall back to direct evaluation, making the cascade setting useless. "
                    f"Consider setting 'cascade_evaluation: false' or implementing cascade functions."
                )
            elif not has_stage2:
                logger.warning(
                    f"Evaluator '{self.evaluation_file}' defines 'evaluate_stage1' but no "
                    f"'evaluate_stage2'. Consider implementing multi-stage evaluation for "
                    f"better cascade benefits."
                )
            else:
                logger.debug(
                    "Cascade evaluation properly configured with available stage functions"
                )

    async def evaluate_program(
        self,
        program_solution: str,
        program_id: str = "",
    ) -> EvaluationResult:
        """
        Evaluate a program and return scores with optional artifacts.

        Args:
            program_solution: Code to evaluate, or an image file path when
                ``is_image_mode`` is True.
            program_id: Optional ID for logging

        Returns:
            EvaluationResult with metrics dict and artifacts dict
        """
        start_time = time.time()
        program_id_str = f" {program_id}" if program_id else ""
        logging.info(f"Evaluating program{program_id_str}")

        last_exception = None
        for attempt in range(self.config.max_retries + 1):
            try:
                with tempfile.NamedTemporaryFile(suffix=self.program_suffix, delete=False) as temp_file:
                    temp_file.write(program_solution.encode("utf-8"))
                    temp_file_path = temp_file.name
            except OSError as e:
                if e.errno == errno.ENOSPC:
                    logger.error("No space left on device in /tmp. Cannot create temporary file.")
                    return EvaluationResult(metrics={"error": 0.0, "disk_space_error": True})
                raise

            # Write sidecar file for image mode so evaluator can find the image
            # In image mode, program_solution is the image file path
            sidecar_path = None
            if self.is_image_mode:
                sidecar_path = temp_file_path + ".image_path"
                try:
                    with open(sidecar_path, "w") as sf:
                        sf.write(program_solution)
                except Exception as e:
                    logger.warning(f"Failed to write image sidecar: {e}")

            try:
                # Run evaluation
                if self.config.cascade_evaluation:
                    logging.info("Running cascade evaluation")
                    result = await self._cascade_evaluate(temp_file_path)
                else:
                    result = await self._direct_evaluate(temp_file_path)

                eval_result = self._process_evaluation_result(result)

                # LLM feedback if configured
                if self.config.use_llm_feedback and self.llm_pool:
                    llm_eval_result = await self._llm_evaluate(program_solution, program_id=program_id)
                    if llm_eval_result:
                        llm_eval_result = self._process_evaluation_result(llm_eval_result)
                        llm_scores = []
                        for name, value in llm_eval_result.metrics.items():
                            if isinstance(value, (int, float)):
                                eval_result.metrics[f"llm_{name}"] = float(value)
                                llm_scores.append(float(value))

                        # Blend into combined_score: (1 - weight) * evaluator + weight * llm_average
                        if llm_scores:
                            w = self.config.llm_feedback_weight
                            llm_avg = sum(llm_scores) / len(llm_scores)
                            eval_result.metrics["llm_average"] = llm_avg
                            if "combined_score" in eval_result.metrics:
                                eval_score = eval_result.metrics["combined_score"]
                                eval_result.metrics["combined_score"] = (1 - w) * eval_score + w * llm_avg

                elapsed = time.time() - start_time
                logger.info(
                    f"Evaluated program{program_id_str} in {elapsed:.2f}s: "
                    f"{format_metrics_safe(eval_result.metrics)}"
                )

                return eval_result

            except asyncio.TimeoutError:
                elapsed = time.time() - start_time
                logger.error(
                    f"Program{program_id_str} TIMED OUT after {elapsed:.2f}s "
                    f"(timeout limit: {self.config.timeout}s). "
                    f"The generated code took too long to execute — likely an expensive "
                    f"solver/optimizer that did not converge within the time limit."
                )

                return EvaluationResult(metrics={"error": 0.0, "timeout": True})

            except Exception as e:
                last_exception = e
                logger.warning(
                    f"Evaluation attempt {attempt + 1}/{self.config.max_retries + 1} failed for program{program_id_str}: {str(e)}"
                )
                traceback.print_exc()

                if attempt < self.config.max_retries:
                    await asyncio.sleep(1.0)

            finally:
                # Clean up temporary file and sidecar
                if os.path.exists(temp_file_path):
                    os.unlink(temp_file_path)
                if sidecar_path and os.path.exists(sidecar_path):
                    os.unlink(sidecar_path)

        logger.error(
            f"All evaluation attempts failed for program{program_id_str}. Last error: {str(last_exception)}"
        )
        return EvaluationResult(metrics={"error": 0.0})

    async def evaluate_multiple(
        self,
        programs: List[Tuple[str, str]],
    ) -> List[Dict[str, float]]:
        """Evaluate multiple programs in parallel."""
        tasks = [
            self.task_pool.create_task(self.evaluate_program, program_solution, program_id)
            for program_solution, program_id in programs
        ]
        return await asyncio.gather(*tasks)

    def _process_evaluation_result(self, result: Any) -> EvaluationResult:
        """Normalize a dict or EvaluationResult into an EvaluationResult."""
        if isinstance(result, dict):
            return EvaluationResult.from_dict(result)
        elif isinstance(result, EvaluationResult):
            return result
        else:
            logger.warning(f"Unexpected evaluation result type: {type(result)}")
            return EvaluationResult(metrics={"error": 0.0})

    async def _direct_evaluate(
        self, program_path: str
    ) -> Union[Dict[str, float], EvaluationResult]:
        """Run the evaluation function in an executor with timeout."""

        async def run_evaluation():
            loop = asyncio.get_event_loop()
            return await loop.run_in_executor(None, self.evaluate_function, program_path)

        return await asyncio.wait_for(run_evaluation(), timeout=self.config.timeout)

    async def _cascade_evaluate(
        self, program_path: str
    ) -> Union[Dict[str, float], EvaluationResult]:
        """Run cascade evaluation with increasingly challenging test cases."""
        try:
            module = self._eval_module

            if not hasattr(module, "evaluate_stage1"):
                return await self._direct_evaluate(program_path)

            # Run stage 1 evaluation
            try:
                async def run_stage1():
                    loop = asyncio.get_event_loop()
                    return await loop.run_in_executor(None, module.evaluate_stage1, program_path)

                stage1_result = await asyncio.wait_for(run_stage1(), timeout=self.config.timeout)
                stage1_eval_result = self._process_evaluation_result(stage1_result)

            except asyncio.TimeoutError:
                logger.error(f"Stage 1 evaluation TIMED OUT after {self.config.timeout}s — generated code too slow")
                return EvaluationResult(
                    metrics={"stage1_passed": 0.0, "error": 0.0, "timeout": True},
                    artifacts={
                        "failure_stage": "stage1",
                        "timeout": True,
                    },
                )
            except Exception as e:
                logger.error(f"Error in stage 1 evaluation: {type(e).__name__}: {str(e)}")
                error_context = self._create_cascade_error_context("stage1", e)
                return EvaluationResult(
                    metrics={"stage1_passed": 0.0, "error": 0.0},
                    artifacts={
                        "stderr": str(e),
                        "traceback": traceback.format_exc(),
                        **error_context,
                    },
                )

            if not self._passes_threshold(
                stage1_eval_result.metrics, self.config.cascade_thresholds[0]
            ):
                return stage1_eval_result

            if not hasattr(module, "evaluate_stage2"):
                return stage1_eval_result

            # Run stage 2 evaluation
            try:
                async def run_stage2():
                    loop = asyncio.get_event_loop()
                    return await loop.run_in_executor(None, module.evaluate_stage2, program_path)

                stage2_result = await asyncio.wait_for(run_stage2(), timeout=self.config.timeout)
                stage2_eval_result = self._process_evaluation_result(stage2_result)

            except asyncio.TimeoutError:
                logger.error(f"Stage 2 evaluation TIMED OUT after {self.config.timeout}s — generated code too slow")
                stage1_eval_result.artifacts.update(
                    {
                        "stage2_timeout": True,
                        "failure_stage": "stage2",
                    }
                )
                stage1_eval_result.metrics["stage2_passed"] = 0.0
                stage1_eval_result.metrics["timeout"] = True
                return stage1_eval_result
            except Exception as e:
                logger.error(f"Error in stage 2 evaluation: {type(e).__name__}: {str(e)}")
                stage1_eval_result.artifacts.update(
                    {
                        "stage2_stderr": str(e),
                        "stage2_traceback": traceback.format_exc(),
                        "failure_stage": "stage2",
                    }
                )
                stage1_eval_result.metrics["stage2_passed"] = 0.0
                return stage1_eval_result

            merged_metrics = {}
            for name, value in stage1_eval_result.metrics.items():
                if isinstance(value, (int, float)) and name != "error":
                    merged_metrics[name] = float(value)

            for name, value in stage2_eval_result.metrics.items():
                if isinstance(value, (int, float)) and name != "error":
                    merged_metrics[name] = float(value)

            merged_artifacts = {**stage1_eval_result.artifacts, **stage2_eval_result.artifacts}

            return EvaluationResult(metrics=merged_metrics, artifacts=merged_artifacts)

        except asyncio.TimeoutError:
            raise
        except Exception as e:
            logger.error(f"Error in cascade evaluation: {type(e).__name__}: {str(e)}")
            error_context = self._create_cascade_error_context("cascade_setup", e)
            return EvaluationResult(
                metrics={"stage1_passed": 0.0, "error": 0.0},
                artifacts={
                    "stderr": str(e),
                    "traceback": traceback.format_exc(),
                    **error_context,
                },
            )

    async def _llm_evaluate(self, program_solution: str, program_id: str = "") -> Optional[Union[Dict[str, float], EvaluationResult]]:
        """Use LLM to evaluate code quality and optionally log to database."""
        try:
            tm = self.prompt_sampler.template_manager
            eval_sys = self.prompt_sampler.config.evaluator_system_message
            system_msg = tm.get_template(eval_sys) if eval_sys in tm.templates else eval_sys
            user_msg = tm.get_template("evaluation").format(current_program=program_solution)
            logger.info(f"LLM judge system message: {system_msg}")
            logger.info(f"LLM judge user message: {user_msg[:200]}...\n")
            prompt = {"system": system_msg, "user": user_msg}
            responses = await self.llm_pool.generate_all(
                system_msg, [{"role": "user", "content": user_msg}]
            )

            if self.database and program_id:
                self.database.log_prompt(
                    program_id=program_id,
                    template_key="evaluation",
                    prompt=prompt,
                    responses=responses,
                )

            avg_metrics = {}
            artifacts = {}
            for i, response in enumerate(responses):
                json_match = re.search(r"```json\n(.*?)\n```", response, re.DOTALL)
                json_str = json_match.group(1) if json_match else response
                if not json_match:
                    start, end = json_str.find("{"), json_str.rfind("}") + 1
                    if start >= 0 and end > start:
                        json_str = json_str[start:end]

                parsed = json.loads(json_str)
                weight = self.llm_pool.weights[i] if self.llm_pool.weights else 1.0
                for key, value in parsed.items():
                    if isinstance(value, (int, float)):
                        avg_metrics[key] = avg_metrics.get(key, 0.0) + float(value) * weight
                    else:
                        artifacts[key] = value

            return EvaluationResult(metrics=avg_metrics, artifacts=artifacts)
        except Exception as e:
            logger.warning(f"LLM evaluation failed: {e}")
            return None

    def _create_cascade_error_context(self, stage: str, error: Exception) -> dict:
        return {
            "failure_stage": stage,
            "error_type": type(error).__name__,
            "error_message": str(error),
        }

    def _passes_threshold(self, metrics: Dict[str, float], threshold: float) -> bool:
        """Uses combined_score if available, otherwise averages all numeric metrics except 'error'."""
        if not metrics:
            return False

        if "combined_score" in metrics:
            score = metrics["combined_score"]
            if isinstance(score, (int, float)):
                return float(score) >= threshold

        valid_metrics = [
            float(v) for k, v in metrics.items()
            if k != "error" and isinstance(v, (int, float))
        ]

        if not valid_metrics:
            return False

        return (sum(valid_metrics) / len(valid_metrics)) >= threshold