from collections.abc import Mapping
from copy import deepcopy
from dataclasses import InitVar, asdict, dataclass, field
from typing import Any, Protocol, TypeAlias, TypedDict, runtime_checkable
from uuid import uuid4

from relai.mocker.base_mocker import BaseMocker

SimulationConfigT: TypeAlias = Mapping[str, BaseMocker]

AgentInputs: TypeAlias = dict[str, Any]
Extras: TypeAlias = dict[str, Any]
MockedData: TypeAlias = dict[str, Any]
AgentOutputs: TypeAlias = dict[str, Any]


class EvaluatorOutputs(TypedDict):
    """
    Outputs generated by an evaluator.

    Attributes:
        score (float): The score assigned by the evaluator.
        feedback (str): Any natural language feedback provided by the evaluator.
    """

    score: float
    feedback: str


@runtime_checkable
class Serializable(Protocol):
    def serialize(self) -> dict[str, Any]: ...


@dataclass(slots=True, frozen=True)
class RELAISample:
    """
    Represents a single sample in a RELAI benchmark.

    Attributes:
        benchmark_id (str): The identifier of the benchmark this sample belongs to.
        id (str): The unique identifier for this sample.
        split (str): The data split this sample belongs to (e.g., "Train", "Validation", "Test").
        agent_inputs (AgentInputs): The inputs provided to the agent from this sample.
        extras (Extras): Any additional metadata or information associated with this sample.
            Use this field to also store evaluator-specific inputs.
        serialized_simulation_config (dict[str, Any]): The serialized simulation configuration for this sample.
            Can be used to reconstruct any mockers used in a previous simulation.
    """

    benchmark_id: str = field(default="default")
    id: str = field(default_factory=lambda: uuid4().hex)
    split: str = field(default="All")
    agent_inputs: AgentInputs = field(default_factory=dict)
    extras: Extras = field(default_factory=dict)
    serialized_simulation_config: dict[str, Any] = field(default_factory=dict)  # `SimulationConfigT` as dict


@dataclass(slots=True)
class SimulationTape:
    """
    A simulation tape records any inputs, outputs, and any other relevant data during a simulation.

    Attributes:
        id (str): The unique identifier for this simulation tape.
        benchmark_id (str): The identifier of the benchmark from which the sample used to initialize this tape was
            taken. If no sample was provided, defaults to "default".
        sample_id (str): The identifier of the sample used to initialize this tape. If no sample was provided,
            defaults to the tape's id.
        split (str): The data split of the sample used to initialize this tape. If no sample was provided,
            defaults to "All".
        agent_inputs (AgentInputs): The inputs provided to the agent from the sample used to initialize this tape. If
            no sample was provided, defaults to an empty dictionary.
        extras (Extras): Any additional metadata or information associated with this tape. If no sample was provided,
            defaults to an empty dictionary.
        evaluator_group (str): The evaluator group associated with this tape, typically set to the benchmark_id of the
            sample. If no sample was provided, defaults to "default". Can be modified in the `agent_fn`.
        simulation_config (SimulationConfigT): The simulation configuration used during this simulation which is a
            mapping from qualified function names to their respective mocker instances (Persona, MockTool etc.,).
    """

    sample: InitVar[RELAISample | None] = field(default=None)
    id: str = field(default_factory=lambda: uuid4().hex)
    benchmark_id: str = field(init=False)
    sample_id: str = field(init=False)
    split: str = field(init=False)
    agent_inputs: AgentInputs = field(init=False)
    extras: Extras = field(init=False)
    evaluator_group: str = field(init=False)
    simulation_config: SimulationConfigT = field(default_factory=dict)

    def __post_init__(self, sample: RELAISample | None) -> None:
        if sample is not None:
            self.benchmark_id = sample.benchmark_id
            self.sample_id = sample.id
            self.split = sample.split
            self.agent_inputs = deepcopy(sample.agent_inputs)
            self.extras = deepcopy(sample.extras)
            self.evaluator_group = "default"
        else:
            self.benchmark_id = "default"
            self.sample_id = self.id
            self.split = "All"
            self.agent_inputs = {}
            self.extras = {}
            self.evaluator_group = "default"

    def add_record(self, key: str, value: Any) -> None:
        self.extras[key] = value

    def set_evaluator_group(self, evaluator_group: str) -> None:
        self.evaluator_group = evaluator_group

    @property
    def data(self) -> Extras:
        return self.agent_inputs | self.extras

    def extract_sample(self) -> RELAISample:
        return RELAISample(
            benchmark_id=self.benchmark_id,
            id=self.sample_id,
            agent_inputs=deepcopy(self.agent_inputs),
            extras=deepcopy(self.extras),
            serialized_simulation_config={
                k: v.serialize() for k, v in self.simulation_config.items() if isinstance(v, Serializable)
            },
        )

    def copy(self) -> "SimulationTape":
        ret = SimulationTape(sample=self.extract_sample())
        ret.simulation_config = self.simulation_config
        return ret


@dataclass(slots=True)
class AgentLog:
    """
    Log of a single agent simulation run.

    Attributes:
        simulation_tape (SimulationTape): The simulation tape containing inputs and metadata.
        agent_outputs (AgentOutputs): The outputs generated by the agent during the simulation.
        trace_id (str | None): An optional trace identifier for the simulation run.
    """

    simulation_tape: SimulationTape
    agent_outputs: AgentOutputs = field(default_factory=dict)
    trace_id: str | None = field(default=None)


@dataclass(slots=True)
class EvaluatorLog:
    """
    Log of a single evaluator run.

    Attributes:
        evaluator_id (str): The ID of the evaluator.
        name (str): The name of the evaluator.
        outputs (EvaluatorOutputs): The outputs generated by the evaluator.
        config (dict[str, Any]): The configuration settings used for the evaluator.
    """

    evaluator_id: str
    name: str
    outputs: EvaluatorOutputs
    config: dict[str, Any] = field(default_factory=dict)


@dataclass(slots=True)
class CriticoLog:
    """
    Log of a Critico evaluation run.

    Attributes:
        agent_log (AgentLog): The log of the agent simulation run.
        evaluator_logs (list[EvaluatorLog]): A list of logs from individual evaluators.
        aggregate_score (float): The aggregate score computed from all the evaluator logs.
        aggregate_feedback (str): The aggregate feedback compiled from all the evaluator logs.
        trace_id (str | None): An optional trace identifier for the corresponding agent simulation run.
    """

    agent_log: AgentLog
    evaluator_logs: list[EvaluatorLog] = field(default_factory=list)
    aggregate_score: float = field(default=0.0)
    aggregate_feedback: str = field(default="")
    trace_id: str | None = field(default=None)

    def serialize(self) -> dict[str, Any]:
        return {
            "trace_id": self.trace_id,
            "agent_log": {
                "sample": asdict(self.agent_log.simulation_tape.extract_sample()),
                "agent_outputs": self.agent_log.agent_outputs,
            },
            "evaluator_logs": [asdict(evaluator_log) for evaluator_log in self.evaluator_logs],
            "aggregate_score": self.aggregate_score,
            "aggregate_feedback": self.aggregate_feedback,
        }
