import json
from enum import Enum
from threading import Lock
from typing import TYPE_CHECKING, Any, Dict, List, Union

from pydantic import BaseModel, Field, model_validator
from typing_extensions import Self

from kiln_ai.datamodel.basemodel import (
    ID_TYPE,
    FilenameString,
    FilenameStringShort,
    KilnParentedModel,
    KilnParentModel,
)
from kiln_ai.datamodel.datamodel_enums import TaskOutputRatingType
from kiln_ai.datamodel.dataset_filters import DatasetFilterId
from kiln_ai.datamodel.json_schema import string_to_json_key
from kiln_ai.datamodel.task_run import Usage
from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error

if TYPE_CHECKING:
    from kiln_ai.datamodel.task import Task

EvalScores = Dict[str, float]

# Module-level set to track evals currently being migrated (to prevent recursion)
# Protected by _migration_lock to ensure thread-safe access
_migration_lock = Lock()
_currently_migrating_eval_ids: set[ID_TYPE] = set()


class EvalTemplateId(str, Enum):
    """
    An eval template is a pre-defined eval that can be used as a starting point for a new eval.
    """

    kiln_requirements = "kiln_requirements"
    issue = "kiln_issue"
    tool_call = "tool_call"
    toxicity = "toxicity"
    bias = "bias"
    maliciousness = "maliciousness"
    factual_correctness = "factual_correctness"
    jailbreak = "jailbreak"
    rag = "rag"


class EvalConfigType(str, Enum):
    g_eval = "g_eval"
    llm_as_judge = "llm_as_judge"


class EvalOutputScore(BaseModel):
    """
    A definition of a score that an evaluator will produce.

    Very similar to TaskRequirement, but conceptually different keeping in a separate models.
    """

    name: FilenameStringShort = Field(
        description="The name of the score. Will be provided to the model so use a descriptive name. Should align to the model's TaskRequirement name if you want to use human evals to evaluate the evaluator's performance."
    )
    instruction: str | None = Field(
        default=None,
        description="A description of the score, used to help the model understand the goal of the score. Will be provided to evaluator models, so should be written for the model, not the team/user.",
    )
    type: TaskOutputRatingType = Field(
        description="The type of rating to use ('five_star', 'pass_fail', 'pass_fail_critical')."
    )

    def json_key(self) -> str:
        """
        The JSON key for the score, used when running the evaluator with a LLM and we need JSON output.

        For example, "Overall Rating" -> "overall_rating"
        """
        return string_to_json_key(self.name)

    @model_validator(mode="after")
    def validate_type(self) -> Self:
        if self.type == TaskOutputRatingType.custom:
            raise ValueError(
                f"Custom scores are not supported in evaluators. Score '{self.name}' was set to a custom score."
            )
        return self


