from __future__ import annotations

import asyncio
import logging
import os
import random
from typing import Any

from autogen_ext.models.openai import AzureOpenAIChatCompletionClient
from autogen_ext.models.openai import BaseOpenAIChatCompletionClient
from prompt_auto_lab.core.papyrus_chat_completion_client import PapyrusChatCompletionClient
from prompt_auto_lab.core.copilot_proxy_client import CopilotProxyClient
from autogen_core.models import ModelFamily
from openai import AsyncOpenAI

logger = logging.getLogger(__name__)
 

def _is_rate_limit_error(exc: Exception) -> bool:
    status_code = getattr(exc, "status_code", None)
    if status_code == 429:
        return True
    response = getattr(exc, "response", None)
    if response is not None and getattr(response, "status_code", None) == 429:
        return True
    if exc.__class__.__name__ == "RateLimitError":
        return True
    text = str(exc)
    return "RateLimit" in text or "TooManyRequests" in text or "RateLimitReached" in text


def _extract_retry_after_seconds(exc: Exception) -> float | None:
    response = getattr(exc, "response", None)
    if response is not None:
        headers = getattr(response, "headers", None) or {}
        retry_after = headers.get("retry-after") or headers.get("Retry-After")
        if retry_after:
            try:
                return float(retry_after)
            except ValueError:
                pass
        retry_after_ms = headers.get("retry-after-ms") or headers.get("x-ms-retry-after-ms")
        if retry_after_ms:
            try:
                return float(retry_after_ms) / 1000.0
            except ValueError:
                pass

    match = None
    try:
        import re

        match = re.search(r"retry after ([\d.]+) seconds", str(exc), re.IGNORECASE)
    except Exception:
        match = None
    if match:
        try:
            return float(match.group(1))
        except ValueError:
            return None
    return None


def _compute_retry_delay(
    attempt: int,
    retry_after_s: float | None,
    base_delay_s: float,
    max_delay_s: float,
) -> float:
    if retry_after_s is not None:
        return max(retry_after_s, base_delay_s)
    delay = base_delay_s * (2 ** attempt)
    if max_delay_s > 0:
        delay = min(delay, max_delay_s)
    return delay


class RetryingAzureOpenAIChatCompletionClient(AzureOpenAIChatCompletionClient):
    def __init__(
        self,
        *args: Any,
        rate_limit_retries: int = 0,
        rate_limit_base_delay_s: float = 1.0,
        rate_limit_max_delay_s: float = 30.0,
        rate_limit_jitter_s: float = 0.0,
        **kwargs: Any,
    ) -> None:
        super().__init__(*args, **kwargs)
        self._rate_limit_retries = max(0, int(rate_limit_retries))
        self._rate_limit_base_delay_s = float(rate_limit_base_delay_s)
        self._rate_limit_max_delay_s = float(rate_limit_max_delay_s)
        self._rate_limit_jitter_s = float(rate_limit_jitter_s)

    async def create(self, *args: Any, **kwargs: Any) -> Any:
        if self._rate_limit_retries <= 0:
            return await super().create(*args, **kwargs)

        attempt = 0
        while True:
            try:
                return await super().create(*args, **kwargs)
            except Exception as exc:
                if not _is_rate_limit_error(exc) or attempt >= self._rate_limit_retries:
                    raise
                retry_after_s = _extract_retry_after_seconds(exc)
                delay_s = _compute_retry_delay(
                    attempt,
                    retry_after_s,
                    self._rate_limit_base_delay_s,
                    self._rate_limit_max_delay_s,
                )
                if self._rate_limit_jitter_s > 0:
                    delay_s += random.random() * self._rate_limit_jitter_s
                logger.info(
                    "Rate limit hit, retrying in {:.2f}s (attempt {}/{})".format(
                        delay_s, attempt + 1, self._rate_limit_retries
                    )
                )
                await asyncio.sleep(delay_s)
                attempt += 1


