from collections import defaultdict
from pathlib import Path
from uuid import UUID

import yaml
from pydantic import BaseModel

import railtracks as rt
from railtracks.utils.logging import get_rt_logger

from ..point import AgentDataPoint
from ..result import (
    AggregateForest,
    CategoricalAggregateNode,
    EvaluatorResult,
    MetricResult,
)
from .evaluator import Evaluator
from .metrics import Categorical, Metric, Numerical

logger = get_rt_logger(__name__)


class JudgeResponseSchema(BaseModel):
    metric_value: str | float | int
    reasoning: str | None = None


class JudgeEvaluator(Evaluator):
    def __init__(
        self,
        llm: rt.llm.ModelBase,
        metrics: list[Metric],
        system_prompt: str | None = None,
        timeout: float | None = None,
        reasoning: bool = True,
    ):
        """
        The JudgeEvaluator with a system prompt, LLM, metric, and reasoning flag.

        Args:
            system_prompt: The system prompt template for the judge LLM.
            llm: The LLM model to be used as the judge.
            metrics: A list of Metrics to guide the evaluation.
            reasoning: A flag indicating whether the judge should provide reasoning for its evaluations.
        """
        # These are config not state
        self._metrics: dict[str, Metric] = {m.identifier: m for m in metrics}
        for m in metrics:
            if isinstance(m, Categorical):
                self._metrics[m.identifier] = m
            else:
                logger.warning(
                    f"JudgeEvaluator currently only supports Categorical metrics, metric {m.name} of type {type(m)} will be skipped."
                )
        self._llm = llm
        self._reasoning: bool = reasoning
        self._template = self._load_yaml()
        self._system_prompt = (
            system_prompt
            if system_prompt is not None
            else self._template["system_prompt"]
        )
        super().__init__()

        self.timeout = timeout
        self._judge = rt.agent_node(
            llm=self._llm,
            output_schema=JudgeResponseSchema,
            tool_nodes=[],
        )

    def run(
        self, data: list[AgentDataPoint]
    ) -> EvaluatorResult[Metric, MetricResult, CategoricalAggregateNode]:
        # (metric_id, adp_id, JudgeResponseSchema)
        judge_outputs: list[tuple[str, str, JudgeResponseSchema]] = self._invoke(data)

        self.agent_data_ids = {adp.identifier for adp in data}
        results: dict[Metric, list[MetricResult]] = defaultdict(list)
        forest = AggregateForest[CategoricalAggregateNode, MetricResult]()

        for output in judge_outputs:
            metric = self._metrics[output[0]]

            metric_result = MetricResult(
                result_name=f"JudgeResult/{metric.name}",
                metric_id=metric.identifier,
                agent_data_id=[UUID(output[1])],
                value=output[2].metric_value,
            )
            results[metric].append(metric_result)
            forest.add_node(metric_result)

            if self._reasoning:
                reasoning_metric = Metric(name=f"{metric.name}_reasoning")
                if output[2].reasoning is not None:
                    results[reasoning_metric].append(
                        MetricResult(
                            result_name=f"JudgeReasoning/{metric.name}",
                            metric_id=reasoning_metric.identifier,
                            agent_data_id=[UUID(output[1])],
                            value=output[2].reasoning,
                        )
                    )
                else:
                    logger.warning(
                        f"No reasoning returned for Judge Evaluator Metric: {metric.name}, AgentDataPoint ID: {output[1]}"
                    )

        self._aggregate_metrics(results, forest)

        self._result = EvaluatorResult(
            evaluator_name=self.name,
            evaluator_id=self.identifier,
            agent_data_ids=self.agent_data_ids,
            metric_results=[item for sublist in results.values() for item in sublist],
            aggregate_results=forest,
            metrics=list(self._metrics.values()),
        )
        return self._result

    def __repr__(self) -> str:
        return (
            f"JudgeEvaluator, "
            f"llm={self._llm}, "
            f"metrics={list(self._metrics.values())}, "
            f"reasoning={self._reasoning})"
        )

    def _invoke(
        self, data: list[AgentDataPoint]
    ) -> list[tuple[str, str, JudgeResponseSchema]]:
        @rt.function_node
        async def judge_flow():
            output = []
            for metric in self._metrics.values():
                logger.info(
                    f"START Evaluating Metric: {metric.name} for {len(data)} AgentDataPoints"
                )

                for idx, adp in enumerate(data):
                    user_message = self._generate_user_prompt(adp)
                    system_message = self._generate_system_prompt(metric)
                    message_history = rt.llm.MessageHistory(
                        [
                            rt.llm.SystemMessage(system_message),
                            rt.llm.UserMessage(user_message),
                        ]
                    )
                    res = await rt.call(
                        self._judge,
                        message_history,
                    )
                    output.append(
                        (metric.identifier, str(adp.identifier), res.structured)
                    )

                    logger.info(
                        f"AgentDataPoint ID: {adp.identifier} {idx + 1}/{len(data)} DONE"
                    )

            return output

        judge_evaluator_flow = rt.Flow(
            name="JudgeEvaluatorFlow",
            entry_point=judge_flow,
            timeout=self.timeout,
            save_state=False,
        )

        return judge_evaluator_flow.invoke()

    def _aggregate_metrics(
        self,
        results: dict[Metric, list[MetricResult]],
        forest: AggregateForest[CategoricalAggregateNode, MetricResult],
    ) -> None:
        for metric in results:
            if isinstance(metric, Numerical):
                continue
            elif isinstance(metric, Categorical):
                aggregate_node = CategoricalAggregateNode(
                    name=f"Aggregate/{metric.name}",
                    metric=metric,
                    children=[val.identifier for val in results[metric]],
                    forest=forest,
                )

                forest.roots.append(aggregate_node.identifier)
                forest.add_node(aggregate_node)

    def _generate_user_prompt(self, data: AgentDataPoint) -> str:
        return self._template["user"].format(
            agent_input=data.agent_input,
            agent_output=data.agent_output.get("message_history", ""),
        )

    def _generate_system_prompt(self, metric: Metric) -> str:
        system_prompt: str = self._template["system_prompt"]

        system_prompt += "\n" + self._template["metric"].format(metric=str(metric))

        if self._reasoning:
            system_prompt += self._template["reasoning"]

        return system_prompt

    def _load_yaml(self):
        yaml_path = Path(__file__).parent / "judge_evaluator.yaml"
        with open(yaml_path, "r") as f:
            template = yaml.safe_load(f)

        return template

    def _get_config(self) -> dict:
        return {
            "llm": self._llm.model_name(),
            "llm_provider": self._llm.model_provider(),
            "system_prompt": self._system_prompt,
            "metrics": sorted(self._metrics.keys()),
            "reasoning": self._reasoning,
        }