class EvalRun(KilnParentedModel):
    """
    The results of running an eval on a single dataset item.

    This is a child of an EvalConfig, which specifies how the scores were generated.

    Eval runs can be one of 2 types:
    1) eval_config_eval=False: we were evaluating a task run (a method of running the task). We get the task input from the dataset_id.input, run the task with the task_run_config, then ran the evaluator on that output. task_run_config_id must be set. The output saved in this model is the output of the task run.
    2) eval_config_eval=True: we were evaluating an eval config (a method of evaluating the task). We used the existing dataset item input/output, and ran the evaluator on it. task_run_config_id must be None. The input/output saved in this model is the input/output of the dataset item.
    """

    dataset_id: ID_TYPE = Field(
        description="The ID of the dataset item that was used for this run. Must belong to the same Task as the grand-parent eval of this EvalRun."
    )
    task_run_config_id: ID_TYPE | None = Field(
        description="The ID of the TaskRunConfig that was run, if this eval run was based on a task run. Must belong to the same Task as this eval. Can be None if this eval run is based on an eval config."
    )
    eval_config_eval: bool = Field(
        description="Whether this eval run to evaluate the parent eval config (evaluating the config using an existing dataset item). If true, task_run_config_id must be None, as we're not running the task.",
        default=False,
    )
    # These two may duplicate the dataset_id.input/output, but we're denormalizing intentionally.
    input: str = Field(
        description="The input to the task. JSON formatted for structured input, plaintext for unstructured input."
    )
    output: str = Field(
        description="The output of the task. JSON formatted for structured output, plaintext for unstructured output."
    )
    reference_answer: str | None = Field(
        default=None,
        description="The reference answer for the input. JSON formatted for structured reference answer, plaintext for unstructured reference answer. Used for reference answer evals.",
    )
    intermediate_outputs: Dict[str, str] | None = Field(
        default=None,
        description="The intermediate outputs of the task (example, eval thinking).",
    )
    task_run_trace: str | None = Field(
        default=None,
        description="The JSON formatted trace of the task run that produced the output.",
    )
    scores: EvalScores = Field(
        description="The output scores of the evaluator (aligning to those required by the grand-parent Eval this object is a child of)."
    )
    task_run_usage: Usage | None = Field(
        default=None,
        description="The usage of the task run that produced this eval run output (not the usage by the evaluation model).",
    )

    def parent_eval_config(self) -> Union["EvalConfig", None]:
        if self.parent is not None and self.parent.__class__.__name__ != "EvalConfig":
            raise ValueError("parent must be an EvalConfig")
        return self.parent  # type: ignore

    @model_validator(mode="after")
    def validate_output_fields(self) -> Self:
        parent_eval_config = self.parent_eval_config()
        parent_eval = parent_eval_config.parent_eval() if parent_eval_config else None
        if not parent_eval:
            return self

        evaluation_data_type = parent_eval.evaluation_data_type
        if (
            evaluation_data_type == EvalDataType.final_answer
            and self.task_run_trace is not None
        ):
            raise ValueError("final_answer runs should not set trace")
        elif (
            not self.eval_config_eval
            and evaluation_data_type == EvalDataType.full_trace
            and self.task_run_trace is None
        ):
            raise ValueError("full_trace task run eval runs should include trace")

        return self

    @model_validator(mode="after")
    def validate_eval_run_types(self) -> Self:
        if self.eval_config_eval and self.task_run_config_id is not None:
            raise ValueError(
                "task_run_config_id must be None if eval_config_eval is true"
            )
        if not self.eval_config_eval and self.task_run_config_id is None:
            raise ValueError(
                "task_run_config_id must be set if eval_config_eval is false"
            )
        return self

    @model_validator(mode="after")
    def validate_scores(self) -> Self:
        # We're checking the scores have the expected keys from the grand-parent eval
        if self.scores is None or len(self.scores) == 0:
            raise ValueError("scores are required, and must have at least one score.")

        parent_eval_config = self.parent_eval_config()
        eval = parent_eval_config.parent_eval() if parent_eval_config else None
        if not eval:
            # Can't validate without the grand-parent eval, allow it to be validated later
            return self

        output_score_keys = [score.json_key() for score in eval.output_scores]
        if set(output_score_keys) != set(self.scores.keys()):
            raise ValueError(
                f"The scores produced by the evaluator must match the scores expected by the eval. Got: [{', '.join(self.scores.keys())}] and expected: [{', '.join(output_score_keys)}]"
            )

        # Check that each score is expected in this eval and the correct type
        for output_score in eval.output_scores:
            match output_score.type:
                case TaskOutputRatingType.five_star:
                    five_star_score = self.scores[output_score.json_key()]
                    if (
                        not isinstance(five_star_score, float)
                        or five_star_score < 1.0
                        or five_star_score > 5.0
                    ):
                        raise ValueError(
                            f"Score {output_score.name} is a five_star rating and must be a float between 1.0 and 5.0 inclusive. Got: {five_star_score}"
                        )
                case TaskOutputRatingType.pass_fail:
                    pass_fail_score = self.scores[output_score.json_key()]
                    if (
                        not isinstance(pass_fail_score, float)
                        or pass_fail_score < 0.0
                        or pass_fail_score > 1.0
                    ):
                        raise ValueError(
                            f"Score {output_score.name} is a pass_fail rating and must be a float between 0.0 and 1.0 inclusive. Got: {pass_fail_score}"
                        )
                case TaskOutputRatingType.pass_fail_critical:
                    pass_fail_critical_score = self.scores[output_score.json_key()]
                    if (
                        not isinstance(pass_fail_critical_score, float)
                        or pass_fail_critical_score < -1.0
                        or pass_fail_critical_score > 1.0
                    ):
                        raise ValueError(
                            f"Score {output_score.name} is a pass_fail_critical rating and must be a float between -1.0 and 1.0 inclusive. Got: {pass_fail_critical_score}"
                        )
                case TaskOutputRatingType.custom:
                    raise ValueError(
                        f"Custom scores are not supported in evaluators. '{output_score.name}' was set to a custom score."
                    )
                case _:
                    # Catch missing cases
                    raise_exhaustive_enum_error(output_score.type)
        return self

    @model_validator(mode="after")
    def validate_reference_answer(self) -> Self:
        parent_eval_config = self.parent_eval_config()
        parent_eval = parent_eval_config.parent_eval() if parent_eval_config else None
        if not parent_eval:
            # Can't validate without the grand-parent eval, allow it to be validated later
            return self

        evaluation_data_type = parent_eval.evaluation_data_type
        if (
            self.reference_answer is not None
            and evaluation_data_type != EvalDataType.reference_answer
        ):
            raise ValueError(
                f"reference_answer is only valid for reference answer evals. Got: {evaluation_data_type.value}"
            )
        return self


