from __future__ import annotations

import json

from pydantic import BaseModel

from safe_genai.exceptions import SafeGenAIFormatterError
from safe_genai.providers.base import BaseProvider
from safe_genai.result import Usage

_FORMATTER_SYSTEM_TEMPLATE = """\
You are a data formatter. Your job is to extract and coerce the provided text into the following JSON schema.
Return only valid JSON that conforms to the schema. Do not include any explanation.

Schema ({schema_name}):
{schema_json}
{extra}"""


class Formatter:
    def __init__(
        self,
        provider: BaseProvider,
        model: str,
        output_schema: type[BaseModel],
        additional_instructions: str | None = None,
    ) -> None:
        self._provider = provider
        self._model = model
        self._output_schema = output_schema
        self._system_prompt = _FORMATTER_SYSTEM_TEMPLATE.format(
            schema_name=output_schema.__name__,
            schema_json=json.dumps(output_schema.model_json_schema(), indent=2),
            extra=f"\nAdditional instructions:\n{additional_instructions}"
            if additional_instructions
            else "",
        )

    def format(self, raw_text: str) -> tuple[BaseModel, Usage]:
        try:
            return self._provider.generate_structured(
                self._model, self._system_prompt, raw_text, self._output_schema
            )
        except Exception as e:
            raise SafeGenAIFormatterError(raw_text, self._output_schema.__name__, e) from e

    async def aformat(self, raw_text: str) -> tuple[BaseModel, Usage]:
        try:
            return await self._provider.agenerate_structured(
                self._model, self._system_prompt, raw_text, self._output_schema
            )
        except Exception as e:
            raise SafeGenAIFormatterError(raw_text, self._output_schema.__name__, e) from e
