"""Copilot proxy model client that forwards requests through VS Code extension."""
from __future__ import annotations

import asyncio
import os
import uuid
import logging
import json
import random
import re
from typing import Any, AsyncGenerator, Sequence, Mapping, Literal, Optional, Union, TYPE_CHECKING

from pydantic import BaseModel
from autogen_core import CancellationToken
from autogen_core.models import (
    ChatCompletionClient,
    CreateResult,
    LLMMessage,
    RequestUsage,
    ModelInfo,
    ModelFamily,
    SystemMessage,
    UserMessage,
    AssistantMessage,
    FunctionExecutionResultMessage,
)
from autogen_core.tools import Tool, ToolSchema

if TYPE_CHECKING:
    from prompt_auto_lab.api.services.copilot_queue_service import CopilotQueueService

logger = logging.getLogger(__name__)

# Default timeout for Copilot requests (5 minutes)
DEFAULT_TIMEOUT_S = 300.0
# Max timeout (30 minutes for long reasoning)
MAX_TIMEOUT_S = 1800.0


_RETRYABLE_ERROR_RE = re.compile(r"\b(429|500|502|503|504)\b", re.IGNORECASE)


def _is_cancelled(token: CancellationToken | None) -> bool:
    if token is None:
        return False
    is_cancelled = getattr(token, "is_cancelled", None)
    if callable(is_cancelled):
        try:
            return bool(is_cancelled())
        except Exception:
            return False
    return bool(getattr(token, "cancelled", False) or getattr(token, "_cancelled", False))


def _is_retryable_error_text(text: str) -> bool:
    return True
    # if not text:
    #     return False
    # lowered = text.lower()
    # return bool(
    #     _RETRYABLE_ERROR_RE.search(text)
    #     or "too many requests" in lowered
    #     or "rate limit" in lowered
    #     or "service unavailable" in lowered
    #     or "bad gateway" in lowered
    #     or "gateway timeout" in lowered
    #     or "econnreset" in lowered
    #     or "etimedout" in lowered
    #     or "enotfound" in lowered
    #     or "err_http2" in lowered
    #     or "protocol_error" in lowered
    #     or "network" in lowered and "error" in lowered
    #     or "connection" in lowered and ("reset" in lowered or "refused" in lowered or "closed" in lowered)
    # )


def _compute_retry_delay_s(attempt: int, base_delay_s: float, max_delay_s: float) -> float:
    exp = min(max(attempt - 1, 0), 6)
    base = min(max_delay_s, base_delay_s * (2**exp))
    jitter = 0.2 + random.random() * 0.4  # 20%–60%
    return max(0.0, base * jitter)


# Local data classes to avoid circular imports with api.models
from enum import Enum
from datetime import datetime
from pydantic import Field


class CopilotMessageRole(str, Enum):
    """Role of a message in a Copilot conversation."""
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"


class CopilotToolCallPart(BaseModel):
    """Tool call embedded in a message."""
    call_id: str
    name: str
    arguments: str  # JSON string


class CopilotToolResultPart(BaseModel):
    """Tool execution result embedded in a message."""
    call_id: str
    content: str
    is_error: bool | None = None


class CopilotMessage(BaseModel):
    """A message in a Copilot conversation."""
    role: CopilotMessageRole
    content: str
    tool_calls: list[CopilotToolCallPart] | None = None
    tool_results: list[CopilotToolResultPart] | None = None


class CopilotToolDefinition(BaseModel):
    """Definition of a tool for Copilot."""
    name: str
    description: str
    parameters: dict[str, Any]


class CopilotToolCall(BaseModel):
    """A tool call from Copilot."""
    id: str
    name: str
    arguments: str  # JSON string


class CopilotRequest(BaseModel):
    """A request to be processed by the Copilot proxy."""
    request_id: str
    messages: list[CopilotMessage]
    model: str
    tools: list[CopilotToolDefinition] | None = None
    temperature: float | None = None
    max_tokens: int | None = None
    created_at: datetime = Field(default_factory=datetime.utcnow)


def _safe_json_dumps(value: Any) -> str:
    if isinstance(value, str):
        return value
    try:
        return json.dumps(value)
    except Exception:
        return str(value)


def _as_tool_call(obj: Any) -> CopilotToolCallPart | None:
    call_id = getattr(obj, "call_id", None) or getattr(obj, "id", None)
    name = getattr(obj, "name", None) or getattr(obj, "tool_name", None)
    args = getattr(obj, "arguments", None)
    if args is None:
        args = getattr(obj, "input", None)
    if call_id and name and args is not None:
        return CopilotToolCallPart(
            call_id=str(call_id),
            name=str(name),
            arguments=_safe_json_dumps(args),
        )
    return None