class EvalConfig(KilnParentedModel, KilnParentModel, parent_of={"runs": EvalRun}):
    """
    A configuration for running an eval. This includes anything needed to run the eval on a dataset like the prompt, model, thresholds, etc.

    A eval might have many configs, example running the same eval with 2 different models. Comparing eval results is only valid within the scope of the same config.
    """

    name: FilenameString = Field(description="The name of the eval config.")
    model_name: str = Field(
        description="The name of the model to use for this eval config. ",
    )
    model_provider: str = Field(
        description="The provider of the model to use for this eval config.",
    )
    config_type: EvalConfigType = Field(
        default=EvalConfigType.g_eval,
        description="This is used to determine the type of eval to run.",
    )
    properties: dict[str, Any] = Field(
        default={},
        description="Properties to be used to execute the eval config. This is config_type specific and should serialize to a json dict.",
    )

    def parent_eval(self) -> Union["Eval", None]:
        if self.parent is not None and self.parent.__class__.__name__ != "Eval":
            raise ValueError("parent must be an Eval")
        return self.parent  # type: ignore

    def runs(self, readonly: bool = False) -> list[EvalRun]:
        return super().runs(readonly=readonly)  # type: ignore

    @model_validator(mode="after")
    def validate_properties(self) -> Self:
        if (
            self.config_type == EvalConfigType.g_eval
            or self.config_type == EvalConfigType.llm_as_judge
        ):
            if "eval_steps" not in self.properties or not isinstance(
                self.properties["eval_steps"], list
            ):
                raise ValueError("eval_steps is required and must be a list for g_eval")
            if "task_description" in self.properties and not isinstance(
                self.properties["task_description"], str
            ):
                raise ValueError(
                    "task_description is optional, but if provided must be a string"
                )
            return self
        else:
            raise ValueError(f"Invalid eval config type: {self.config_type}")

    @model_validator(mode="after")
    def validate_json_serializable(self) -> "EvalConfig":
        try:
            # This will raise a TypeError if the dict contains non-JSON-serializable objects
            json.dumps(self.properties)
        except TypeError as e:
            raise ValueError(f"Properties must be JSON serializable: {e!s}")
        return self


class EvalDataType(str, Enum):
    final_answer = "final_answer"
    full_trace = "full_trace"
    reference_answer = "reference_answer"


