import json
import logging
from abc import ABC, abstractmethod
from collections.abc import Callable, Coroutine
from typing import Any

import anthropic
import httpx

from bot.config import settings

log = logging.getLogger(__name__)

ProgressCallback = Callable[[str], Coroutine[Any, Any, None]]


class AIClient(ABC):
    @abstractmethod
    async def complete(
        self,
        system: str,
        user: str,
        max_tokens: int,
        on_progress: ProgressCallback | None = None,
    ) -> str: ...


class AnthropicClient(AIClient):
    def __init__(self, model_id: str) -> None:
        self.model_id = model_id
        self.client = anthropic.AsyncAnthropic(api_key=settings.anthropic_api_key)

    async def complete(
        self,
        system: str,
        user: str,
        max_tokens: int,
        on_progress: ProgressCallback | None = None,
    ) -> str:
        kwargs: dict[str, Any] = {
            "model": self.model_id,
            "max_tokens": max_tokens,
            "system": system,
            "messages": [{"role": "user", "content": user}],
        }
        if "opus-4-6" in self.model_id:
            kwargs["thinking"] = {"type": "adaptive"}

        async with self.client.messages.stream(**kwargs) as stream:
            response = await stream.get_final_message()

        return "".join(
            block.text for block in response.content if block.type == "text"
        )


class OpenRouterClient(AIClient):
    BASE_URL = "https://openrouter.ai/api/v1/chat/completions"

    def __init__(self, model_id: str) -> None:
        self.model_id = model_id

    async def complete(
        self,
        system: str,
        user: str,
        max_tokens: int,
        on_progress: ProgressCallback | None = None,
    ) -> str:
        headers = {
            "Authorization": f"Bearer {settings.openrouter_api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "model": self.model_id,
            "max_tokens": max_tokens,
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
            "stream": True,
        }

        collected = []
        async with httpx.AsyncClient(timeout=300) as client:
            async with client.stream("POST", self.BASE_URL, headers=headers, json=payload) as resp:
                resp.raise_for_status()
                async for line in resp.aiter_lines():
                    if not line.startswith("data: "):
                        continue
                    data = line[6:]
                    if data == "[DONE]":
                        break
                    chunk = json.loads(data)
                    delta = chunk.get("choices", [{}])[0].get("delta", {})
                    content = delta.get("content", "")
                    if content:
                        collected.append(content)

        return "".join(collected)


def get_ai_client(model_id: str) -> AIClient:
    if "/" in model_id:
        return OpenRouterClient(model_id)
    return AnthropicClient(model_id)