"""The `llm.format` decorator for defining response formats as classes."""

from __future__ import annotations

import inspect
import json
from dataclasses import dataclass
from typing import Any, Generic, cast

from ..tools import FORMAT_TOOL_NAME, ToolFn, ToolParameterSchema, ToolSchema
from ..types import NoneType
from .output_parser import OutputParser, is_output_parser
from .primitives import create_wrapper_model, is_primitive_type
from .types import FormatSpec, FormattableT, FormattingMode, HasFormattingInstructions

TOOL_MODE_INSTRUCTIONS = f"""Always respond to the user's query using the {FORMAT_TOOL_NAME} tool for structured output."""


JSON_MODE_INSTRUCTIONS = (
    "Respond only with valid JSON that matches this exact schema:\n{json_schema}"
)


@dataclass(kw_only=True)
class Format(Generic[FormattableT]):
    """Class representing a structured output format for LLM responses.

    A `Format` contains metadata needed to describe a structured output type
    to the LLM, including the expected schema. This class is not instantiated directly,
    but is created by calling `llm.format`, or is automatically generated by LLM
    providers when a `Formattable` is passed to a call method.

    Example:

      ```python
      from mirascope import llm

      class Book:
          title: str
          author: str

      print(llm.format(Book, mode="tool"))
      ```
    """

    name: str
    """The name of the response format."""

    description: str | None
    """A description of the response format, if available."""

    schema: dict[str, object]
    """JSON schema representation of the structured output format."""

    mode: FormattingMode
    """The decorator-provided mode of the response format.

    Determines how the LLM call may be modified in order to extract the expected format.
    """

    formattable: type[FormattableT] | OutputParser[FormattableT]
    """The formattable type or custom output parser.

    Can be one of:
    - type[BaseModel]: A Pydantic model class for structured output
    - PrimitiveType: A primitive type (str, int, list, etc.) for simple output
    - OutputParser[FormattableT]: A custom parser created with @llm.output_parser

    The type determines how the response will be parsed in response.parse().
    OutputParser uses Any for the response type since it works with any response.
    """

    @property
    def formatting_instructions(self) -> str | None:
        """The formatting instructions that will be added to the LLM system prompt.

        If the format has a custom `OutputParser`, its formatting instructions will be used.
        Otherwise, if the format type has a `formatting_instructions` class method,
        the output of that call will be used. Otherwise, instructions may be
        auto-generated based on the formatting mode.
        """
        if is_output_parser(self.formattable) or isinstance(
            self.formattable, HasFormattingInstructions
        ):
            return self.formattable.formatting_instructions()

        if self.mode == "tool":
            return TOOL_MODE_INSTRUCTIONS
        elif self.mode == "json":
            json_schema = json.dumps(self.schema, indent=2)
            instructions = JSON_MODE_INSTRUCTIONS.format(json_schema=json_schema)
            return inspect.cleandoc(instructions)
        elif self.mode == "parser":
            return None  # pragma: no cover

    def create_tool_schema(
        self,
    ) -> ToolSchema[ToolFn[..., None]]:
        """Generate a `ToolSchema` for parsing this format.

        Returns:
            `ToolSchema` for the format tool
        """

        schema_dict: dict[str, Any] = self.schema.copy()
        schema_dict["type"] = "object"

        properties = schema_dict.get("properties")
        if not properties or not isinstance(properties, dict):
            properties = {}  # pragma: no cover
        properties = cast(dict[str, Any], properties)
        required: list[str] = list(properties.keys())

        description = (
            f"Use this tool to extract data in {self.name} format for a final response."
        )
        if self.description:
            description += "\n" + self.description

        parameters = ToolParameterSchema(
            properties=properties,
            required=required,
            additionalProperties=False,
        )
        if "$defs" in schema_dict and isinstance(schema_dict["$defs"], dict):
            parameters.defs = schema_dict["$defs"]

        def _unused_format_fn() -> None:
            raise TypeError(
                "Format tool function should not be called."
            )  # pragma: no cover

        return ToolSchema(
            fn=cast(ToolFn[..., None], _unused_format_fn),
            name=FORMAT_TOOL_NAME,
            description=description,
            parameters=parameters,
            strict=None,  # Provider determines whether to use strict mode.
        )


