"""LLM base class and concrete implementations (OpenAI, OpenRouter)."""

import os
import time
from typing import Any, Dict, List

from openai import OpenAI

# Retry API calls on failure (e.g. rate limit, transient error, network).
MAX_RETRIES = 4
RETRY_DELAY_SECONDS = 3


def _chat_once(client: OpenAI, **create_kwargs: Any) -> str:
    r = client.chat.completions.create(**create_kwargs)
    return (r.choices[0].message.content or "").strip()


def _chat_with_retry(client: OpenAI, **create_kwargs: Any) -> str:
    last_exc = None
    for attempt in range(MAX_RETRIES):
        try:
            return _chat_once(client, **create_kwargs)
        except Exception as e:
            last_exc = e
            if attempt < MAX_RETRIES - 1:
                time.sleep(RETRY_DELAY_SECONDS)
    raise last_exc


class LLM:
    """Base class for any LLM. Subclass and implement chat(messages)."""

    def chat(self, messages: List[Dict[str, Any]]) -> str:
        """
        Override this.
        messages: list of {"role": "user"|"system"|"assistant", "content": "..."}.
        Returns: the assistant reply as a string.
        """
        raise NotImplementedError("Subclass must implement chat(messages)")


class OpenAILLM(LLM):
    """Uses OpenAI API (or any OpenAI-compatible endpoint)."""

    def __init__(
        self,
        model: str,
        api_key: str | None = None,
        base_url: str | None = None,
        temperature: float = 0.7,
    ) -> None:
        self.model = model
        self.temperature = temperature
        kwargs = {"api_key": api_key or os.environ.get("OPENAI_API_KEY", "")}
        if base_url is not None:
            kwargs["base_url"] = base_url
        self._client = OpenAI(**kwargs)

    def chat(self, messages: List[Dict[str, Any]]) -> str:
        return _chat_with_retry(
            self._client,
            model=self.model,
            messages=[{"role": m["role"], "content": m["content"]} for m in messages],
            temperature=self.temperature,
        )


class OpenRouterLLM(LLM):
    """Uses OpenRouter API (https://openrouter.ai). Same as OpenAI client with different base_url."""

    BASE_URL = "https://openrouter.ai/api/v1"

    def __init__(
        self,
        model: str,
        api_key: str | None = None,
        temperature: float = 0.7,
    ) -> None:
        self.model = model
        self.temperature = temperature
        self._client = OpenAI(
            api_key=api_key or os.environ.get("OPENROUTER_API_KEY", ""),
            base_url=self.BASE_URL,
        )

    def chat(self, messages: List[Dict[str, Any]]) -> str:
        return _chat_with_retry(
            self._client,
            model=self.model,
            messages=[{"role": m["role"], "content": m["content"]} for m in messages],
            temperature=self.temperature,
        )