def build_azure_model_client(
    *,
    token_provider: Any,
    timeout_s: float | None = None,
    max_retries: int | None = None,
    deployment: str,
    endpoint: str,
    api_version: str,
) -> AzureOpenAIChatCompletionClient:
    if timeout_s is None:
        timeout_s = float(os.environ.get("AUTOPROMPTLAB_TIMEOUT_S", "60"))
    if max_retries is None:
        max_retries = int(os.environ.get("AUTOPROMPTLAB_MAX_RETRIES", "2"))
    rate_limit_retries = int(os.environ.get("AUTOPROMPTLAB_RATE_LIMIT_RETRIES", "3"))
    rate_limit_base_delay_s = float(os.environ.get("AUTOPROMPTLAB_RATE_LIMIT_BASE_DELAY_S", "1.0"))
    rate_limit_max_delay_s = float(os.environ.get("AUTOPROMPTLAB_RATE_LIMIT_MAX_DELAY_S", "30.0"))
    rate_limit_jitter_s = float(os.environ.get("AUTOPROMPTLAB_RATE_LIMIT_JITTER_S", "0.5"))
    # _ = max_retries  # Papyrus client currently does not implement retries.
    # endpoint = os.environ.get("AUTOPROMPTLAB_PAPYRUS_ENDPOINT", "https://hk.papyrus.binginternal.com/chat/completions")
    # model = os.environ.get("AUTOPROMPTLAB_PAPYRUS_MODEL", "SyndicationAPI-TextToText-ModelF-2")
    # quota_id = os.environ.get("AUTOPROMPTLAB_PAPYRUS_QUOTA_ID", "")
    # timeout_ms = int(os.environ.get("AUTOPROMPTLAB_PAPYRUS_TIMEOUT_MS", "100000"))

    # return PapyrusChatCompletionClient(
    #     endpoint=endpoint,
    #     model=model,
    #     token_provider=token_provider,
    #     papyrus_quota_id=quota_id,
    #     papyrus_timeout_ms=timeout_ms,
    #     request_timeout_s=timeout_s,
    # )


    return RetryingAzureOpenAIChatCompletionClient(
        azure_deployment=deployment,
        model="gpt-5.1-chat-2025-11-13",
        api_version=api_version,
        azure_endpoint=endpoint,
        azure_ad_token_provider=token_provider,
        timeout=timeout_s,
        max_retries=max_retries,
        rate_limit_retries=rate_limit_retries,
        rate_limit_base_delay_s=rate_limit_base_delay_s,
        rate_limit_max_delay_s=rate_limit_max_delay_s,
        rate_limit_jitter_s=rate_limit_jitter_s,
        # temperature= 0.1,
        model_info={
            "vision": True,
            "function_calling": True,
            "json_output": True,
            "family": ModelFamily.GPT_5,
            "structured_output": True,
            "multiple_system_messages": True,
        }
    )


