from collections.abc import AsyncIterator, Awaitable
from typing import Any, Protocol, cast, runtime_checkable

import litellm

from otto.acp.registry import is_acp_model
from otto.auth import AuthStorage
from otto.auth.base import OAuthProvider
from otto.auth.registry import (
    get_litellm_overrides,
    get_oauth_provider,
    is_oauth_model,
    needs_custom_provider,
)
from otto.errors import OttoConfigError, raise_as_otto
from otto.log import get_logger

log = get_logger(__name__)


def format_acp_chat_rejection(model: str) -> str:
    return (
        f"ACP model '{model}' is delegation-only and cannot be used in interactive chat. "
        "Use delegation instead, or switch to a standard chat model with /model."
    )


@runtime_checkable
class StreamProvider(Protocol):
    async def stream_completion(
        self,
        model: str,
        messages: list[dict],
        *,
        tools: list[dict] | None = None,
        tool_choice: str | dict | None = None,
        max_tokens: int | None = None,
        reasoning_effort: str | None = None,
    ) -> AsyncIterator[Any]: ...


class LiteLLMStreamProvider:
    """Wraps litellm.acompletion with stream=True and optional OAuth overrides."""

    def __init__(self, overrides: dict[str, Any] | None = None) -> None:
        self._overrides = overrides

    async def stream_completion(
        self,
        model: str,
        messages: list[dict],
        *,
        tools: list[dict] | None = None,
        tool_choice: str | dict | None = None,
        max_tokens: int | None = None,
        reasoning_effort: str | None = None,
    ) -> AsyncIterator[Any]:
        kwargs: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": True,
            "num_retries": 3,
            "request_timeout": 120,
        }

        if tools:
            kwargs["tools"] = tools
        if tool_choice is not None:
            kwargs["tool_choice"] = tool_choice
        if max_tokens is not None:
            kwargs["max_tokens"] = max_tokens
        if reasoning_effort is not None:
            kwargs["reasoning_effort"] = reasoning_effort
        if self._overrides is not None:
            kwargs.update(
                {key: value for key, value in self._overrides.items() if key != "extra_headers"}
            )
            override_headers = self._overrides.get("extra_headers")
            if isinstance(override_headers, dict):
                kwargs["extra_headers"] = {
                    **kwargs.get("extra_headers", {}),
                    **override_headers,
                }
        return cast(AsyncIterator[Any], await litellm.acompletion(**kwargs))


class OAuthStreamProvider:
    """Wraps an OAuthProvider instance, handling prefix injection and Awaitable streams."""

    def __init__(
        self,
        oauth_provider: OAuthProvider,
        *,
        system_prompt_prefix: str | None = None,
    ) -> None:
        self._provider = oauth_provider
        self._system_prompt_prefix = system_prompt_prefix

    def effective_system(self, system: str) -> str:
        if self._system_prompt_prefix:
            return f"{self._system_prompt_prefix.strip()}\n\n{system}"
        return system

    async def stream_completion(
        self,
        model: str,
        messages: list[dict],
        *,
        tools: list[dict] | None = None,
        tool_choice: str | dict | None = None,
        max_tokens: int | None = None,
        reasoning_effort: str | None = None,
    ) -> AsyncIterator[Any]:
        system_prompt = None
        for msg in messages:
            if msg.get("role") == "system":
                system_prompt = msg.get("content")
                break

        provider_kwargs: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "tools": tools,
            "tool_choice": tool_choice,
            "system_prompt": system_prompt,
            "max_tokens": max_tokens,
        }
        if reasoning_effort is not None:
            provider_kwargs["reasoning_effort"] = reasoning_effort

        stream = self._provider.stream_completion(**provider_kwargs)
        if isinstance(stream, Awaitable):
            stream = await stream
        return stream


async def resolve_provider(
    model: str,
    system: str,
    auth_storage: AuthStorage | None = None,
) -> StreamProvider:
    """Detect the right provider for a model and return a StreamProvider adapter.

    Absorbs the provider-detection logic from agent.py run() lines 105-130.
    """
    if is_acp_model(model):
        raise OttoConfigError(format_acp_chat_rejection(model))

    if auth_storage is None:
        return LiteLLMStreamProvider()

    is_claude_code = model.lower().startswith("claude-code/")

    # Try LiteLLM overrides first (for OAuth models that need provider-specific kwargs)
    try:
        oauth_overrides = await _get_oauth_overrides(model, auth_storage)
    except Exception as exc:
        log.error("oauth override setup failed", extra={"model": model, "error": str(exc)})
        raise_as_otto(exc)

    if oauth_overrides is not None:
        log.info(
            "oauth litellm overrides selected",
            extra={"model": model, "keys": list(oauth_overrides)},
        )
        return LiteLLMStreamProvider(overrides=oauth_overrides)

    # Try custom OAuth provider (claude-code, codex, gemini-cli)
    if needs_custom_provider(model) and is_oauth_model(model, auth_storage):
        try:
            oauth_provider = await get_oauth_provider(model, auth_storage=auth_storage)
        except Exception as exc:
            log.error("oauth provider setup failed", extra={"model": model, "error": str(exc)})
            raise_as_otto(exc)

        # Determine system_prompt_prefix — exclude for claude-code + anthropic_oauth
        raw_prefix = getattr(oauth_provider, "system_prompt_prefix", None)
        provider_name = str(getattr(oauth_provider, "provider_name", ""))
        prefix: str | None = None
        if raw_prefix and not (is_claude_code and provider_name == "anthropic_oauth"):
            prefix = raw_prefix

        log.info(
            "oauth provider selected",
            extra={"model": model, "provider": provider_name},
        )
        return OAuthStreamProvider(oauth_provider, system_prompt_prefix=prefix)

    return LiteLLMStreamProvider()


async def _get_oauth_overrides(
    model: str,
    auth_storage: AuthStorage,
) -> dict[str, Any] | None:
    if needs_custom_provider(model):
        return None
    return await get_litellm_overrides(model, auth_storage)
