from typing import AsyncGenerator, List, Dict, Any

from openai import OpenAI

from analogai.foundation.settings import config as infra_config
from analogai.foundation.common.logging.app_logger import app_logger
from analogai.deepthink.infrastructure.gateways.llm.gateway import LlmGateway


class OllamaLlmGatewayImpl(LlmGateway):
    def __init__(
        self,
        system_prompt: str | None = None,
        max_tokens: int = 1024,          # raised — 300 is too restrictive nowadays
        model: str | None = None,
        temperature: float = 0.75,       # ← key change: was 1.0 → much safer
        top_p: float = 0.92,             # ← key change: was 1.0 → reduces blanks
    ):
        if not model:
            raise ValueError("Model name must be provided for Ollama gateway")
        self.model = model
        self.default_system_prompt = system_prompt or ""
        self.default_max_tokens = max_tokens
        self.default_temperature = temperature
        self.default_top_p = top_p

        # Common stop tokens for many models (prevents runaway or instant stop)
        self.default_stop = [
            "<|eot_id|>", "<|end_of_text|>", "<|im_end|>", "</s>",
            "\n\nUser:", "\nUser:", "User:"
        ]

        host = infra_config.get("OLLAMA_HOST", "http://localhost:11434")
        base_url = self._normalize_host(host)
        api_base = base_url if base_url.endswith("/v1") else f"{base_url.rstrip('/')}/v1"

        timeout_seconds = float(infra_config.get("LLM_TIMEOUT_SECONDS", 90))
        self.client = OpenAI(
            base_url=api_base,
            api_key="ollama",  # dummy — ignored by Ollama
            timeout=timeout_seconds,
        )
        self.keep_alive = infra_config.get("OLLAMA_KEEP_ALIVE", "10m")

    def set_model(self, model: str) -> None:
        self.model = model

    def set_system_prompt(self, system_prompt: str) -> None:
        self.default_system_prompt = system_prompt

    @staticmethod
    def _normalize_host(host: str) -> str:
        if not host:
            return "http://localhost:11434"
        normalized = host.strip()
        if normalized.startswith("0.0.0.0"):
            normalized = normalized.replace("0.0.0.0", "localhost", 1)
        if not normalized.startswith(("http://", "https://")):
            normalized = f"http://{normalized}"
        return normalized

    def _build_messages(self, user_message: str) -> List[Dict[str, str]]:
        messages: List[Dict[str, str]] = []
        if self.default_system_prompt:
            messages.append({"role": "system", "content": self.default_system_prompt})
        messages.append({"role": "user", "content": user_message})
        return messages

    @staticmethod
    def _estimate_prompt_usage(messages: List[Dict[str, str]]) -> tuple[int, int]:
        total_chars = sum(len(msg.get("content") or "") for msg in messages)
        estimated_tokens = max(1, total_chars // 4) if total_chars else 0
        return total_chars, estimated_tokens

    @staticmethod
    def _extract_content(obj: Any) -> str | None:
        """Very defensive — handles null, [], "", whitespace, reasoning field, etc."""
        choices = getattr(obj, "choices", None)
        if not choices or len(choices) == 0:
            return None

        choice = choices[0]
        if not hasattr(choice, "message"):
            return None

        msg = choice.message
        content = getattr(msg, "content", None)

        if content is None:
            pass
        elif isinstance(content, list):  # seen in some Ollama bugs
            content = " ".join(str(item).strip() for item in content if item)
        elif not isinstance(content, str):
            content = str(content)

        if isinstance(content, str):
            stripped = content.strip()
            if stripped:
                return stripped

        # Fallback for reasoning models (e.g. some Nemotron / experimental variants)
        if hasattr(msg, "reasoning") and msg.reasoning:
            reasoning = str(msg.reasoning).strip()
            if reasoning:
                app_logger.info("[Ollama] Extracted from 'reasoning' field")
                return reasoning

        return None

    def generate_completion(self, user_message: str) -> str:
        messages = self._build_messages(user_message)
        max_retries = int(infra_config.get("OLLAMA_EMPTY_RETRIES", 2))

        total_chars, est_tokens = self._estimate_prompt_usage(messages)
        app_logger.debug(
            "[OllamaGateway] prompt_usage chars=%d est_tokens=%d model=%s",
            total_chars, est_tokens, self.model
        )

        for attempt in range(max_retries + 1):
            try:
                response = self.client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                    max_tokens=self.default_max_tokens,
                    temperature=self.default_temperature,
                    top_p=self.default_top_p,
                    stop=self.default_stop,
                    extra_body={"keep_alive": self.keep_alive},
                    stream=False,
                )
                content = self._extract_content(response)
                if content:
                    return content

                app_logger.warning(
                    "[OllamaGateway] Empty content attempt %d/%d | model=%s | resp=%r",
                    attempt + 1, max_retries + 1, self.model, response
                )

            except Exception as exc:
                app_logger.error(
                    "[OllamaGateway] Completion error attempt %d: %s",
                    attempt + 1, exc, exc_info=True
                )
                if attempt == max_retries:
                    raise

        raise RuntimeError(f"Ollama gave empty/invalid response after {max_retries + 1} tries")

    async def generate_streaming_completion(
        self, user_message: str
    ) -> AsyncGenerator[str, None]:
        messages = self._build_messages(user_message)
        max_retries = int(infra_config.get("OLLAMA_EMPTY_RETRIES", 1))
        collected = ""

        for attempt in range(max_retries + 1):
            try:
                stream = self.client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                    max_tokens=self.default_max_tokens,
                    temperature=self.default_temperature,
                    top_p=self.default_top_p,
                    stop=self.default_stop,
                    extra_body={"keep_alive": self.keep_alive},
                    stream=True,
                )

                async for chunk in stream:  # note: assumes async iteration works; if sync client → wrap
                    content = self._extract_content(chunk)
                    if content:
                        collected += content
                        yield content

                if collected.strip():
                    return

                app_logger.warning(
                    "[OllamaGateway] Streaming empty after attempt %d/%d",
                    attempt + 1, max_retries + 1
                )

            except Exception as exc:
                app_logger.error(
                    "[OllamaGateway] Streaming error attempt %d: %s",
                    attempt + 1, exc, exc_info=True
                )

        raise RuntimeError("Streaming yielded no content after retries")

if __name__ == "__main__":
    gateway = OllamaLlmGatewayImpl(model="llama3.1")
    try:
        result = gateway.generate_completion("Hello, how are you?")
        print(result)
    except Exception as e:
        print("Error:", e)