from collections.abc import Iterator

from ollama import Client, ResponseError

from yera.config2.dataclasses import LLMConfig
from yera.config2.keyring import DevKeyring
from yera.models.data_classes import Message
from yera.models.llm_interfaces.base import BaseLLMInterface, TBaseModel


class OllamaLLM(BaseLLMInterface):
    def __init__(self, config: LLMConfig, overrides, creds_map):
        kws = {**config.inference.model_dump(), **overrides}
        self.client = None
        self.model_id = config.model_id
        self.creds_map = creds_map
        super().__init__(**kws)

    def start(self):
        base_url = DevKeyring.get(f"providers.{self.creds_map['base_url']}")
        self.client = Client(host=base_url)

        # Check if Ollama server is running
        try:
            self.client.list()  # Simple check to see if server responds
        except Exception as e:
            raise ConnectionError(
                f"Cannot connect to Ollama server at {base_url}. "
                f"Make sure Ollama is installed and running. "
                f"Error: {e}"
            ) from None

    def stop(self):
        self.client = None

    def chat(
        self,
        messages: list[Message],
        **ollama_kw,
    ) -> Iterator[str]:
        ollama_messages = [
            {"role": msg.role, "content": msg.content} for msg in messages
        ]

        try:
            stream = self.client.chat(
                model=self.model_id,
                messages=ollama_messages,
                stream=True,
                **ollama_kw,
            )

            in_thought_block = False
            for chunk in stream:
                thought_token = chunk.get("message", {}).get("thinking")
                if thought_token:
                    if not in_thought_block:
                        yield "\n<thinking>\n"
                        in_thought_block = True
                    yield thought_token
                content_token = chunk.get("message", {}).get("content")
                if content_token:
                    if in_thought_block:
                        yield "\n<thinking>\n\n"
                        in_thought_block = False

                    yield content_token

        except ResponseError as e:
            if e.status_code == 404:
                raise ValueError(
                    f"Model '{self.model_id}' not found. "
                    f"Pull it with: ollama pull {self.model_id}"
                ) from e
            raise

    def make_struct(
        self,
        messages: list[Message],
        cls: type[TBaseModel],
        **ollama_kw,
    ) -> Iterator[str]:
        ollama_messages = [
            {"role": msg.role, "content": msg.content} for msg in messages
        ]

        try:
            # Use Pydantic's model_json_schema() for structured outputs
            stream = self.client.chat(
                model=self.model_id,
                messages=ollama_messages,
                format=cls.model_json_schema(),
                stream=True,
                **ollama_kw,
            )

            for chunk in stream:
                if not chunk.get("message"):
                    continue

                msg = chunk["message"]

                # For structured outputs, thinking is usually separate
                # but we will yield it if present once it won't bugger the JSON.
                # if msg.get("thinking"):
                #    yield msg["thinking"]

                # Yield JSON content
                if msg.get("content"):
                    yield msg["content"]

        except ResponseError as e:
            if e.status_code == 404:
                raise ValueError(
                    f"Model '{self.model_id}' not found. "
                    f"Pull it with: ollama pull {self.model_id}"
                ) from e
            raise