def _as_tool_result(obj: Any) -> CopilotToolResultPart | None:
    call_id = getattr(obj, "call_id", None) or getattr(obj, "id", None)
    content = getattr(obj, "content", None)
    if content is None:
        content = getattr(obj, "result", None)
    if content is None:
        content = getattr(obj, "output", None)
    if call_id and content is not None:
        is_error = getattr(obj, "is_error", None)
        if is_error is None:
            is_error = getattr(obj, "error", None)
        return CopilotToolResultPart(
            call_id=str(call_id),
            content=_safe_json_dumps(content),
            is_error=bool(is_error) if is_error is not None else None,
        )
    return None


def _extract_text_part(part: Any) -> str | None:
    if isinstance(part, str):
        return part
    if hasattr(part, "text") and isinstance(getattr(part, "text", None), str):
        return str(getattr(part, "text"))
    if hasattr(part, "content") and isinstance(getattr(part, "content", None), str):
        return str(getattr(part, "content"))
    return None


def _convert_message(msg: LLMMessage) -> CopilotMessage:
    """Convert AutoGen LLMMessage to CopilotMessage."""
    # Handle different message types (autogen_core.models)
    if isinstance(msg, SystemMessage):
        role = CopilotMessageRole.SYSTEM
    elif isinstance(msg, AssistantMessage):
        role = CopilotMessageRole.ASSISTANT
    elif isinstance(msg, FunctionExecutionResultMessage):
        # Tool results should be sent as user tool result parts to VS Code
        role = CopilotMessageRole.USER
    else:
        role = CopilotMessageRole.USER

    # Extract content and tool calls/results
    content = ""
    tool_calls: list[CopilotToolCallPart] = []
    tool_results: list[CopilotToolResultPart] = []

    if hasattr(msg, "tool_calls"):
        for tc in getattr(msg, "tool_calls") or []:
            parsed = _as_tool_call(tc)
            if parsed:
                tool_calls.append(parsed)

    if hasattr(msg, "tool_results"):
        for tr in getattr(msg, "tool_results") or []:
            parsed = _as_tool_result(tr)
            if parsed:
                tool_results.append(parsed)

    if hasattr(msg, 'content'):
        if isinstance(msg.content, str):
            content = msg.content
        elif isinstance(msg.content, list):
            # Handle multi-part content
            parts: list[str] = []
            for part in msg.content:
                tool_call = _as_tool_call(part)
                if tool_call:
                    tool_calls.append(tool_call)
                    continue
                tool_result = _as_tool_result(part)
                if tool_result:
                    tool_results.append(tool_result)
                    continue
                text = _extract_text_part(part)
                if text:
                    parts.append(text)
            content = "\n".join(parts)

    # Some AutoGen message types (notably FunctionExecutionResultMessage) may carry the tool call_id
    # on the message itself (and the result in `content`), rather than in `tool_results` or parts.
    # If we don't preserve call_id -> result, the VS Code LM tool pipeline can't correlate tool
    # results, and models may repeat the same tool call.
    if isinstance(msg, FunctionExecutionResultMessage) and not tool_results:
        msg_call_id = getattr(msg, "call_id", None) or getattr(msg, "id", None)
        if msg_call_id and content:
            is_error = getattr(msg, "is_error", None)
            if is_error is None:
                is_error = getattr(msg, "error", None)
            tool_results.append(
                CopilotToolResultPart(
                    call_id=str(msg_call_id),
                    content=_safe_json_dumps(content),
                    is_error=bool(is_error) if is_error is not None else None,
                )
            )
            # Prefer structured tool result over duplicating the same payload in freeform text.
            content = ""

    # Similarly, some assistant messages may have the tool call metadata on the message object.
    if isinstance(msg, AssistantMessage) and not tool_calls:
        parsed = _as_tool_call(msg)
        if parsed:
            tool_calls.append(parsed)
            if content == parsed.arguments:
                content = ""

    # If tool results are present, map to user role for VS Code tool-result handling.
    if tool_results and role == CopilotMessageRole.ASSISTANT:
        role = CopilotMessageRole.USER

    return CopilotMessage(
        role=role,
        content=content,
        tool_calls=tool_calls or None,
        tool_results=tool_results or None,
    )