class Eval(KilnParentedModel, KilnParentModel, parent_of={"configs": EvalConfig}):
    name: FilenameString = Field(description="The name of the eval.")
    description: str | None = Field(
        default=None, description="The description of the eval"
    )
    template: EvalTemplateId | None = Field(
        default=None,
        description="The template selected when creating this eval. Useful for suggesting eval steps and output scores.",
    )
    current_config_id: ID_TYPE = Field(
        default=None,
        description="The id of the current config to use for this eval. This can be changed over time to run the same eval with different configs.",
    )
    eval_set_filter_id: DatasetFilterId = Field(
        description="The id of the dataset filter which defines which dataset items are included when running this eval. Should be mutually exclusive with eval_configs_filter_id."
    )
    eval_configs_filter_id: DatasetFilterId | None = Field(
        default=None,
        description="The id of the dataset filter which defines which dataset items are included when comparing the quality of the eval configs under this eval. Should consist of dataset items with ratings. Should be mutually exclusive with eval_set_filter_id.",
    )
    output_scores: List[EvalOutputScore] = Field(
        description="The scores this evaluator should produce."
    )
    favourite: bool = Field(
        default=False,
        description="Whether this eval is a favourite of the user. Rendered as a star icon in the UI.",
    )
    template_properties: dict[str, str | int | bool | float] = Field(
        default={},
        description="Properties to be used to execute the eval. This is template_type specific and should serialize to a json dict.",
    )
    evaluation_data_type: EvalDataType = Field(
        default=EvalDataType.final_answer,
        description="The output of the task run to evaluate. Can be final answer or full trace.",
    )

    # Workaround to return typed parent without importing Task
    def parent_task(self) -> Union["Task", None]:
        if self.parent is not None and self.parent.__class__.__name__ != "Task":
            raise ValueError("parent must be a Task")
        return self.parent  # type: ignore

    def configs(self, readonly: bool = False) -> list[EvalConfig]:
        return super().configs(readonly=readonly)  # type: ignore

    @model_validator(mode="after")
    def upgrade_old_reference_answer_eval_config(self) -> Self:
        """
        Migration: Set the first judge config as the default for existing reference answer evals that don't have a current_config_id set.

        For reference_answer evals that don't have a current_config_id set, this migration
        will set the first config (by created_at) as the default.
        """
        if self.id is None:
            return self

        # Only run during file loading
        if not self._loaded_from_file:
            return self

        # Skip if already migrated (has a current_config_id set)
        if self.current_config_id is not None:
            return self

        # Only migrate reference_answer evals
        if self.evaluation_data_type != EvalDataType.reference_answer:
            return self

        # Prevent recursion: self.configs() loads child files, which re-loads this parent
        # (see basemodel.py where we iterate_children_paths_of_parent_path calls load_from_file)
        # This causes the validator to run again, creating an infinite loop without this guard.
        with _migration_lock:
            if self.id in _currently_migrating_eval_ids:
                return self
            _currently_migrating_eval_ids.add(self.id)

        try:
            # Get the configs - these are loaded from child files
            configs_list = self.configs(readonly=True)
            if configs_list and len(configs_list) > 0:
                # Sort by created_at to get the oldest (first created) config
                sorted_configs = sorted(configs_list, key=lambda c: c.created_at)
                self.current_config_id = sorted_configs[0].id
        finally:
            with _migration_lock:
                _currently_migrating_eval_ids.discard(self.id)

        return self

    @model_validator(mode="after")
    def validate_scores(self) -> Self:
        if self.output_scores is None or len(self.output_scores) == 0:
            raise ValueError(
                "output_scores are required, and must have at least one score."
            )

        # check for duplicate names (once transformed to JSON keys)
        output_score_keys = [score.json_key() for score in self.output_scores]
        if len(output_score_keys) != len(set(output_score_keys)):
            raise ValueError(
                f"output_scores must have unique names (once transformed to JSON keys). Got: [{', '.join(output_score_keys)}]"
            )
        return self

    @model_validator(mode="after")
    def validate_template_properties(self) -> Self:
        # eval_configs_filter_id is required for all templates except "rag"
        if (
            self.template is not EvalTemplateId.rag
            and self.eval_configs_filter_id is None
        ):
            raise ValueError(
                "eval_configs_filter_id is required for all templates except 'rag'"
            )

        # Check for properties that are required for the issue template
        if self.template == EvalTemplateId.issue:
            if "issue_prompt" not in self.template_properties or not isinstance(
                self.template_properties["issue_prompt"], str
            ):
                raise ValueError("issue_prompt is required for issue template")
            if "failure_example" in self.template_properties and not isinstance(
                self.template_properties["failure_example"], str
            ):
                raise ValueError(
                    "failure_example is optional for issue template, but if provided must be a string"
                )
            if "pass_example" in self.template_properties and not isinstance(
                self.template_properties["pass_example"], str
            ):
                raise ValueError(
                    "pass_example is optional for issue template, but if provided must be a string"
                )

        if self.template == EvalTemplateId.tool_call:
            if self.evaluation_data_type != EvalDataType.full_trace:
                raise ValueError(
                    "tool_call template should have evaluation_data_type set to full_trace"
                )
            if (
                "tool" not in self.template_properties
                or not isinstance(self.template_properties["tool"], str)
                or not self.template_properties["tool"].strip()
            ):
                raise ValueError("tool is required for tool call template")
            if "tool_function_name" not in self.template_properties or not isinstance(
                self.template_properties["tool_function_name"], str
            ):
                raise ValueError(
                    "tool_function_name is required for tool call template"
                )
            if (
                "appropriate_tool_use_guidelines" not in self.template_properties
                or not isinstance(
                    self.template_properties["appropriate_tool_use_guidelines"], str
                )
                or not self.template_properties[
                    "appropriate_tool_use_guidelines"
                ].strip()
            ):
                raise ValueError(
                    "appropriate_tool_use_guidelines is required for tool call template"
                )
            if (
                "inappropriate_tool_use_guidelines" in self.template_properties
                and not isinstance(
                    self.template_properties["inappropriate_tool_use_guidelines"], str
                )
            ):
                raise ValueError(
                    "inappropriate_tool_use_guidelines is optional for tool call template, but if provided must be a string"
                )
        return self
