from typing import Any

from jsonschema import ValidationError as SchemaValidationError
from jsonschema import validate
from pydantic import Field, field_validator

from core.domain._autogenerated_id import AutoGeneratedId
from core.domain.exceptions import JSONSchemaValidationError
from core.domain.message import Message
from core.domain.reasoning_effort import ReasoningEffort
from core.domain.tool import HostedTool, Tool
from core.domain.tool_choice import ToolChoice
from core.utils.schemas import JsonSchema, make_optional, remove_extra_keys, sanitize_empty_values


class Version(AutoGeneratedId):
    """Properties that described a way a task run was executed.
    Although some keys are provided as an example, any key:value are accepted"""

    model: str | None = Field(default=None, description="The LLM model used for the run")
    provider: str | None = Field(default=None, description="The LLM provider used for the run")
    temperature: float | None = Field(default=None, description="The temperature for generation")

    max_output_tokens: int | None = Field(
        default=None,
        description="The maximum tokens to generate in the prompt",
    )

    # A set would have been nicer but isn't JSON serializable for storage and would require custom code.
    # Pass a full tool to enable external tools

    # We use a str fallback to let the runner decide how to handle the tool
    enabled_tools: list[HostedTool | Tool] | None = None

    use_structured_generation: bool | None = Field(
        default=None,
        description="Whether to use structured generation for the task",
    )

    tool_choice: ToolChoice | None = None

    top_p: float | None = None

    presence_penalty: float | None = None

    frequency_penalty: float | None = None

    parallel_tool_calls: bool | None = Field(
        default=None,
        description="Whether to allow the model to output mutliple tool calls",
    )

    prompt: list[Message] | None = None

    reasoning_effort: ReasoningEffort | None = None

    reasoning_budget: int | None = None

    input_variables_schema: dict[str, Any] | None = Field(
        default=None,
        description="A JSON schema for the variables used to template the instructions during the inference."
        "Auto generated from the prompt if the prompt is a Jinja2 template",
    )

    class OutputSchema(AutoGeneratedId):
        json_schema: dict[str, Any]

    output_schema: OutputSchema | None = Field(
        default=None,
        description="A JSON schema for the output of the model, aka the schema in the response format",
    )

    def validate_input(self, obj: dict[str, Any] | None):
        if self.input_variables_schema is None:
            if obj:
                raise JSONSchemaValidationError("Input variables are provided but the version does not support them")
            return
        if obj is None:
            raise JSONSchemaValidationError("Input variables are not provided but the version requires them")
        try:
            validate(obj, self.input_variables_schema)
        except SchemaValidationError as e:
            kp = ".".join([str(p) for p in e.path])
            raise JSONSchemaValidationError(
                f"Input variables are not compatible with the version's input variables schema.\nat [{kp}], {e.message}",
            ) from e

    def validate_output(
        self,
        obj: Any,
        partial: bool = False,
        strip_extras: bool = False,
        sanitize_empties: bool = False,
    ):
        """Enforce validates that an object matches the schema. Object is updated in place."""
        if not self.output_schema:
            return obj

        schema = make_optional(self.output_schema.json_schema) if partial else self.output_schema.json_schema

        navigators: list[JsonSchema.Navigator] = []
        if sanitize_empties:
            navigators.append(sanitize_empty_values)
        if strip_extras:
            navigators.append(remove_extra_keys)

        if navigators:
            JsonSchema(schema).navigate(obj, navigators=navigators)

        try:
            validate(obj, schema)
        except SchemaValidationError as e:
            kp = ".".join([str(p) for p in e.path])
            raise JSONSchemaValidationError(f"at [{kp}], {e.message}") from e
        return obj

    def should_use_auto_cache(self):
        return self.temperature == 0 and not self.enabled_tools

    @field_validator("enabled_tools", mode="before")
    @classmethod
    def validate_enabled_tools(cls, v: Any) -> Any:
        if not v:
            return None
        if not isinstance(v, list):
            raise ValueError("enabled_tools must be a list")

        return [_validate_tool(t) for t in v]


def _validate_tool(v: Any) -> Tool | HostedTool:
    if isinstance(v, str):
        try:
            return HostedTool(v)
        except ValueError:
            return Tool(name=v, input_schema={}, output_schema={})
    return Tool.model_validate(v)