def _convert_tool(tool: Tool | ToolSchema) -> CopilotToolDefinition:
    """Convert AutoGen Tool to CopilotToolDefinition."""
    if isinstance(tool, Tool):
        schema = tool.schema
    else:
        schema = tool

    return CopilotToolDefinition(
        name=schema.get("name", ""),
        description=schema.get("description", ""),
        parameters=schema.get("parameters", {}),
    )


class CopilotProxyClient(ChatCompletionClient):
    """Model client that proxies requests through VS Code extension to Copilot API.

    This client queues LLM requests and waits for the VS Code extension to
    process them using the vscode.lm API, then returns the responses.
    """

    def __init__(
        self,
        queue_service: CopilotQueueService,
        model: str,
        timeout_s: float = DEFAULT_TIMEOUT_S,
    ) -> None:
        self._queue = queue_service
        self._model = model
        self._timeout_s = min(timeout_s, MAX_TIMEOUT_S)
        self._total_usage = RequestUsage(prompt_tokens=0, completion_tokens=0)

    @property
    def model_info(self) -> ModelInfo:
        """Return model information."""
        return ModelInfo(
            vision=False,
            function_calling=True,
            json_output=True,
            family=ModelFamily.UNKNOWN,
            structured_output=False,
            multiple_system_messages=True,
        )

    @property
    def capabilities(self) -> ModelInfo:
        """Alias for model_info for compatibility."""
        return self.model_info

    def _create_request(
        self,
        messages: Sequence[LLMMessage],
        tools: Sequence[Tool | ToolSchema] | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
    ) -> CopilotRequest:
        """Create a CopilotRequest from AutoGen inputs."""
        request_id = str(uuid.uuid4())

        converted_messages = [_convert_message(m) for m in messages]
        logger.info(
            "CopilotProxyClient._create_request: request_id=%s msgs=%d roles=%s tool_calls=%s tool_results=%s",
            request_id,
            len(converted_messages),
            "->".join([m.role for m in converted_messages]),
            ",".join([str(len(m.tool_calls or [])) for m in converted_messages]),
            ",".join([str(len(m.tool_results or [])) for m in converted_messages]),
        )
        converted_tools = [_convert_tool(t) for t in tools] if tools else None

        return CopilotRequest(
            request_id=request_id,
            messages=converted_messages,
            model=self._model,
            tools=converted_tools,
            temperature=temperature,
            max_tokens=max_tokens,
        )

    async def create(
        self,
        messages: Sequence[LLMMessage],
        *,
        tools: Sequence[Tool | ToolSchema] = [],
        tool_choice: Tool | Literal["auto", "required", "none"] = "auto",
        json_output: Optional[bool | type[BaseModel]] = None,
        extra_create_args: Mapping[str, Any] = {},
        cancellation_token: Optional[CancellationToken] = None,
    ) -> CreateResult:
        """Send a request to Copilot via VS Code extension and wait for response."""
        logger.info(
            "CopilotProxyClient.create: model=%s tool_choice=%s tools=%d extra_create_args_keys=%s",
            self._model,
            getattr(tool_choice, "schema", None) and getattr(tool_choice, "schema").get("name") or tool_choice,
            len(tools) if tools else 0,
            sorted(list(extra_create_args.keys())) if extra_create_args else [],
        )
        if not self._queue.is_extension_connected():
            raise RuntimeError(
                "Copilot provider requires VS Code extension connection. "
                "Please open VS Code with the AutoPromptLab extension."
            )

        retry_attempts = int(
            extra_create_args.get("copilot_retry_attempts")
            or os.getenv("AUTOPROMPTLAB_COPILOT_RETRY_ATTEMPTS")
            or 5
        )
        retry_attempts = max(1, min(retry_attempts, 10))

        retry_base_delay_s = float(
            extra_create_args.get("copilot_retry_base_delay_s")
            or os.getenv("AUTOPROMPTLAB_COPILOT_RETRY_BASE_DELAY_S")
            or 0.5
        )
        retry_max_delay_s = float(
            extra_create_args.get("copilot_retry_max_delay_s")
            or os.getenv("AUTOPROMPTLAB_COPILOT_RETRY_MAX_DELAY_S")
            or 10.0
        )
        retry_base_delay_s = max(0.0, retry_base_delay_s)
        retry_max_delay_s = max(0.0, retry_max_delay_s)

        last_error: str | None = None
        response = None

        for attempt in range(1, retry_attempts + 1):
            if _is_cancelled(cancellation_token):
                raise asyncio.CancelledError()
            if not self._queue.is_extension_connected():
                raise RuntimeError(
                    "Copilot provider requires VS Code extension connection. "
                    "Please open VS Code with the AutoPromptLab extension."
                )

            request = self._create_request(
                messages=messages,
                tools=list(tools) if tools else None,
                temperature=extra_create_args.get("temperature"),
                max_tokens=extra_create_args.get("max_tokens"),
            )

            await self._queue.queue_request(request)
            logger.info(
                "Queued Copilot request: %s (attempt %d/%d)",
                request.request_id,
                attempt,
                retry_attempts,
            )

            try:
                response = await self._queue.wait_for_response(
                    request.request_id,
                    timeout=self._timeout_s,
                )
            except TimeoutError:
                last_error = f"Copilot request timed out after {self._timeout_s}s"
                logger.error("Copilot request timeout: %s", request.request_id)
                if attempt >= retry_attempts:
                    raise
                delay_s = _compute_retry_delay_s(attempt, retry_base_delay_s, retry_max_delay_s)
                logger.warning(
                    "Retrying Copilot request after timeout (attempt %d/%d); sleeping %.2fs",
                    attempt,
                    retry_attempts,
                    delay_s,
                )
                await asyncio.sleep(delay_s)
                continue

            if response.error:
                err_text = str(response.error)
                last_error = err_text
                if attempt < retry_attempts and _is_retryable_error_text(err_text):
                    delay_s = _compute_retry_delay_s(attempt, retry_base_delay_s, retry_max_delay_s)
                    logger.warning(
                        "Retryable Copilot error (attempt %d/%d): %s | sleeping %.2fs",
                        attempt,
                        retry_attempts,
                        err_text,
                        delay_s,
                    )
                    await asyncio.sleep(delay_s)
                    continue
                raise RuntimeError(f"Copilot API error: {response.error}")

            break

        if response is None:
            raise RuntimeError(f"Copilot API error: {last_error or 'Unknown error'}")

        # Convert response to CreateResult
        # Handle tool calls if present
        from autogen_core import FunctionCall

        function_calls = None
        if response.tool_calls:
            function_calls = [
                FunctionCall(
                    id=tc.id,
                    name=tc.name,
                    arguments=tc.arguments,
                )
                for tc in response.tool_calls
            ]

        content = response.content or ""

        # Create result
        result = CreateResult(
            content=content,
            finish_reason=response.finish_reason or "stop",
            usage=RequestUsage(prompt_tokens=0, completion_tokens=0),
            cached=False,
        )

        # Add function calls if present
        if function_calls:
            result = CreateResult(
                content=function_calls,
                finish_reason=response.finish_reason or "tool_calls",
                usage=RequestUsage(prompt_tokens=0, completion_tokens=0),
                cached=False,
            )

        return result

    async def create_stream(
        self,
        messages: Sequence[LLMMessage],
        *,
        tools: Sequence[Tool | ToolSchema] = [],
        tool_choice: Tool | Literal["auto", "required", "none"] = "auto",
        json_output: Optional[bool | type[BaseModel]] = None,
        extra_create_args: Mapping[str, Any] = {},
        cancellation_token: Optional[CancellationToken] = None,
    ) -> AsyncGenerator[Union[str, CreateResult], None]:
        """Streaming not supported - falls back to non-streaming."""
        result = await self.create(
            messages=messages,
            tools=tools,
            json_output=json_output,
            extra_create_args=extra_create_args,
            cancellation_token=cancellation_token,
        )
        if isinstance(result.content, str):
            yield result.content
        else:
            yield str(result.content)

    def actual_usage(self) -> RequestUsage:
        """Return actual token usage."""
        return self._total_usage

    def total_usage(self) -> RequestUsage:
        """Return total token usage."""
        return self._total_usage

    def count_tokens(
        self,
        messages: Sequence[LLMMessage],
        *,
        tools: Sequence[Tool | ToolSchema] = [],
    ) -> int:
        """Estimate token count (rough approximation)."""
        total = 0
        for msg in messages:
            if hasattr(msg, 'content') and isinstance(msg.content, str):
                total += len(msg.content) // 4  # Rough estimate
        return total

    def remaining_tokens(
        self,
        messages: Sequence[LLMMessage],
        *,
        tools: Sequence[Tool | ToolSchema] = [],
    ) -> int:
        """Return remaining tokens (estimate based on 128k context window)."""
        used = self.count_tokens(messages, tools=tools)
        # Assume 128k context window for most Copilot models
        return max(0, 128000 - used)

    async def close(self) -> None:
        """Close the client (no-op for proxy client)."""
        pass
