"""This module provides an extremely high-level interface to LiteLLM,
with support for streaming text output using LocalTextStream.
"""

import json
from typing import Optional

import litellm
from pydantic import BaseModel, Field

litellm.drop_params = True


class LLMConfig(BaseModel):
    """Configuration for the LLM."""

    model: str = Field(
        default="openai/gpt-5.1",
        description="The full-sized LLM model to use for formatting",
    )
    small_model: str = Field(
        default="openai/gpt-5.1-mini",
        description="The smaller LLM model to use for faster cleaning / processing",
    )
    temperature: float = Field(
        default=1, description="The temperature to use for the LLM"
    )
    seed: int = Field(default=42, description="The seed to use for the LLM")


CLEAN_PROMPT = """Below is some text from a document or webpage. Clean it up by removing HTML tags, boilerplate, artifacts, and obvious copy/paste noise. Return only the text that would be read aloud by a speaker. Preserve the original wording as closely as possible otherwise."""

FORMAT_PROMPT = """You are given cleaned text that should be transformed into a natural spoken transcript suitable for text-to-speech. Rewrite the text so it sounds like a fluent spoken narration while preserving meaning and key facts. Use natural phrasing, add appropriate sentence breaks, and remove web/markdown artifacts. Do not add new facts. Return only the final transcript."""


class Response(BaseModel):
    """Response from the LLM."""

    cleaned_text: str = Field(description="The cleaned or formatted text")


class LiteLLM:
    """LiteLLM class for interacting with LLMs via LiteLLM."""

    def __init__(self, config: Optional[LLMConfig] = None) -> None:
        self.config = config or LLMConfig()

    async def clean_text(
        self,
        text: str,
        model: str | None = None,
        temperature: float | None = None,
        seed: int | None = None,
    ) -> str:
        """Clean the text for the LLM."""
        response = await litellm.acompletion(
            model=model or self.config.small_model,
            messages=[
                {"role": "system", "content": CLEAN_PROMPT},
                {"role": "user", "content": text},
            ],
            temperature=temperature or self.config.temperature,
            seed=seed or self.config.seed,
            response_format=Response,
        )

        if not response.choices or not response.choices[0].message.content:
            raise RuntimeError("Could not parse response from LLM")

        return json.loads(response.choices[0].message.content)["cleaned_text"]

    async def format_text(
        self,
        text: str,
        model: str | None = None,
        temperature: float | None = None,
        seed: int | None = None,
    ) -> str:
        """Format the text into a natural spoken transcript."""
        response = await litellm.acompletion(
            model=model or self.config.model,
            messages=[
                {"role": "system", "content": FORMAT_PROMPT},
                {"role": "user", "content": text},
            ],
            temperature=temperature or self.config.temperature,
            seed=seed or self.config.seed,
            response_format=Response,
        )

        if not response.choices or not response.choices[0].message.content:
            raise RuntimeError("Could not parse response from LLM")

        return json.loads(response.choices[0].message.content)["cleaned_text"]
