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

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 .types import 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]
    """The `Formattable` type that this `Format` describes.
    
    While the `FormattbleT` typevar allows for `None`, a `Format` will never be
    constructed when the `FormattableT` is `None`, so you may treat this as 
    a `RequiredFormattableT` in practice.
    """

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

        If the format type has a `formatting_instructions` class method, the output of that
        call will be used for instructions. Otherwise, instructions may be auto-generated
        based on the formatting mode.
        """
        if 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)

    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=True,
        )


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

    This function converts a Formattable type (e.g. Pydantic BaseModel) into a `Format`
    object that describes how the object should be formatted. Calling `llm.format`
    is optional, as all the APIs that expect a `Format` can also take the Formattable
    type directly. However, calling `llm.format` is necessary in order to specify the
    formatting mode that will be used.

    Args:
        mode: The format mode to use, one of the following:
            - "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.

    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 Formattable type.

    Example:
      Using with an LLM call:

      ```python
      from pydantic import BaseModel

      from mirascope import llm


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

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

      @llm.call(
          provider_id="openai",
          model_id="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()
      print(f"{book.title} by {book.author}")
      ```
    """
    # TODO: Add caching or memoization to this function (e.g. functools.lru_cache)

    if formattable is None or formattable is NoneType:
        return None

    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: type[FormattableT] | Format[FormattableT] | None,
    default_mode: FormattingMode,
) -> Format[FormattableT] | None:
    """Resolve a `Format` (or None) from a possible `Format` or Formattable."""
    if isinstance(formattable, Format):
        return formattable
    else:
        return format(formattable, mode=default_mode)
