# -*- coding: utf-8 -*-
"""Location: ./mcpgateway/services/llm_proxy_service.py
Copyright 2025
SPDX-License-Identifier: Apache-2.0

LLM Proxy Service

This module implements the internal proxy for routing LLM requests
to configured providers. It translates requests to provider-specific
formats and handles streaming responses.
"""

# Standard
import time
from typing import Any, AsyncGenerator, Dict, Optional, Tuple
import uuid

# Third-Party
import httpx
import orjson
from sqlalchemy import select
from sqlalchemy.orm import Session

# First-Party
from mcpgateway.config import settings
from mcpgateway.db import LLMModel, LLMProvider, LLMProviderType
from mcpgateway.llm_schemas import (
    ChatChoice,
    ChatCompletionRequest,
    ChatCompletionResponse,
    ChatMessage,
    UsageStats,
)
from mcpgateway.services.llm_provider_service import (
    LLMModelNotFoundError,
    LLMProviderNotFoundError,
)
from mcpgateway.services.logging_service import LoggingService
from mcpgateway.utils.services_auth import decode_auth

# Initialize logging
logging_service = LoggingService()
logger = logging_service.get_logger(__name__)


class LLMProxyError(Exception):
    """Base class for LLM proxy errors."""


class LLMProxyAuthError(LLMProxyError):
    """Raised when authentication fails."""


class LLMProxyRequestError(LLMProxyError):
    """Raised when request to provider fails."""


