"""Kimi.com provider implementation for Oclawma.

Kimi (Moonshot AI) provides cloud-based LLM inference with the K2.5 model.
API documentation: https://platform.moonshot.cn/docs
"""

from __future__ import annotations

import json
import os
from collections.abc import AsyncIterator
from typing import Any

import httpx

from oclawma.providers.base import (
    AuthenticationError,
    BaseProvider,
    CompletionError,
    CompletionRequest,
    CompletionResponse,
    Message,
    ModelNotFoundError,
    ProviderConnectionError,
    ProviderError,
    RateLimitError,
    UsageStats,
)


class KimiProvider(BaseProvider):
    """Kimi.com (Moonshot AI) provider for cloud LLM inference.

    Supports the K2.5 model and provides streaming responses,
    proper token counting via API, and automatic fallback handling.
    """

    DEFAULT_BASE_URL = "https://api.moonshot.cn/v1"
    DEFAULT_MODEL = "kimi-k2-5"

    # Model name mappings
    MODEL_ALIASES = {
        "k2.5": "kimi-k2-5",
        "k2p5": "kimi-k2-5",
        "k1.5": "kimi-k1-5",
        "k1p5": "kimi-k1-5",
    }

    def __init__(
        self,
        base_url: str | None = None,
        api_key: str | None = None,
        timeout: float = 120.0,
    ) -> None:
        """Initialize the Kimi provider.

        Args:
            base_url: The base URL for Kimi API (default: https://api.moonshot.cn/v1).
            api_key: API key for authentication. If not provided, will look for
                    KIMI_API_KEY environment variable.
            timeout: Request timeout in seconds (default: 120).

        Raises:
            AuthenticationError: If no API key is provided or found in environment.
        """
        resolved_api_key = api_key or os.environ.get("KIMI_API_KEY")
        if not resolved_api_key:
            raise AuthenticationError(
                "Kimi API key is required. Set KIMI_API_KEY environment variable "
                "or pass api_key parameter."
            )

        super().__init__(base_url or self.DEFAULT_BASE_URL, resolved_api_key)
        self._timeout = timeout
        self._client: httpx.AsyncClient | None = None

    async def _get_client(self) -> httpx.AsyncClient:
        """Get or create the HTTP client."""
        if self._client is None or self._client.is_closed:
            self._client = httpx.AsyncClient(
                base_url=self.base_url,
                timeout=httpx.Timeout(self._timeout, connect=30.0),
                headers={
                    "Content-Type": "application/json",
                    "Authorization": f"Bearer {self.api_key}",
                },
            )
        return self._client

    async def close(self) -> None:
        """Close the HTTP client connection."""
        if self._client and not self._client.is_closed:
            await self._client.aclose()
            self._client = None

    def _resolve_model(self, model: str) -> str:
        """Resolve model alias to full model name.

        Args:
            model: Model name or alias.

        Returns:
            Full model name for API.
        """
        return self.MODEL_ALIASES.get(model, model)

    def _convert_messages(self, messages: list[Message]) -> list[dict[str, Any]]:
        """Convert messages to Kimi API format.

        Kimi uses OpenAI-compatible message format.

        Returns:
            List of messages in Kimi API format.
        """
        return [
            {
                "role": msg.role,
                "content": msg.content,
                **({"name": msg.name} if msg.name else {}),
            }
            for msg in messages
        ]

    def _handle_error(self, error: httpx.HTTPError) -> ProviderError:
        """Convert HTTP errors to provider errors."""
        if isinstance(error, httpx.ConnectError):
            return ProviderConnectionError(
                f"Cannot connect to Kimi API at {self.base_url}. " "Check your network connection.",
                cause=error,
            )
        elif isinstance(error, httpx.TimeoutException):
            return ProviderConnectionError(
                "Request to Kimi API timed out. The model may be loading or the request too complex.",
                cause=error,
            )
        elif isinstance(error, httpx.HTTPStatusError):
            status_code = error.response.status_code
            try:
                error_data = error.response.json()
                error_message = error_data.get("error", {}).get("message", str(error))
            except (json.JSONDecodeError, AttributeError):
                error_message = error.response.text or str(error)

            if status_code == 401:
                return AuthenticationError(
                    "Authentication failed. Check your KIMI_API_KEY.",
                    cause=error,
                )
            elif status_code == 404:
                return ModelNotFoundError(
                    f"Model not found: {error_message}",
                    cause=error,
                )
            elif status_code == 429:
                return RateLimitError(
                    f"Rate limit exceeded: {error_message}. Please retry later.",
                    cause=error,
                )
            elif status_code == 400:
                return CompletionError(
                    f"Bad request: {error_message}",
                    cause=error,
                )
            else:
                return CompletionError(
                    f"Kimi API returned HTTP {status_code}: {error_message}",
                    cause=error,
                )
        else:
            return CompletionError(f"Unexpected error: {error}", cause=error)

    def count_tokens(self, text: str, model: str | None = None) -> int:
        """Estimate the number of tokens in the given text.

        Uses a character-based heuristic optimized for the K2.5 model.
        Kimi models typically use ~3.5-4 characters per token for English text.

        Args:
            text: The text to count tokens for.
            model: Optional model name (ignored for estimation).

        Returns:
            Estimated token count.
        """
        if not text:
            return 0

        # Kimi uses a tokenizer similar to other modern LLMs
        # Roughly 4 characters per token for English, ~2-3 for CJK
        # We use a conservative estimate
        chars_per_token = 3.5
        return max(1, int((len(text) + chars_per_token - 1) // chars_per_token))

    async def complete(self, request: CompletionRequest) -> CompletionResponse:
        """Generate a chat completion.

        Args:
            request: The completion request.

        Returns:
            The completion response.

        Raises:
            ConnectionError: If unable to connect to Kimi API.
            AuthenticationError: If authentication fails.
            ModelNotFoundError: If the model is not available.
            CompletionError: If completion generation fails.
        """
        client = await self._get_client()
        model = self._resolve_model(request.model)
        messages = self._convert_messages(request.messages)

        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": False,
            "temperature": request.temperature,
            "top_p": request.top_p,
        }

        if request.max_tokens:
            payload["max_tokens"] = request.max_tokens

        try:
            response = await client.post("/chat/completions", json=payload)
            response.raise_for_status()
            data = response.json()

            choice = data.get("choices", [{}])[0]
            message = choice.get("message", {})
            content = message.get("content", "")
            finish_reason = choice.get("finish_reason")

            # Get token usage from API response
            usage_data = data.get("usage", {})
            usage = UsageStats(
                prompt_tokens=usage_data.get("prompt_tokens", 0),
                completion_tokens=usage_data.get("completion_tokens", 0),
                total_tokens=usage_data.get("total_tokens", 0),
            )

            return CompletionResponse(
                content=content,
                model=data.get("model", model),
                usage=usage,
                finish_reason=finish_reason,
            )

        except httpx.HTTPError as e:
            raise self._handle_error(e) from e

    async def stream_complete(
        self, request: CompletionRequest
    ) -> AsyncIterator[CompletionResponse]:
        """Generate a streaming chat completion.

        Args:
            request: The completion request.

        Yields:
            Completion response chunks.

        Raises:
            ConnectionError: If unable to connect to Kimi API.
            AuthenticationError: If authentication fails.
            ModelNotFoundError: If the model is not available.
            CompletionError: If completion generation fails.
        """
        client = await self._get_client()
        model = self._resolve_model(request.model)
        messages = self._convert_messages(request.messages)

        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": True,
            "temperature": request.temperature,
            "top_p": request.top_p,
        }

        if request.max_tokens:
            payload["max_tokens"] = request.max_tokens

        accumulated_content = ""
        prompt_tokens = 0

        try:
            async with client.stream("POST", "/chat/completions", json=payload) as response:
                response.raise_for_status()

                async for line in response.aiter_lines():
                    line = line.strip()
                    if not line:
                        continue

                    # Handle SSE format: "data: {...}"
                    if line.startswith("data: "):
                        line = line[6:]  # Remove "data: " prefix

                    # Skip SSE comments and keep-alive
                    if line.startswith(":") or not line:
                        continue

                    # Handle stream end
                    if line == "[DONE]":
                        break

                    try:
                        chunk = json.loads(line)
                    except json.JSONDecodeError as e:
                        raise CompletionError(f"Invalid JSON in stream: {line}") from e

                    # Check for errors in stream
                    if "error" in chunk:
                        error_data = chunk["error"]
                        raise CompletionError(
                            f"Stream error: {error_data.get('message', 'Unknown error')}"
                        )

                    choice = chunk.get("choices", [{}])[0]
                    delta = choice.get("delta", {})
                    content_chunk = delta.get("content", "")
                    finish_reason = choice.get("finish_reason")

                    # Update accumulated content
                    accumulated_content += content_chunk

                    # Get usage from final chunk if available
                    usage_data = chunk.get("usage", {})
                    if usage_data:
                        prompt_tokens = usage_data.get("prompt_tokens", prompt_tokens)
                        completion_tokens = usage_data.get("completion_tokens", 0)
                        total_tokens = usage_data.get("total_tokens", 0)
                    else:
                        # Estimate tokens if not provided
                        prompt_tokens = self.count_message_tokens(request.messages, model)
                        completion_tokens = self.count_tokens(accumulated_content, model)
                        total_tokens = prompt_tokens + completion_tokens

                    usage = UsageStats(
                        prompt_tokens=prompt_tokens,
                        completion_tokens=completion_tokens,
                        total_tokens=total_tokens,
                    )

                    # Yield chunk if there's content or it's the final chunk
                    if content_chunk or finish_reason:
                        yield CompletionResponse(
                            content=content_chunk,
                            model=chunk.get("model", model),
                            usage=usage,
                            finish_reason=finish_reason,
                        )

        except httpx.HTTPError as e:
            raise self._handle_error(e) from e

    async def list_models(self) -> list[str]:
        """List available models from Kimi.

        Returns:
            List of available model names.

        Raises:
            ConnectionError: If unable to connect to Kimi API.
            AuthenticationError: If authentication fails.
        """
        client = await self._get_client()

        try:
            response = await client.get("/models")
            response.raise_for_status()
            data = response.json()

            models = []
            for model in data.get("data", []):
                model_id = model.get("id", "")
                if model_id:
                    models.append(model_id)

            return sorted(models)

        except httpx.HTTPError as e:
            raise self._handle_error(e) from e

    async def health_check(self) -> dict[str, Any]:
        """Check the health of the Kimi API connection.

        Returns:
            Health status information including available models.
        """
        result: dict[str, Any] = {
            "status": "unknown",
            "base_url": self.base_url,
            "api_key_configured": bool(self.api_key),
            "models": [],
            "error": None,
        }

        try:
            # Try to list models as a health check
            models = await self.list_models()
            result["models"] = models
            result["status"] = "healthy"
        except AuthenticationError as e:
            result["status"] = "unauthenticated"
            result["error"] = str(e.message)
        except ProviderConnectionError as e:
            result["status"] = "unhealthy"
            result["error"] = str(e.message)
        except Exception as e:
            result["status"] = "error"
            result["error"] = str(e)

        return result

    async def get_token_count(self, messages: list[Message], model: str | None = None) -> int:
        """Get accurate token count from Kimi API.

        This uses the token counting endpoint if available, otherwise
        falls back to local estimation.

        Args:
            messages: List of messages to count tokens for.
            model: Optional model name for accurate counting.

        Returns:
            Token count.
        """
        # Kimi doesn't have a dedicated token counting endpoint,
        # so we use the local estimation method
        return self.count_message_tokens(messages, model)
