import threading
from typing import Any

from crewai.experimental.evaluation.base_evaluator import AgentEvaluationResult, AggregationStrategy
from crewai.agent import Agent
from crewai.task import Task
from crewai.experimental.evaluation.evaluation_display import EvaluationDisplayFormatter
from crewai.utilities.events.agent_events import AgentEvaluationStartedEvent, AgentEvaluationCompletedEvent, AgentEvaluationFailedEvent
from crewai.experimental.evaluation import BaseEvaluator, create_evaluation_callbacks
from collections.abc import Sequence
from crewai.utilities.events.crewai_event_bus import crewai_event_bus
from crewai.utilities.events.utils.console_formatter import ConsoleFormatter
from crewai.utilities.events.task_events import TaskCompletedEvent
from crewai.utilities.events.agent_events import LiteAgentExecutionCompletedEvent
from crewai.experimental.evaluation.base_evaluator import AgentAggregatedEvaluationResult, EvaluationScore, MetricCategory

class ExecutionState:
    def __init__(self):
        self.traces = {}
        self.current_agent_id: str | None = None
        self.current_task_id: str | None = None
        self.iteration = 1
        self.iterations_results = {}
        self.agent_evaluators = {}

class AgentEvaluator:
    def __init__(
        self,
        agents: list[Agent],
        evaluators: Sequence[BaseEvaluator] | None = None,
    ):
        self.agents: list[Agent] = agents
        self.evaluators: Sequence[BaseEvaluator] | None = evaluators

        self.callback = create_evaluation_callbacks()
        self.console_formatter = ConsoleFormatter()
        self.display_formatter = EvaluationDisplayFormatter()

        self._thread_local: threading.local = threading.local()

        for agent in self.agents:
            self._execution_state.agent_evaluators[str(agent.id)] = self.evaluators

        self._subscribe_to_events()

    @property
    def _execution_state(self) -> ExecutionState:
        if not hasattr(self._thread_local, 'execution_state'):
            self._thread_local.execution_state = ExecutionState()
        return self._thread_local.execution_state

    def _subscribe_to_events(self) -> None:
        from typing import cast
        crewai_event_bus.register_handler(TaskCompletedEvent, cast(Any, self._handle_task_completed))
        crewai_event_bus.register_handler(LiteAgentExecutionCompletedEvent, cast(Any, self._handle_lite_agent_completed))

    def _handle_task_completed(self, source: Any, event: TaskCompletedEvent) -> None:
        assert event.task is not None
        agent = event.task.agent
        if agent and str(getattr(agent, 'id', 'unknown')) in self._execution_state.agent_evaluators:
            self.emit_evaluation_started_event(agent_role=agent.role, agent_id=str(agent.id), task_id=str(event.task.id))

            state = ExecutionState()
            state.current_agent_id = str(agent.id)
            state.current_task_id = str(event.task.id)

            assert state.current_agent_id is not None and state.current_task_id is not None
            trace = self.callback.get_trace(state.current_agent_id, state.current_task_id)

            if not trace:
                return

            result = self.evaluate(
                agent=agent,
                task=event.task,
                execution_trace=trace,
                final_output=event.output,
                state=state
            )

            current_iteration = self._execution_state.iteration
            if current_iteration not in self._execution_state.iterations_results:
                self._execution_state.iterations_results[current_iteration] = {}

            if agent.role not in self._execution_state.iterations_results[current_iteration]:
                self._execution_state.iterations_results[current_iteration][agent.role] = []

            self._execution_state.iterations_results[current_iteration][agent.role].append(result)

    def _handle_lite_agent_completed(self, source: object, event: LiteAgentExecutionCompletedEvent) -> None:
        agent_info = event.agent_info
        agent_id = str(agent_info["id"])

        if agent_id in self._execution_state.agent_evaluators:
            state = ExecutionState()
            state.current_agent_id = agent_id
            state.current_task_id = "lite_task"

            target_agent = None
            for agent in self.agents:
                if str(agent.id) == agent_id:
                    target_agent = agent
                    break

            if not target_agent:
                return

            assert state.current_agent_id is not None and state.current_task_id is not None
            trace = self.callback.get_trace(state.current_agent_id, state.current_task_id)

            if not trace:
                return

            result = self.evaluate(
                agent=target_agent,
                execution_trace=trace,
                final_output=event.output,
                state=state
            )

            current_iteration = self._execution_state.iteration
            if current_iteration not in self._execution_state.iterations_results:
                self._execution_state.iterations_results[current_iteration] = {}

            agent_role = target_agent.role
            if agent_role not in self._execution_state.iterations_results[current_iteration]:
                self._execution_state.iterations_results[current_iteration][agent_role] = []

            self._execution_state.iterations_results[current_iteration][agent_role].append(result)

    def set_iteration(self, iteration: int) -> None:
        self._execution_state.iteration = iteration

    def reset_iterations_results(self) -> None:
        self._execution_state.iterations_results = {}

    def get_evaluation_results(self) -> dict[str, list[AgentEvaluationResult]]:
        if self._execution_state.iterations_results and self._execution_state.iteration in self._execution_state.iterations_results:
            return self._execution_state.iterations_results[self._execution_state.iteration]
        return {}

    def display_results_with_iterations(self) -> None:
        self.display_formatter.display_summary_results(self._execution_state.iterations_results)

    def get_agent_evaluation(self, strategy: AggregationStrategy = AggregationStrategy.SIMPLE_AVERAGE, include_evaluation_feedback: bool = True) -> dict[str, AgentAggregatedEvaluationResult]:
        agent_results = {}
        with crewai_event_bus.scoped_handlers():
            task_results = self.get_evaluation_results()
            for agent_role, results in task_results.items():
                if not results:
                    continue

                agent_id = results[0].agent_id

                aggregated_result = self.display_formatter._aggregate_agent_results(
                    agent_id=agent_id,
                    agent_role=agent_role,
                    results=results,
                    strategy=strategy
                )

                agent_results[agent_role] = aggregated_result


            if self._execution_state.iterations_results and self._execution_state.iteration == max(self._execution_state.iterations_results.keys(), default=0):
                self.display_results_with_iterations()

            if include_evaluation_feedback:
                self.display_evaluation_with_feedback()

        return agent_results

    def display_evaluation_with_feedback(self) -> None:
        self.display_formatter.display_evaluation_with_feedback(self._execution_state.iterations_results)

    def evaluate(
        self,
        agent: Agent,
        execution_trace: dict[str, Any],
        final_output: Any,
        state: ExecutionState,
        task: Task | None = None,
    ) -> AgentEvaluationResult:
        result = AgentEvaluationResult(
            agent_id=state.current_agent_id or str(agent.id),
            task_id=state.current_task_id or (str(task.id) if task else "unknown_task")
        )

        assert self.evaluators is not None
        task_id = str(task.id) if task else None
        for evaluator in self.evaluators:
            try:
                self.emit_evaluation_started_event(agent_role=agent.role, agent_id=str(agent.id), task_id=task_id)
                score = evaluator.evaluate(
                    agent=agent,
                    task=task,
                    execution_trace=execution_trace,
                    final_output=final_output
                )
                result.metrics[evaluator.metric_category] = score
                self.emit_evaluation_completed_event(agent_role=agent.role, agent_id=str(agent.id), task_id=task_id, metric_category=evaluator.metric_category, score=score)
            except Exception as e:
                self.emit_evaluation_failed_event(agent_role=agent.role, agent_id=str(agent.id), task_id=task_id, error=str(e))
                self.console_formatter.print(f"Error in {evaluator.metric_category.value} evaluator: {str(e)}")

        return result

    def emit_evaluation_started_event(self, agent_role: str, agent_id: str, task_id: str | None = None):
        crewai_event_bus.emit(
            self,
            AgentEvaluationStartedEvent(agent_role=agent_role, agent_id=agent_id, task_id=task_id, iteration=self._execution_state.iteration)
        )

    def emit_evaluation_completed_event(self, agent_role: str, agent_id: str, task_id: str | None = None, metric_category: MetricCategory | None = None, score: EvaluationScore | None = None):
        crewai_event_bus.emit(
            self,
            AgentEvaluationCompletedEvent(agent_role=agent_role, agent_id=agent_id, task_id=task_id, iteration=self._execution_state.iteration, metric_category=metric_category, score=score)
        )

    def emit_evaluation_failed_event(self, agent_role: str, agent_id: str, error: str, task_id: str | None = None):
        crewai_event_bus.emit(
            self,
            AgentEvaluationFailedEvent(agent_role=agent_role, agent_id=agent_id, task_id=task_id, iteration=self._execution_state.iteration, error=error)
        )

def create_default_evaluator(agents: list[Agent], llm: None = None):
    from crewai.experimental.evaluation import (
        GoalAlignmentEvaluator,
        SemanticQualityEvaluator,
        ToolSelectionEvaluator,
        ParameterExtractionEvaluator,
        ToolInvocationEvaluator,
        ReasoningEfficiencyEvaluator
    )

    evaluators = [
        GoalAlignmentEvaluator(llm=llm),
        SemanticQualityEvaluator(llm=llm),
        ToolSelectionEvaluator(llm=llm),
        ParameterExtractionEvaluator(llm=llm),
        ToolInvocationEvaluator(llm=llm),
        ReasoningEfficiencyEvaluator(llm=llm),
    ]

    return AgentEvaluator(evaluators=evaluators, agents=agents)
