from __future__ import annotations

import asyncio
import logging
import random
from abc import ABC, abstractmethod
from typing import Any, AsyncIterator, Callable, Coroutine, Literal, TypeVar

from openai import AsyncOpenAI, APIStatusError, APIConnectionError, APITimeoutError
from google import genai
from google.genai import types

from clawagents.config.config import EngineConfig

logger = logging.getLogger(__name__)

T = TypeVar("T")

# ─── Public Types ──────────────────────────────────────────────────────────


class LLMMessage:
    def __init__(self, role: Literal["system", "user", "assistant"], content: str):
        self.role = role
        self.content = content


class LLMResponse:
    def __init__(
        self,
        content: str,
        model: str,
        tokens_used: int,
        partial: bool = False,
    ):
        self.content = content
        self.model = model
        self.tokens_used = tokens_used
        self.partial = partial


OnChunkCallback = (
    Callable[[str], Coroutine[Any, Any, None]] | Callable[[str], None] | None
)


class LLMProvider(ABC):
    name: str

    @abstractmethod
    async def chat(
        self,
        messages: list[LLMMessage],
        on_chunk: OnChunkCallback = None,
        cancel_event: asyncio.Event | None = None,
    ) -> LLMResponse:
        pass


# ─── Streaming Robustness Internals ───────────────────────────────────────

_MAX_RETRIES = 3
_INITIAL_DELAY_S = 1.0
_MAX_DELAY_S = 16.0
_CHUNK_STALL_S = 60.0
_RETRYABLE_STATUS_CODES = frozenset({429, 500, 502, 503, 504})


def _is_retryable(err: BaseException) -> bool:
    if isinstance(err, APIStatusError):
        return err.status_code in _RETRYABLE_STATUS_CODES
    if isinstance(err, (APIConnectionError, APITimeoutError)):
        return True
    if isinstance(err, Exception):
        msg = str(err).lower()
        return any(
            tok in msg
            for tok in (
                "econnreset", "network", "timeout", "stream stalled",
                "rate limit", "too many requests", "service unavailable",
                "429", "500", "502", "503", "504",
            )
        )
    return False


def _jittered_delay(attempt: int) -> float:
    base = _INITIAL_DELAY_S * (2 ** attempt)
    return min(base + random.random() * base * 0.1, _MAX_DELAY_S)


async def _stall_guarded_stream(
    aiter: AsyncIterator[T],
    timeout_s: float,
) -> AsyncIterator[T]:
    """Yield items from *aiter*, raising TimeoutError if no item arrives
    within *timeout_s* seconds (stall detection)."""
    ait = aiter.__aiter__()
    while True:
        try:
            chunk = await asyncio.wait_for(ait.__anext__(), timeout=timeout_s)
            yield chunk
        except StopAsyncIteration:
            return


async def _invoke_callback(
    cb: OnChunkCallback,
    text: str,
) -> None:
    """Call *cb* with *text*, isolating errors so a broken callback
    never kills the stream."""
    if cb is None:
        return
    try:
        if asyncio.iscoroutinefunction(cb):
            await cb(text)
        else:
            cb(text)
    except Exception:
        logger.debug("onChunk callback raised — isolated", exc_info=True)


async def _with_retry(
    tag: str,
    fn: Callable[[], Coroutine[Any, Any, T]],
) -> T:
    last_error: BaseException | None = None
    for attempt in range(_MAX_RETRIES + 1):
        if attempt > 0:
            delay = _jittered_delay(attempt - 1)
            logger.warning(
                "  [%s] Retry %d/%d after %.1fs", tag, attempt, _MAX_RETRIES, delay,
            )
            await asyncio.sleep(delay)
        try:
            return await fn()
        except Exception as exc:
            last_error = exc
            if not _is_retryable(exc):
                break
    raise last_error  # type: ignore[misc]


# ─── OpenAI Provider ──────────────────────────────────────────────────────


