"""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
    base_url: str | None = None

    @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",
}


def resolve_model_alias(model: str) -> str:
    """Resolve user-defined model alias to full model ID. Passthrough if not an alias."""
    from fliiq.runtime.config import global_fliiq_dir

    models_file = global_fliiq_dir() / "models.yaml"
    if not models_file.is_file():
        return model
    try:
        import yaml

        data = yaml.safe_load(models_file.read_text()) or {}
        aliases = data.get("aliases", {})
        return aliases.get(model, model)
    except Exception:
        return model


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

        kwargs: dict = {"api_key": config.api_key}
        if config.base_url:
            kwargs["base_url"] = config.base_url
        self._client = anthropic.AsyncAnthropic(**kwargs)

    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

        kwargs: dict = {"api_key": config.api_key}
        if config.base_url:
            kwargs["base_url"] = config.base_url
        self._client = openai.AsyncOpenAI(**kwargs)

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

        # --- Convert messages: Anthropic format → OpenAI format ---
        msgs: list[dict] = []
        if system:
            msgs.append({"role": "system", "content": system})

        for msg in messages:
            role = msg["role"]
            content = msg.get("content", "")

            if role == "assistant" and isinstance(content, list):
                text_parts = []
                oai_tool_calls = []
                for block in content:
                    if block.get("type") == "text":
                        text_parts.append(block["text"])
                    elif block.get("type") == "tool_use":
                        oai_tool_calls.append({
                            "id": block["id"],
                            "type": "function",
                            "function": {
                                "name": block["name"],
                                "arguments": json.dumps(block["input"]),
                            },
                        })
                oai_msg: dict = {"role": "assistant", "content": "".join(text_parts) or None}
                if oai_tool_calls:
                    oai_msg["tool_calls"] = oai_tool_calls
                msgs.append(oai_msg)

            elif role == "user" and isinstance(content, list):
                for block in content:
                    if block.get("type") == "tool_result":
                        msgs.append({
                            "role": "tool",
                            "tool_call_id": block["tool_use_id"],
                            "content": block.get("content", ""),
                        })

            else:
                msgs.append({"role": role, "content": content})

        # --- Convert tools: Anthropic format → OpenAI format ---
        oai_tools = None
        if tools:
            oai_tools = [
                {
                    "type": "function",
                    "function": {
                        "name": t["name"],
                        "description": t.get("description", ""),
                        "parameters": t.get("input_schema", {}),
                    },
                }
                for t in tools
            ]

        # --- API call ---
        kwargs: dict = {
            "model": self.model,
            "messages": msgs,
            "max_tokens": self.config.max_tokens,
            "temperature": self.config.temperature,
        }
        if oai_tools:
            kwargs["tools"] = oai_tools

        response = await self._client.chat.completions.create(**kwargs)
        choice = response.choices[0]

        # --- Convert response: OpenAI format → Anthropic format ---
        text = choice.message.content or ""
        tool_calls: list[ToolCall] = []
        raw_content: list[dict] = []

        if text:
            raw_content.append({"type": "text", "text": text})

        if choice.message.tool_calls:
            for tc in choice.message.tool_calls:
                try:
                    tool_input = json.loads(tc.function.arguments)
                except (json.JSONDecodeError, TypeError):
                    tool_input = {}
                tool_calls.append(ToolCall(id=tc.id, name=tc.function.name, input=tool_input))
                raw_content.append({
                    "type": "tool_use",
                    "id": tc.id,
                    "name": tc.function.name,
                    "input": tool_input,
                })

        stop_reason = "tool_use" if choice.finish_reason == "tool_calls" else "end_turn"

        usage = None
        if response.usage:
            usage = {"input_tokens": response.usage.prompt_tokens, "output_tokens": response.usage.completion_tokens}

        return LLMResponse(
            content=text,
            model=response.model,
            provider=LLMProvider.OPENAI,
            usage=usage,
            tool_calls=tool_calls,
            stop_reason=stop_reason,
            raw_content=raw_content if raw_content else None,
        )

    async def generate_stream(self, messages: list[dict], system: str | None = None) -> AsyncIterator[str]:
        # Convert messages — strip tool blocks for streaming
        msgs: list[dict] = []
        if system:
            msgs.append({"role": "system", "content": system})

        for msg in messages:
            role = msg["role"]
            content = msg.get("content", "")

            if role == "assistant" and isinstance(content, list):
                text = "".join(b.get("text", "") for b in content if b.get("type") == "text")
                if text:
                    msgs.append({"role": "assistant", "content": text})
            elif role == "user" and isinstance(content, list):
                # Skip tool_result blocks in streaming
                continue
            else:
                msgs.append({"role": role, "content": content})

        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


def _find_tool_name(messages: list[dict], tool_use_id: str) -> str:
    """Find the tool name for a given tool_use_id by scanning assistant messages.

    Gemini FunctionResponse requires the function name, but Anthropic-format
    tool results only carry tool_use_id. This scans backward to find the match.
    """
    for msg in reversed(messages):
        if msg.get("role") != "assistant":
            continue
        content = msg.get("content", [])
        if not isinstance(content, list):
            continue
        for block in content:
            if block.get("type") == "tool_use" and block.get("id") == tool_use_id:
                return block["name"]
    return "unknown"


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:
        import uuid

        from google.genai import types

        # --- Convert messages: Anthropic format → Gemini format ---
        contents = []
        for msg in messages:
            role = "model" if msg["role"] == "assistant" else "user"
            content = msg.get("content", "")

            if msg["role"] == "assistant" and isinstance(content, list):
                parts = []
                for block in content:
                    if block.get("type") == "text" and block.get("text"):
                        parts.append(types.Part.from_text(text=block["text"]))
                    elif block.get("type") == "tool_use":
                        parts.append(types.Part.from_function_call(
                            name=block["name"],
                            args=block["input"],
                        ))
                if parts:
                    contents.append(types.Content(role="model", parts=parts))

            elif msg["role"] == "user" and isinstance(content, list):
                parts = []
                for block in content:
                    if block.get("type") == "tool_result":
                        tool_name = _find_tool_name(messages, block["tool_use_id"])
                        parts.append(types.Part.from_function_response(
                            name=tool_name,
                            response={"result": block.get("content", "")},
                        ))
                if parts:
                    contents.append(types.Content(role="user", parts=parts))

            else:
                text = content if isinstance(content, str) else str(content)
                if text:
                    contents.append(types.Content(role=role, parts=[types.Part.from_text(text=text)]))

        # --- Convert tools: Anthropic format → Gemini format ---
        gemini_tools = None
        if tools:
            func_decls = []
            for t in tools:
                schema = dict(t.get("input_schema", {}))
                schema.pop("additionalProperties", None)
                func_decls.append(types.FunctionDeclaration(
                    name=t["name"],
                    description=t.get("description", ""),
                    parameters=schema if schema else None,
                ))
            gemini_tools = [types.Tool(function_declarations=func_decls)]

        # --- API call ---
        gen_config = types.GenerateContentConfig(
            max_output_tokens=self.config.max_tokens,
            temperature=self.config.temperature,
        )
        if system:
            gen_config.system_instruction = system
        if gemini_tools:
            gen_config.tools = gemini_tools

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

        # --- Convert response: Gemini format → Anthropic format ---
        text_parts: list[str] = []
        tool_calls: list[ToolCall] = []
        raw_content: list[dict] = []

        if response.candidates and response.candidates[0].content:
            for part in response.candidates[0].content.parts:
                if part.text:
                    text_parts.append(part.text)
                    raw_content.append({"type": "text", "text": part.text})
                elif part.function_call:
                    fc = part.function_call
                    tc_id = f"toolu_{uuid.uuid4().hex[:12]}"
                    args = dict(fc.args) if fc.args else {}
                    tool_calls.append(ToolCall(id=tc_id, name=fc.name, input=args))
                    raw_content.append({
                        "type": "tool_use",
                        "id": tc_id,
                        "name": fc.name,
                        "input": args,
                    })

        stop_reason = "tool_use" if tool_calls else "end_turn"

        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="".join(text_parts),
            model=self.model,
            provider=LLMProvider.GEMINI,
            usage=usage,
            tool_calls=tool_calls,
            stop_reason=stop_reason,
            raw_content=raw_content if raw_content else None,
        )

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

        # Convert messages — strip tool blocks for streaming
        contents = []
        for msg in messages:
            role = "model" if msg["role"] == "assistant" else "user"
            content = msg.get("content", "")

            if isinstance(content, list):
                parts = []
                for block in content:
                    if block.get("type") == "text" and block.get("text"):
                        parts.append(types.Part.from_text(text=block["text"]))
                if parts:
                    contents.append(types.Content(role=role, parts=parts))
            else:
                text = content if isinstance(content, str) else str(content)
                if text:
                    contents.append(types.Content(role=role, parts=[types.Part.from_text(text=text)]))

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

        async for chunk in await self._client.aio.models.generate_content_stream(
            model=self.model,
            contents=contents,
            config=gen_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)