def build_azure_model_thought_client(
    *,
    token_provider: Any,
    timeout_s: float | None = None,
    max_retries: int | None = None,
    endpoint: str,
    api_version: str,
) -> AzureOpenAIChatCompletionClient:
    if timeout_s is None:
        timeout_s = float(os.environ.get("AUTOPROMPTLAB_TIMEOUT_S", "60"))
    if max_retries is None:
        max_retries = int(os.environ.get("AUTOPROMPTLAB_MAX_RETRIES", "2"))
    rate_limit_retries = int(os.environ.get("AUTOPROMPTLAB_RATE_LIMIT_RETRIES", "3"))
    rate_limit_base_delay_s = float(os.environ.get("AUTOPROMPTLAB_RATE_LIMIT_BASE_DELAY_S", "1.0"))
    rate_limit_max_delay_s = float(os.environ.get("AUTOPROMPTLAB_RATE_LIMIT_MAX_DELAY_S", "30.0"))
    rate_limit_jitter_s = float(os.environ.get("AUTOPROMPTLAB_RATE_LIMIT_JITTER_S", "0.5"))
    # _ = max_retries  # Papyrus client currently does not implement retries.
    # endpoint = os.environ.get("AUTOPROMPTLAB_PAPYRUS_ENDPOINT", "https://hk.papyrus.binginternal.com/chat/completions")
    # model = os.environ.get("AUTOPROMPTLAB_PAPYRUS_MODEL", "SyndicationAPI-TextToText-ModelF-2")
    # quota_id = os.environ.get("AUTOPROMPTLAB_PAPYRUS_QUOTA_ID", "")
    # timeout_ms = int(os.environ.get("AUTOPROMPTLAB_PAPYRUS_TIMEOUT_MS", "100000"))

    # return PapyrusChatCompletionClient(
    #     endpoint=endpoint,
    #     model=model,
    #     token_provider=token_provider,
    #     papyrus_quota_id=quota_id,
    #     papyrus_timeout_ms=timeout_ms,
    #     request_timeout_s=timeout_s,
    # )


    return RetryingAzureOpenAIChatCompletionClient(
        azure_deployment="gpt-5.1-chat",
        model="gpt-5.1-chat-2025-11-13",
        api_version=api_version,
        azure_endpoint=endpoint,
        azure_ad_token_provider=token_provider,
        timeout=timeout_s,
        max_retries=max_retries,
        rate_limit_retries=rate_limit_retries,
        rate_limit_base_delay_s=rate_limit_base_delay_s,
        rate_limit_max_delay_s=rate_limit_max_delay_s,
        rate_limit_jitter_s=rate_limit_jitter_s,
        reasoning={"effort": "high"},
        reasoning_effort="high",
        # temperature= 0.1,
        model_info={
            "vision": True,
            "function_calling": True,
            "json_output": True,
            "family": ModelFamily.GPT_5,
            "structured_output": True,
            "multiple_system_messages": True,
        },
    )


def build_papyrus_model_client(
    *,
    token_provider: Any,
    endpoint: str,
    model: str,
    quota_id: str,
    papyrus_timeout_ms: int = 100000,
    request_timeout_s: int = 60,
    max_retries: int | None = None,
) -> BaseOpenAIChatCompletionClient:
    """Build a Papyrus model client based on ``BaseOpenAIChatCompletionClient``.

    Usage is similar to:

        kimi_client = openai.Client(base_url=KIMI_BASE_URL, api_key=KIMI_API_KEY, timeout=REQUEST_TIMEOUT)
        response = kimi_client.chat.completions.create(
            model=KIMI_MODEL,
            messages=kimi_messages,
            temperature=0.3,
            stream=stream,
        )

    Here we use ``AsyncOpenAI`` together with ``BaseOpenAIChatCompletionClient`` to wrap Papyrus.
    """

    if max_retries is None:
        max_retries = int(os.environ.get("AUTOPROMPTLAB_MAX_RETRIES", "2"))

    # Get a one-time token from ``token_provider`` and use it as the OpenAI SDK ``api_key``
    provider = token_provider if callable(token_provider) else None
    api_key = provider() if provider is not None else token_provider

    # Support endpoints that already include the "/chat/completions" path
    # to avoid duplicated paths such as ".../chat/completions/chat/completions".
    base_url = endpoint.rstrip("/")
    suffix = "/chat/completions"
    if base_url.endswith(suffix):
        base_url = base_url[: -len(suffix)]

    default_headers = {
        "papyrus-model-name": model,
        "papyrus-quota-id": quota_id,
        "papyrus-timeout-ms": str(papyrus_timeout_ms),
    }

    async_client = AsyncOpenAI(
        api_key=api_key,
        base_url=base_url,
        timeout=request_timeout_s,
        max_retries=max_retries,
        default_headers=default_headers,
    )

    model_info = {
        "vision": False,
        "function_calling": True,
        "json_output": True,
        "family": ModelFamily.GPT_5,
        "structured_output": True,
        "multiple_system_messages": True,
    }

    return PapyrusChatCompletionClient(
        client=async_client,
        create_args={"model": model},
        model_info=model_info,
        token_provider=provider,
    )


