"""LLM provider abstraction — BaseLLM ABC + concrete implementations + factory."""

from abc import ABC, abstractmethod
from enum import StrEnum
from typing import AsyncIterator

import structlog
from pydantic import BaseModel, field_validator

log = structlog.get_logger()


class LLMProvider(StrEnum):
    ANTHROPIC = "anthropic"
    OPENAI = "openai"
    GEMINI = "gemini"


class LLMConfig(BaseModel):
    provider: LLMProvider
    api_key: str
    model: str | None = None
    max_tokens: int = 4096
    temperature: float = 0.7

    @field_validator("api_key")
    @classmethod
    def api_key_not_empty(cls, v: str) -> str:
        if not v.strip():
            raise ValueError("api_key must not be empty")
        return v


class ToolCall(BaseModel):
    """A tool call extracted from the LLM response."""

    id: str
    name: str
    input: dict


class LLMResponse(BaseModel):
    content: str
    model: str
    provider: LLMProvider
    usage: dict | None = None
    tool_calls: list[ToolCall] = []
    stop_reason: str = "end_turn"
    raw_content: list[dict] | None = None


# Default models per provider
DEFAULT_MODELS = {
    LLMProvider.ANTHROPIC: "claude-sonnet-4-5-20250929",
    LLMProvider.OPENAI: "gpt-4o",
    LLMProvider.GEMINI: "gemini-2.0-flash",
}


class BaseLLM(ABC):
    def __init__(self, config: LLMConfig):
        self.config = config
        self.model = config.model or DEFAULT_MODELS[config.provider]

    @abstractmethod
    async def generate(
        self, messages: list[dict], system: str | None = None, tools: list[dict] | None = None
    ) -> LLMResponse: ...

    @abstractmethod
    async def generate_stream(self, messages: list[dict], system: str | None = None) -> AsyncIterator[str]: ...


class AnthropicLLM(BaseLLM):
    def __init__(self, config: LLMConfig):
        super().__init__(config)
        import anthropic

        self._client = anthropic.AsyncAnthropic(api_key=config.api_key)

    async def generate(
        self, messages: list[dict], system: str | None = None, tools: list[dict] | None = None
    ) -> LLMResponse:
        kwargs = {
            "model": self.model,
            "max_tokens": self.config.max_tokens,
            "temperature": self.config.temperature,
            "messages": messages,
        }
        if system:
            kwargs["system"] = system
        if tools:
            kwargs["tools"] = tools

        response = await self._client.messages.create(**kwargs)

        # Parse content blocks: extract text and tool_use calls
        text_parts = []
        tool_calls = []
        raw_content = []
        for block in response.content:
            if block.type == "text":
                text_parts.append(block.text)
                raw_content.append({"type": "text", "text": block.text})
            elif block.type == "tool_use":
                tool_calls.append(ToolCall(id=block.id, name=block.name, input=block.input))
                raw_content.append({"type": "tool_use", "id": block.id, "name": block.name, "input": block.input})

        return LLMResponse(
            content="".join(text_parts),
            model=response.model,
            provider=LLMProvider.ANTHROPIC,
            usage={"input_tokens": response.usage.input_tokens, "output_tokens": response.usage.output_tokens},
            tool_calls=tool_calls,
            stop_reason=response.stop_reason,
            raw_content=raw_content,
        )

    async def generate_stream(self, messages: list[dict], system: str | None = None) -> AsyncIterator[str]:
        kwargs = {
            "model": self.model,
            "max_tokens": self.config.max_tokens,
            "temperature": self.config.temperature,
            "messages": messages,
        }
        if system:
            kwargs["system"] = system

        async with self._client.messages.stream(**kwargs) as stream:
            async for text in stream.text_stream:
                yield text


class OpenAILLM(BaseLLM):
    def __init__(self, config: LLMConfig):
        super().__init__(config)
        import openai

        self._client = openai.AsyncOpenAI(api_key=config.api_key)

    async def generate(
        self, messages: list[dict], system: str | None = None, tools: list[dict] | None = None
    ) -> LLMResponse:
        if tools:
            log.warning("openai_tools_not_implemented", msg="tools param accepted but not yet wired for OpenAI")

        msgs = []
        if system:
            msgs.append({"role": "system", "content": system})
        msgs.extend(messages)

        response = await self._client.chat.completions.create(
            model=self.model,
            messages=msgs,
            max_tokens=self.config.max_tokens,
            temperature=self.config.temperature,
        )
        choice = response.choices[0]
        usage = None
        if response.usage:
            usage = {"input_tokens": response.usage.prompt_tokens, "output_tokens": response.usage.completion_tokens}

        return LLMResponse(
            content=choice.message.content,
            model=response.model,
            provider=LLMProvider.OPENAI,
            usage=usage,
        )

    async def generate_stream(self, messages: list[dict], system: str | None = None) -> AsyncIterator[str]:
        msgs = []
        if system:
            msgs.append({"role": "system", "content": system})
        msgs.extend(messages)

        stream = await self._client.chat.completions.create(
            model=self.model,
            messages=msgs,
            max_tokens=self.config.max_tokens,
            temperature=self.config.temperature,
            stream=True,
        )
        async for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                yield chunk.choices[0].delta.content


class GeminiLLM(BaseLLM):
    def __init__(self, config: LLMConfig):
        super().__init__(config)
        from google import genai

        self._client = genai.Client(api_key=config.api_key)

    async def generate(
        self, messages: list[dict], system: str | None = None, tools: list[dict] | None = None
    ) -> LLMResponse:
        if tools:
            log.warning("gemini_tools_not_implemented", msg="tools param accepted but not yet wired for Gemini")

        from google.genai import types

        contents = []
        for msg in messages:
            role = "model" if msg["role"] == "assistant" else "user"
            contents.append(types.Content(role=role, parts=[types.Part(text=msg["content"])]))

        config = types.GenerateContentConfig(
            max_output_tokens=self.config.max_tokens,
            temperature=self.config.temperature,
        )
        if system:
            config.system_instruction = system

        response = await self._client.aio.models.generate_content(
            model=self.model,
            contents=contents,
            config=config,
        )

        usage = None
        if response.usage_metadata:
            usage = {
                "input_tokens": response.usage_metadata.prompt_token_count,
                "output_tokens": response.usage_metadata.candidates_token_count,
            }

        return LLMResponse(
            content=response.text,
            model=self.model,
            provider=LLMProvider.GEMINI,
            usage=usage,
        )

    async def generate_stream(self, messages: list[dict], system: str | None = None) -> AsyncIterator[str]:
        from google.genai import types

        contents = []
        for msg in messages:
            role = "model" if msg["role"] == "assistant" else "user"
            contents.append(types.Content(role=role, parts=[types.Part(text=msg["content"])]))

        config = types.GenerateContentConfig(
            max_output_tokens=self.config.max_tokens,
            temperature=self.config.temperature,
        )
        if system:
            config.system_instruction = system

        async for chunk in await self._client.aio.models.generate_content_stream(
            model=self.model,
            contents=contents,
            config=config,
        ):
            if chunk.text:
                yield chunk.text


_PROVIDER_MAP = {
    LLMProvider.ANTHROPIC: AnthropicLLM,
    LLMProvider.OPENAI: OpenAILLM,
    LLMProvider.GEMINI: GeminiLLM,
}


def create_llm(config: LLMConfig) -> BaseLLM:
    """Factory: instantiate the correct LLM provider from config."""
    cls = _PROVIDER_MAP.get(config.provider)
    if not cls:
        raise ValueError(f"Unknown LLM provider: {config.provider}")
    return cls(config)