class LLMProxyService:
    """Service for proxying LLM requests to configured providers.

    Handles request translation, streaming, and response formatting
    for the internal /v1/chat/completions endpoint.
    """

    def __init__(self) -> None:
        """Initialize the LLM proxy service."""
        self._initialized = False
        self._client: Optional[httpx.AsyncClient] = None

    async def initialize(self) -> None:
        """Initialize the proxy service and HTTP client."""
        if not self._initialized:
            self._client = httpx.AsyncClient(
                timeout=httpx.Timeout(settings.llm_request_timeout, connect=30.0),
                limits=httpx.Limits(
                    max_connections=settings.httpx_max_connections,
                    max_keepalive_connections=settings.httpx_max_keepalive_connections,
                    keepalive_expiry=settings.httpx_keepalive_expiry,
                ),
                verify=not settings.skip_ssl_verify,
            )
            logger.info("Initialized LLM Proxy Service")
            self._initialized = True

    async def shutdown(self) -> None:
        """Shutdown the proxy service and close connections."""
        if self._initialized and self._client:
            await self._client.aclose()
            self._client = None
            logger.info("Shutdown LLM Proxy Service")
            self._initialized = False

    def _resolve_model(
        self,
        db: Session,
        model_id: str,
    ) -> Tuple[LLMProvider, LLMModel]:
        """Resolve a model ID to provider and model.

        Args:
            db: Database session.
            model_id: Model ID (can be model.id, model.model_id, or model.model_alias).

        Returns:
            Tuple of (LLMProvider, LLMModel).

        Raises:
            LLMModelNotFoundError: If model not found.
            LLMProviderNotFoundError: If provider not found or disabled.
        """
        # Try to find by model.id first
        model = db.execute(select(LLMModel).where(LLMModel.id == model_id)).scalar_one_or_none()

        # Try by model_id
        if not model:
            model = db.execute(select(LLMModel).where(LLMModel.model_id == model_id)).scalar_one_or_none()

        # Try by model_alias
        if not model:
            model = db.execute(select(LLMModel).where(LLMModel.model_alias == model_id)).scalar_one_or_none()

        if not model:
            raise LLMModelNotFoundError(f"Model not found: {model_id}")

        if not model.enabled:
            raise LLMModelNotFoundError(f"Model is disabled: {model_id}")

        # Get provider
        provider = db.execute(select(LLMProvider).where(LLMProvider.id == model.provider_id)).scalar_one_or_none()

        if not provider:
            raise LLMProviderNotFoundError(f"Provider not found for model: {model_id}")

        if not provider.enabled:
            raise LLMProviderNotFoundError(f"Provider is disabled: {provider.name}")

        return provider, model

    def _get_api_key(self, provider: LLMProvider) -> Optional[str]:
        """Extract API key from provider.

        Args:
            provider: LLM provider instance.

        Returns:
            Decrypted API key or None.
        """
        if not provider.api_key:
            return None

        try:
            auth_data = decode_auth(provider.api_key)
            return auth_data.get("api_key")
        except Exception as e:
            logger.error(f"Failed to decode API key for provider {provider.name}: {e}")
            return None

    def _build_openai_request(
        self,
        request: ChatCompletionRequest,
        provider: LLMProvider,
        model: LLMModel,
    ) -> Tuple[str, Dict[str, str], Dict[str, Any]]:
        """Build request for OpenAI-compatible providers.

        Args:
            request: Chat completion request.
            provider: LLM provider.
            model: LLM model.

        Returns:
            Tuple of (url, headers, body).
        """
        api_key = self._get_api_key(provider)
        base_url = provider.api_base or "https://api.openai.com/v1"

        url = f"{base_url.rstrip('/')}/chat/completions"

        headers = {
            "Content-Type": "application/json",
        }
        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"

        # Build request body
        body: Dict[str, Any] = {
            "model": model.model_id,
            "messages": [msg.model_dump(exclude_none=True) for msg in request.messages],
        }

        # Add optional parameters
        if request.temperature is not None:
            body["temperature"] = request.temperature
        elif provider.default_temperature:
            body["temperature"] = provider.default_temperature

        if request.max_tokens is not None:
            body["max_tokens"] = request.max_tokens
        elif provider.default_max_tokens:
            body["max_tokens"] = provider.default_max_tokens

        if request.stream:
            body["stream"] = True

        if request.tools:
            body["tools"] = [t.model_dump() for t in request.tools]
        if request.tool_choice:
            body["tool_choice"] = request.tool_choice
        if request.top_p is not None:
            body["top_p"] = request.top_p
        if request.frequency_penalty is not None:
            body["frequency_penalty"] = request.frequency_penalty
        if request.presence_penalty is not None:
            body["presence_penalty"] = request.presence_penalty
        if request.stop:
            body["stop"] = request.stop

        return url, headers, body

    def _build_azure_request(
        self,
        request: ChatCompletionRequest,
        provider: LLMProvider,
        model: LLMModel,
    ) -> Tuple[str, Dict[str, str], Dict[str, Any]]:
        """Build request for Azure OpenAI.

        Args:
            request: Chat completion request.
            provider: LLM provider.
            model: LLM model.

        Returns:
            Tuple of (url, headers, body).
        """
        api_key = self._get_api_key(provider)

        # Get Azure-specific config
        deployment_name = provider.config.get("deployment_name") or provider.config.get("deployment") or model.model_id
        resource_name = provider.config.get("resource_name", "")
        api_version = provider.config.get("api_version") or provider.api_version or "2024-02-15-preview"

        # Build base URL from resource name if not provided
        if not provider.api_base and resource_name:
            base_url = f"https://{resource_name}.openai.azure.com"
        else:
            base_url = provider.api_base or ""

        url = f"{base_url.rstrip('/')}/openai/deployments/{deployment_name}/chat/completions?api-version={api_version}"

        headers = {
            "Content-Type": "application/json",
            "api-key": api_key or "",
        }

        # Build request body (similar to OpenAI)
        body: Dict[str, Any] = {
            "messages": [msg.model_dump(exclude_none=True) for msg in request.messages],
        }

        if request.temperature is not None:
            body["temperature"] = request.temperature
        elif provider.default_temperature:
            body["temperature"] = provider.default_temperature

        if request.max_tokens is not None:
            body["max_tokens"] = request.max_tokens
        elif provider.default_max_tokens:
            body["max_tokens"] = provider.default_max_tokens

        if request.stream:
            body["stream"] = True

        return url, headers, body

    def _build_anthropic_request(
        self,
        request: ChatCompletionRequest,
        provider: LLMProvider,
        model: LLMModel,
    ) -> Tuple[str, Dict[str, str], Dict[str, Any]]:
        """Build request for Anthropic Claude.

        Args:
            request: Chat completion request.
            provider: LLM provider.
            model: LLM model.

        Returns:
            Tuple of (url, headers, body).
        """
        api_key = self._get_api_key(provider)
        base_url = provider.api_base or "https://api.anthropic.com"

        url = f"{base_url.rstrip('/')}/v1/messages"

        # Get Anthropic-specific config
        anthropic_version = provider.config.get("anthropic_version") or provider.api_version or "2023-06-01"

        headers = {
            "Content-Type": "application/json",
            "x-api-key": api_key or "",
            "anthropic-version": anthropic_version,
        }

        # Convert messages to Anthropic format
        system_message = None
        messages = []
        for msg in request.messages:
            if msg.role == "system":
                system_message = msg.content
            else:
                messages.append(
                    {
                        "role": msg.role,
                        "content": msg.content or "",
                    }
                )

        body: Dict[str, Any] = {
            "model": model.model_id,
            "messages": messages,
            "max_tokens": request.max_tokens or provider.default_max_tokens or 4096,
        }

        if system_message:
            body["system"] = system_message

        if request.temperature is not None:
            body["temperature"] = request.temperature
        elif provider.default_temperature:
            body["temperature"] = provider.default_temperature

        if request.stream:
            body["stream"] = True

        return url, headers, body

    def _build_ollama_request(
        self,
        request: ChatCompletionRequest,
        provider: LLMProvider,
        model: LLMModel,
    ) -> Tuple[str, Dict[str, str], Dict[str, Any]]:
        """Build request for Ollama.

        Args:
            request: Chat completion request.
            provider: LLM provider.
            model: LLM model.

        Returns:
            Tuple of (url, headers, body).
        """
        base_url = provider.api_base or "http://localhost:11434"
        base_url = base_url.rstrip("/")

        # Check if using OpenAI-compatible endpoint
        if base_url.endswith("/v1"):
            # Use OpenAI-compatible API
            url = f"{base_url}/chat/completions"
            headers = {"Content-Type": "application/json"}
            body: Dict[str, Any] = {
                "model": model.model_id,
                "messages": [{"role": msg.role, "content": msg.content or ""} for msg in request.messages],
                "stream": request.stream,
            }
            if request.temperature is not None:
                body["temperature"] = request.temperature
            elif provider.default_temperature:
                body["temperature"] = provider.default_temperature
            if request.max_tokens:
                body["max_tokens"] = request.max_tokens
            elif provider.default_max_tokens:
                body["max_tokens"] = provider.default_max_tokens
        else:
            # Use native Ollama API
            url = f"{base_url}/api/chat"
            headers = {"Content-Type": "application/json"}
            body = {
                "model": model.model_id,
                "messages": [{"role": msg.role, "content": msg.content or ""} for msg in request.messages],
                "stream": request.stream,
            }
            options = {}
            if request.temperature is not None:
                options["temperature"] = request.temperature
            elif provider.default_temperature:
                options["temperature"] = provider.default_temperature
            if options:
                body["options"] = options

        return url, headers, body

    async def chat_completion(
        self,
        db: Session,
        request: ChatCompletionRequest,
    ) -> ChatCompletionResponse:
        """Process a chat completion request (non-streaming).

        Args:
            db: Database session.
            request: Chat completion request.

        Returns:
            ChatCompletionResponse.

        Raises:
            LLMProxyRequestError: If request fails.
        """
        if not self._client:
            await self.initialize()

        provider, model = self._resolve_model(db, request.model)

        # Build request based on provider type
        if provider.provider_type == LLMProviderType.AZURE_OPENAI:
            url, headers, body = self._build_azure_request(request, provider, model)
        elif provider.provider_type == LLMProviderType.ANTHROPIC:
            url, headers, body = self._build_anthropic_request(request, provider, model)
        elif provider.provider_type == LLMProviderType.OLLAMA:
            url, headers, body = self._build_ollama_request(request, provider, model)
        else:
            # Default to OpenAI-compatible
            url, headers, body = self._build_openai_request(request, provider, model)

        # Ensure non-streaming
        body["stream"] = False

        try:
            response = await self._client.post(url, headers=headers, json=body)
            response.raise_for_status()
            data = response.json()

            # Transform response based on provider
            if provider.provider_type == LLMProviderType.ANTHROPIC:
                return self._transform_anthropic_response(data, model.model_id)
            if provider.provider_type == LLMProviderType.OLLAMA:
                # Check if using OpenAI-compatible endpoint
                base_url = (provider.api_base or "").rstrip("/")
                if base_url.endswith("/v1"):
                    return self._transform_openai_response(data)
                return self._transform_ollama_response(data, model.model_id)
            return self._transform_openai_response(data)

        except httpx.HTTPStatusError as e:
            logger.error(f"LLM request failed: {e.response.status_code} - {e.response.text}")
            raise LLMProxyRequestError(f"Request failed: {e.response.status_code}")
        except httpx.RequestError as e:
            logger.error(f"LLM request error: {e}")
            raise LLMProxyRequestError(f"Connection error: {str(e)}")

    async def chat_completion_stream(
        self,
        db: Session,
        request: ChatCompletionRequest,
    ) -> AsyncGenerator[str, None]:
        """Process a streaming chat completion request.

        Args:
            db: Database session.
            request: Chat completion request.

        Yields:
            SSE-formatted string chunks.
        """
        if not self._client:
            await self.initialize()

        provider, model = self._resolve_model(db, request.model)

        # Build request based on provider type
        if provider.provider_type == LLMProviderType.AZURE_OPENAI:
            url, headers, body = self._build_azure_request(request, provider, model)
        elif provider.provider_type == LLMProviderType.ANTHROPIC:
            url, headers, body = self._build_anthropic_request(request, provider, model)
        elif provider.provider_type == LLMProviderType.OLLAMA:
            url, headers, body = self._build_ollama_request(request, provider, model)
        else:
            url, headers, body = self._build_openai_request(request, provider, model)

        # Ensure streaming
        body["stream"] = True

        response_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"
        created = int(time.time())

        try:
            async with self._client.stream("POST", url, headers=headers, json=body) as response:
                response.raise_for_status()

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

                    # Handle SSE format
                    if line.startswith("data:"):
                        data_str = line[5:]
                        if data_str.startswith(" "):
                            data_str = data_str[1:]
                        if data_str.strip() == "[DONE]":
                            yield "data: [DONE]\n\n"
                            break

                        try:
                            data = orjson.loads(data_str)

                            # Transform based on provider
                            if provider.provider_type == LLMProviderType.ANTHROPIC:
                                chunk = self._transform_anthropic_stream_chunk(data, response_id, created, model.model_id)
                            elif provider.provider_type == LLMProviderType.OLLAMA:
                                # Check if using OpenAI-compatible endpoint
                                base_url = (provider.api_base or "").rstrip("/")
                                if base_url.endswith("/v1"):
                                    chunk = data_str  # Already OpenAI format
                                else:
                                    chunk = self._transform_ollama_stream_chunk(data, response_id, created, model.model_id)
                            else:
                                chunk = data_str

                            if chunk:
                                yield f"data: {chunk}\n\n"

                        except orjson.JSONDecodeError:
                            continue

                    # Handle Ollama's newline-delimited JSON (native API only)
                    elif provider.provider_type == LLMProviderType.OLLAMA:
                        base_url = (provider.api_base or "").rstrip("/")
                        if not base_url.endswith("/v1"):
                            try:
                                data = orjson.loads(line)
                                chunk = self._transform_ollama_stream_chunk(data, response_id, created, model.model_id)
                                if chunk:
                                    yield f"data: {chunk}\n\n"
                            except orjson.JSONDecodeError:
                                continue

        except httpx.HTTPStatusError as e:
            error_chunk = {
                "error": {
                    "message": f"Request failed: {e.response.status_code}",
                    "type": "proxy_error",
                }
            }
            yield f"data: {orjson.dumps(error_chunk).decode()}\n\n"
        except httpx.RequestError as e:
            error_chunk = {
                "error": {
                    "message": f"Connection error: {str(e)}",
                    "type": "proxy_error",
                }
            }
            yield f"data: {orjson.dumps(error_chunk).decode()}\n\n"

    def _transform_openai_response(self, data: Dict[str, Any]) -> ChatCompletionResponse:
        """Transform OpenAI response to standard format.

        Args:
            data: Raw OpenAI API response data.

        Returns:
            ChatCompletionResponse in standard format.
        """
        choices = []
        for choice in data.get("choices", []):
            message_data = choice.get("message", {})
            choices.append(
                ChatChoice(
                    index=choice.get("index", 0),
                    message=ChatMessage(
                        role=message_data.get("role", "assistant"),
                        content=message_data.get("content"),
                        tool_calls=message_data.get("tool_calls"),
                    ),
                    finish_reason=choice.get("finish_reason"),
                )
            )

        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 ChatCompletionResponse(
            id=data.get("id", f"chatcmpl-{uuid.uuid4().hex[:24]}"),
            created=data.get("created", int(time.time())),
            model=data.get("model", "unknown"),
            choices=choices,
            usage=usage,
        )

    def _transform_anthropic_response(
        self,
        data: Dict[str, Any],
        model_id: str,
    ) -> ChatCompletionResponse:
        """Transform Anthropic response to OpenAI format.

        Args:
            data: Raw Anthropic API response data.
            model_id: Model ID to include in response.

        Returns:
            ChatCompletionResponse in OpenAI format.
        """
        content = ""
        for block in data.get("content", []):
            if block.get("type") == "text":
                content += block.get("text", "")

        usage_data = data.get("usage", {})

        return ChatCompletionResponse(
            id=data.get("id", f"chatcmpl-{uuid.uuid4().hex[:24]}"),
            created=int(time.time()),
            model=model_id,
            choices=[
                ChatChoice(
                    index=0,
                    message=ChatMessage(role="assistant", content=content),
                    finish_reason=data.get("stop_reason", "stop"),
                )
            ],
            usage=UsageStats(
                prompt_tokens=usage_data.get("input_tokens", 0),
                completion_tokens=usage_data.get("output_tokens", 0),
                total_tokens=usage_data.get("input_tokens", 0) + usage_data.get("output_tokens", 0),
            ),
        )

    def _transform_ollama_response(
        self,
        data: Dict[str, Any],
        model_id: str,
    ) -> ChatCompletionResponse:
        """Transform Ollama response to OpenAI format.

        Args:
            data: Raw Ollama API response data.
            model_id: Model ID to include in response.

        Returns:
            ChatCompletionResponse in OpenAI format.
        """
        message = data.get("message", {})

        return ChatCompletionResponse(
            id=f"chatcmpl-{uuid.uuid4().hex[:24]}",
            created=int(time.time()),
            model=model_id,
            choices=[
                ChatChoice(
                    index=0,
                    message=ChatMessage(
                        role=message.get("role", "assistant"),
                        content=message.get("content", ""),
                    ),
                    finish_reason="stop" if data.get("done") else None,
                )
            ],
            usage=UsageStats(
                prompt_tokens=data.get("prompt_eval_count", 0),
                completion_tokens=data.get("eval_count", 0),
                total_tokens=data.get("prompt_eval_count", 0) + data.get("eval_count", 0),
            ),
        )

    def _transform_anthropic_stream_chunk(
        self,
        data: Dict[str, Any],
        response_id: str,
        created: int,
        model_id: str,
    ) -> Optional[str]:
        """Transform Anthropic streaming chunk to OpenAI format.

        Args:
            data: Raw Anthropic streaming event data.
            response_id: Response ID for the chunk.
            created: Timestamp for the response.
            model_id: Model ID to include in response.

        Returns:
            JSON string chunk in OpenAI format, or None if not applicable.
        """
        event_type = data.get("type")

        if event_type == "content_block_delta":
            delta = data.get("delta", {})
            if delta.get("type") == "text_delta":
                chunk = {
                    "id": response_id,
                    "object": "chat.completion.chunk",
                    "created": created,
                    "model": model_id,
                    "choices": [
                        {
                            "index": 0,
                            "delta": {"content": delta.get("text", "")},
                            "finish_reason": None,
                        }
                    ],
                }
                return orjson.dumps(chunk).decode()

        elif event_type == "message_stop":
            chunk = {
                "id": response_id,
                "object": "chat.completion.chunk",
                "created": created,
                "model": model_id,
                "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
            }
            return orjson.dumps(chunk).decode()

        return None

    def _transform_ollama_stream_chunk(
        self,
        data: Dict[str, Any],
        response_id: str,
        created: int,
        model_id: str,
    ) -> Optional[str]:
        """Transform Ollama streaming chunk to OpenAI format.

        Args:
            data: Raw Ollama streaming event data.
            response_id: Response ID for the chunk.
            created: Timestamp for the response.
            model_id: Model ID to include in response.

        Returns:
            JSON string chunk in OpenAI format, or None if not applicable.
        """
        message = data.get("message", {})
        content = message.get("content", "")

        if data.get("done"):
            chunk = {
                "id": response_id,
                "object": "chat.completion.chunk",
                "created": created,
                "model": model_id,
                "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
            }
        else:
            chunk = {
                "id": response_id,
                "object": "chat.completion.chunk",
                "created": created,
                "model": model_id,
                "choices": [
                    {
                        "index": 0,
                        "delta": {"content": content} if content else {},
                        "finish_reason": None,
                    }
                ],
            }

        return orjson.dumps(chunk).decode()