def build_copilot_proxy_client(
    *,
    model: str = "",
    timeout_s: float = 300.0,
) -> CopilotProxyClient:
    """Build a Copilot proxy client that forwards requests through VS Code extension.

    Args:
        model: The Copilot model ID to use (e.g., "gpt-4o", "claude-3.5-sonnet").
               If empty, uses the first available model from the extension.
        timeout_s: Request timeout in seconds (default 5 minutes, max 30 minutes).

    Returns:
        CopilotProxyClient instance.

    Raises:
        RuntimeError: If VS Code extension is not connected.
    """
    # Lazy import to avoid circular dependency:
    # model_clients -> copilot_queue_service -> api.models -> api -> routes -> protocol -> model_clients
    from prompt_auto_lab.api.services.copilot_queue_service import get_copilot_queue

    queue = get_copilot_queue()

    if not queue.is_extension_connected():
        raise RuntimeError(
            "Copilot provider requires VS Code extension connection. "
            "Please ensure VS Code is open with the AutoPromptLab extension active."
        )

    # If no model specified, use first available
    if not model:
        available = queue.get_available_models()
        if available:
            model = available[0].id
        else:
            raise RuntimeError("No Copilot models available from VS Code extension.")

    return CopilotProxyClient(
        queue_service=queue,
        model=model,
        timeout_s=timeout_s,
    )


async def fetch_local_models(
    base_url: str = "http://localhost:23333",
) -> list[dict]:
    """Fetch available models from a local LLM service.

    Args:
        base_url: The base URL of the local service (default: http://localhost:23333).

    Returns:
        List of model info dictionaries with keys: id, name, vendor, family, version,
        maxInputTokens, capabilities.

    Raises:
        Exception: If the request fails or returns invalid data.
    """
    import httpx

    url = f"{base_url.rstrip('/')}/api/v1/lm/chatModels"
    async with httpx.AsyncClient(timeout=30.0) as client:
        response = await client.get(url)
        response.raise_for_status()
        return response.json()


def _map_local_model_family(family: str) -> ModelFamily:
    """Map local model family string to AutoGen ModelFamily enum."""
    family_lower = family.lower()
    if "gpt-5" in family_lower or "gpt5" in family_lower:
        return ModelFamily.GPT_5
    if "gpt-4" in family_lower or "gpt4" in family_lower:
        return ModelFamily.GPT_4O
    if "claude" in family_lower:
        return ModelFamily.CLAUDE_3_5_SONNET
    if "gemini" in family_lower:
        return ModelFamily.GEMINI_1_5_FLASH
    # Default to GPT-4o for unknown families
    return ModelFamily.GPT_4O


def build_local_model_client(
    *,
    model: str,
    base_url: str = "http://localhost:23333/api/openai/v1",
    timeout_s: float = 300.0,
    max_retries: int = 2,
    model_info_override: dict | None = None,
) -> BaseOpenAIChatCompletionClient:
    """Build a local model client that connects to a local LLM service.

    The local service should expose an OpenAI-compatible chat completions endpoint.

    Args:
        model: The model ID to use (e.g., "gpt-5.1-codex", "claude-sonnet-4").
        base_url: The base URL for the OpenAI-compatible API
                  (default: http://localhost:23333/api/openai/v1).
        timeout_s: Request timeout in seconds (default 5 minutes).
        max_retries: Number of retries on failure (default 2).
        model_info_override: Optional dict to override model capabilities.
                             Keys: vision, function_calling, json_output, family, etc.

    Returns:
        BaseOpenAIChatCompletionClient instance.
    """
    # Accept either a base `/v1` URL or a full `/chat/completions` URL.
    base_url = base_url.rstrip("/")
    suffix = "/chat/completions"
    if base_url.endswith(suffix):
        base_url = base_url[: -len(suffix)]

    async_client = AsyncOpenAI(
        api_key="local",  # Dummy key for local service
        base_url=base_url,
        timeout=timeout_s,
        max_retries=max_retries,
    )

    # Default model info - can be overridden
    model_info = {
        "vision": True,
        "function_calling": True,
        "json_output": True,
        "family": _map_local_model_family(model),
        "structured_output": True,
        "multiple_system_messages": True,
    }

    if model_info_override:
        model_info.update(model_info_override)

    return BaseOpenAIChatCompletionClient(
        client=async_client,
        create_args={"model": model},
        model_info=model_info,
    )
