"""Integration with Dottxt's API."""

import json
from typing import TYPE_CHECKING, Any, Optional

from pydantic import TypeAdapter

from outlines.models.base import Model, ModelTypeAdapter
from outlines.types import CFG, JsonSchema, Regex
from outlines.types.utils import (
    is_dataclass,
    is_genson_schema_builder,
    is_pydantic_model,
    is_typed_dict,
)

if TYPE_CHECKING:
    from dottxt import Dottxt as DottxtClient

__all__ = ["Dottxt", "from_dottxt"]


class DottxtTypeAdapter(ModelTypeAdapter):
    """Type adapter for the `Dottxt` model."""

    def format_input(self, model_input: str) -> str:
        """Format the prompt to pass to the client.

        Parameters
        ----------
        model_input
            The input provided by the user.

        Returns
        -------
        str
            The input to pass to the client.

        """
        if isinstance(model_input, str):
            return model_input
        raise TypeError(
            f"The input type {model_input} is not available with Dottxt. "
            "The only available type is `str`."
        )

    def format_output_type(self, output_type: Optional[Any] = None) -> str:
        """Format the output type to pass to the client.

        TODO: `int`, `float` and other Python types could be supported via
        JSON Schema.

        Parameters
        ----------
        output_type
            The output type provided by the user.

        Returns
        -------
        str
            The output type to pass to the client.

        """
        # Unsupported languages
        if output_type is None:
            raise TypeError(
                "You must provide an output type. Dottxt only supports "
                "constrained generation."
            )
        elif isinstance(output_type, Regex):
            raise TypeError(
                "Regex-based structured outputs will soon be available with "
                "Dottxt. Use an open source model in the meantime."
            )
        elif isinstance(output_type, CFG):
            raise TypeError(
                "CFG-based structured outputs will soon be available with "
                "Dottxt. Use an open source model in the meantime."
            )

        elif isinstance(output_type, JsonSchema):
            return output_type.schema
        elif is_dataclass(output_type):
            schema = TypeAdapter(output_type).json_schema()
            return json.dumps(schema)
        elif is_typed_dict(output_type):
            schema = TypeAdapter(output_type).json_schema()
            return json.dumps(schema)
        elif is_pydantic_model(output_type):
            schema = output_type.model_json_schema()
            return json.dumps(schema)
        elif is_genson_schema_builder(output_type):
            return output_type.to_json()
        else:
            type_name = getattr(output_type, "__name__", output_type)
            raise TypeError(
                f"The type `{type_name}` is not supported by Dottxt. "
                "Consider using a local mode instead."
            )


class Dottxt(Model):
    """Thin wrapper around the `dottxt.client.Dottxt` client.

    This wrapper is used to convert the input and output types specified by the
    users at a higher level to arguments to the `dottxt.client.Dottxt` client.

    """

    def __init__(
        self,
        client: "DottxtClient",
        model_name: Optional[str] = None,
        model_revision: Optional[str] = None,
    ):
        """
        Parameters
        ----------
        client
            A `dottxt.Dottxt` client.
        model_name
            The name of the model to use.
        model_revision
            The revision of the model to use.

        """
        self.client = client
        self.model_name = model_name
        self.model_revision = model_revision
        self.type_adapter = DottxtTypeAdapter()

    def generate(
        self,
        model_input: str,
        output_type: Optional[Any] = None,
        **inference_kwargs: Any,
    ) -> str:
        """Generate text using Dottxt.

        Parameters
        ----------
        model_input
            The prompt based on which the model will generate a response.
        output_type
            The desired format of the response generated by the model. The
            output type must be of a type that can be converted to a JSON
            schema.
        **inference_kwargs
            Additional keyword arguments to pass to the client.

        Returns
        -------
        str
            The text generated by the model.

        """
        prompt = self.type_adapter.format_input(model_input)
        json_schema = self.type_adapter.format_output_type(output_type)

        if (
            "model_name" not in inference_kwargs
            and self.model_name is not None
        ):
            inference_kwargs["model_name"] = self.model_name

        if (
            "model_revision" not in inference_kwargs
            and self.model_revision is not None
        ):
            inference_kwargs["model_revision"] = self.model_revision

        completion = self.client.json(
            prompt,
            json_schema,
            **inference_kwargs,
        )
        return completion.data

    def generate_batch(
        self,
        model_input,
        output_type = None,
        **inference_kwargs,
    ):
        raise NotImplementedError(
            "Dottxt does not support batch generation."
        )

    def generate_stream(
        self,
        model_input,
        output_type=None,
        **inference_kwargs,
    ):
        """Not available for Dottxt."""
        raise NotImplementedError(
            "Dottxt does not support streaming. Call the model/generator for "
            + "regular generation instead."
        )


def from_dottxt(
    client: "DottxtClient",
    model_name: Optional[str] = None,
    model_revision: Optional[str] = None,
) -> Dottxt:
    """Create an Outlines `Dottxt` model instance from a `dottxt.Dottxt`
    client instance.

    Parameters
    ----------
    client
        A `dottxt.Dottxt` client instance.
    model_name
        The name of the model to use.
    model_revision
        The revision of the model to use.

    Returns
    -------
    Dottxt
        An Outlines `Dottxt` model instance.

    """
    return Dottxt(client, model_name, model_revision)