def format(
    formattable: type[FormattableT] | OutputParser[FormattableT] | None,
    *,
    mode: FormattingMode,
) -> Format[FormattableT] | None:
    """Returns a `Format` that describes structured output or custom parsing.

    This function converts a Formattable type (e.g. Pydantic `BaseModel` or primitive type)
    or an `OutputParser` into a `Format` object that describes how the output should be
    formatted and parsed. Calling `llm.format` is optional, as all the APIs that expect
    a `Format` can also take the Formattable type or `OutputParser` directly. However,
    calling `llm.format` is necessary in order to specify the formatting mode for
    `BaseModel`/primitive types.

    Primitive types are automatically wrapped in a `BaseModel` with an "output" field
    for schema generation, then unwrapped during parsing.

    Args:
        formattable: The type or parser to format:
            - BaseModel type: Uses structured output with JSON schema
            - Primitive type: Wrapped in schema for structured output
            - OutputParser: Uses custom parsing with instructions
        mode: The format mode to use (required):
            - "strict": Use model strict structured outputs, or fail if unavailable.
            - "tool": Use forced tool calling with a special tool that represents a
              formatted response.
            - "json": Use provider json mode if available, or modify prompt to request
              json if not.
            - "parser": Must be used for OutputParser types.

    The Formattable type may provide custom formatting instructions via a
    `formatting_instructions(cls)` classmethod. If that method is present, it will be called,
    and the resulting instructions will automatically be appended to the system prompt.

    If no formatting instructions are present, then Mirascope may auto-generate instructions
    based on the active format mode. To disable this behavior and all prompt modification,
    you can add the `formatting_instructions` classmethod and have it return `None`.

    Returns:
      A `Format` object describing the format type or parser.

    Example:
      Using with a BaseModel:

      ```python
      from pydantic import BaseModel
      from mirascope import llm

      class Book(BaseModel):
          title: str
          author: str

      format = llm.format(Book, mode="strict")

      @llm.call("openai/gpt-5-mini", format=format)
      def recommend_book(genre: str):
          return f"Recommend a {genre} book."

      response = recommend_book("fantasy")
      book: Book = response.parse()
      ```

    Example:

      Using with an `OutputParser`:

      ```python
      @llm.output_parser(
          formatting_instructions="Return XML: <book><title>...</title></book>"
      )
      def parse_book_xml(response: llm.AnyResponse) -> Book:
          # ... parsing logic ...
          return Book(...)

      @llm.call("openai/gpt-5-mini", format=parse_book_xml)
      def recommend_book(genre: str):
          return f"Recommend a {genre} book."
      ```
    """
    # TODO: Add caching or memoization to this function (e.g. functools.lru_cache)

    if formattable is None or formattable is NoneType:
        return None

    if is_output_parser(formattable):
        if mode != "parser":
            raise ValueError(f"mode must be 'parser' for OutputParser, got '{mode}'")
        return Format[Any](
            name=formattable.__name__,
            description=formattable.__doc__,
            schema={},
            mode="parser",
            formattable=formattable,
        )

    if is_primitive_type(formattable):
        wrapper_model = create_wrapper_model(formattable)
        schema = wrapper_model.model_json_schema()
        name = (
            formattable.__name__
            if hasattr(formattable, "__name__")
            else str(formattable)
        )

        return Format[FormattableT](
            name=name,
            description=None,
            schema=schema,
            mode=mode,
            formattable=formattable,
        )

    description = None
    if formattable.__doc__:
        description = inspect.cleandoc(formattable.__doc__)

    schema = formattable.model_json_schema()

    return Format[FormattableT](
        name=formattable.__name__,
        description=description,
        schema=schema,
        mode=mode,
        formattable=formattable,
    )


def resolve_format(
    formattable: FormatSpec[FormattableT] | None,
    default_mode: FormattingMode,
) -> Format[FormattableT] | None:
    """Resolve a `Format` (or None) from a possible `Format`, Formattable, or `OutputParser`.

    Args:
        formattable: The format specification:
            - Format: Returned as-is
            - BaseModel/primitive type: Converted to Format with default_mode
            - OutputParser: Converted to Format with mode='parser'
        default_mode: The mode to use for BaseModel/primitive types.

    Returns:
        A Format object or None.
    """
    if isinstance(formattable, Format):
        return formattable

    if is_output_parser(formattable):
        return format(formattable, mode="parser")

    return format(formattable, mode=default_mode)