class OpenAIProvider(LLMProvider):
    name = "openai"

    def __init__(self, config: EngineConfig):
        self.client = AsyncOpenAI(api_key=config.openai_api_key)
        self.model = config.openai_model

    async def chat(
        self,
        messages: list[LLMMessage],
        on_chunk: OnChunkCallback = None,
        cancel_event: asyncio.Event | None = None,
    ) -> LLMResponse:
        formatted = [{"role": m.role, "content": m.content} for m in messages]

        if not on_chunk:
            return await _with_retry("openai", lambda: self._request_once(formatted))
        return await self._stream_with_retry(formatted, on_chunk, cancel_event)

    async def _request_once(
        self, messages: list[dict[str, str]],
    ) -> LLMResponse:
        resp = await self.client.chat.completions.create(
            model=self.model,
            messages=messages,  # type: ignore[arg-type]
        )
        return LLMResponse(
            content=resp.choices[0].message.content or "",
            model=self.model,
            tokens_used=resp.usage.total_tokens if resp.usage else 0,
        )

    async def _stream_with_retry(
        self,
        messages: list[dict[str, str]],
        on_chunk: OnChunkCallback,
        cancel_event: asyncio.Event | None,
    ) -> LLMResponse:
        last_error: BaseException | None = None

        for attempt in range(_MAX_RETRIES + 1):
            if attempt > 0:
                delay = _jittered_delay(attempt - 1)
                logger.warning(
                    "  [openai] Stream retry %d/%d after %.1fs",
                    attempt, _MAX_RETRIES, delay,
                )
                await asyncio.sleep(delay)

            full_content = ""
            final_tokens = 0

            try:
                stream = await self.client.chat.completions.create(
                    model=self.model,
                    messages=messages,  # type: ignore[arg-type]
                    stream=True,
                    stream_options={"include_usage": True},
                )

                async for chunk in _stall_guarded_stream(stream, _CHUNK_STALL_S):
                    if cancel_event and cancel_event.is_set():
                        await stream.close()
                        return LLMResponse(
                            content=full_content,
                            model=self.model,
                            tokens_used=final_tokens,
                            partial=True,
                        )

                    try:
                        if (
                            chunk.choices
                            and chunk.choices[0].delta
                            and chunk.choices[0].delta.content
                        ):
                            text = chunk.choices[0].delta.content
                            full_content += text
                            await _invoke_callback(on_chunk, text)
                        if chunk.usage:
                            final_tokens = chunk.usage.total_tokens
                    except Exception:
                        pass  # malformed chunk — skip

                return LLMResponse(
                    content=full_content,
                    model=self.model,
                    tokens_used=final_tokens,
                )

            except Exception as exc:
                last_error = exc
                if full_content:
                    logger.warning(
                        "  [openai] Stream interrupted after %d chars — returning partial",
                        len(full_content),
                    )
                    return LLMResponse(
                        content=full_content,
                        model=self.model,
                        tokens_used=final_tokens,
                        partial=True,
                    )
                if not _is_retryable(exc):
                    break

        raise last_error  # type: ignore[misc]


# ─── Gemini Provider ──────────────────────────────────────────────────────


class GeminiProvider(LLMProvider):
    name = "gemini"

    def __init__(self, config: EngineConfig):
        self.client = genai.Client(api_key=config.gemini_api_key)
        self.model = config.gemini_model

    async def chat(
        self,
        messages: list[LLMMessage],
        on_chunk: OnChunkCallback = None,
        cancel_event: asyncio.Event | None = None,
    ) -> LLMResponse:
        system_instruction = "\n".join(
            m.content for m in messages if m.role == "system"
        )
        user_messages = "\n\n".join(
            m.content for m in messages if m.role != "system"
        )

        config_opts: dict[str, Any] = {}
        if system_instruction:
            config_opts["system_instruction"] = system_instruction
        gemini_config = types.GenerateContentConfig(**config_opts)

        if not on_chunk:
            return await _with_retry("gemini", lambda: self._request_once(
                user_messages, gemini_config,
            ))
        return await self._stream_with_retry(
            user_messages, gemini_config, on_chunk, cancel_event,
        )

    async def _request_once(
        self,
        user_messages: str,
        gemini_config: types.GenerateContentConfig,
    ) -> LLMResponse:
        resp = await self.client.aio.models.generate_content(
            model=self.model,
            contents=user_messages,
            config=gemini_config,
        )
        return LLMResponse(
            content=resp.text or "",
            model=self.model,
            tokens_used=(
                resp.usage_metadata.candidates_token_count
                if resp.usage_metadata
                else 0
            ),
        )

    async def _stream_with_retry(
        self,
        user_messages: str,
        gemini_config: types.GenerateContentConfig,
        on_chunk: OnChunkCallback,
        cancel_event: asyncio.Event | None,
    ) -> LLMResponse:
        last_error: BaseException | None = None

        for attempt in range(_MAX_RETRIES + 1):
            if attempt > 0:
                delay = _jittered_delay(attempt - 1)
                logger.warning(
                    "  [gemini] Stream retry %d/%d after %.1fs",
                    attempt, _MAX_RETRIES, delay,
                )
                await asyncio.sleep(delay)

            full_content = ""
            final_tokens = 0

            try:
                stream = await self.client.aio.models.generate_content_stream(
                    model=self.model,
                    contents=user_messages,
                    config=gemini_config,
                )

                async for chunk in _stall_guarded_stream(stream, _CHUNK_STALL_S):
                    if cancel_event and cancel_event.is_set():
                        return LLMResponse(
                            content=full_content,
                            model=self.model,
                            tokens_used=final_tokens,
                            partial=True,
                        )

                    try:
                        if chunk.text:
                            full_content += chunk.text
                            await _invoke_callback(on_chunk, chunk.text)
                        if chunk.usage_metadata:
                            final_tokens = chunk.usage_metadata.candidates_token_count
                    except Exception:
                        pass  # malformed chunk — skip

                return LLMResponse(
                    content=full_content,
                    model=self.model,
                    tokens_used=final_tokens,
                )

            except Exception as exc:
                last_error = exc
                if full_content:
                    logger.warning(
                        "  [gemini] Stream interrupted after %d chars — returning partial",
                        len(full_content),
                    )
                    return LLMResponse(
                        content=full_content,
                        model=self.model,
                        tokens_used=final_tokens,
                        partial=True,
                    )
                if not _is_retryable(exc):
                    break

        raise last_error  # type: ignore[misc]


# ─── Factory ──────────────────────────────────────────────────────────────


def create_provider(config: EngineConfig) -> LLMProvider:
    if config.get_provider() == "gemini":
        return GeminiProvider(config)
    return OpenAIProvider(config)
