"""OpenAI-compatible chat completion endpoint."""
import time
from typing import Optional, List, Dict, Any
from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from ..models import ClientKey
from ..middleware.auth import get_current_client
from ..services.router import model_router
from ..services.key_selector import key_selector
from ..services.logger import async_logger, LogEntry
from ..providers.litellm_wrapper import LiteLLMWrapper

router = APIRouter(prefix="/v1", tags=["chat"])


class ChatMessage(BaseModel):
    """Chat message."""
    role: str
    content: str


class ChatCompletionRequest(BaseModel):
    """Chat completion request."""
    model: str
    messages: List[ChatMessage]
    stream: bool = False
    temperature: Optional[float] = None
    max_tokens: Optional[int] = None
    top_p: Optional[float] = None


@router.post("/chat/completions")
async def chat_completions(
    request: ChatCompletionRequest,
    client: ClientKey = Depends(get_current_client)
):
    """OpenAI-compatible chat completions endpoint."""
    start_time = time.time()

    # Route model to provider
    provider, error = await model_router.route_model(request.model)
    if error:
        raise HTTPException(status_code=400, detail={"error": error})

    # Get next API key
    provider_key = await key_selector.get_next_key(provider.id)
    if not provider_key:
        raise HTTPException(
            status_code=503,
            detail={"error": f"No available API key for provider: {provider.name}"}
        )

    # Build messages
    messages = [{"role": m.role, "content": m.content} for m in request.messages]

    # Build kwargs
    kwargs = {}
    if request.temperature is not None:
        kwargs["temperature"] = request.temperature
    if request.max_tokens is not None:
        kwargs["max_tokens"] = request.max_tokens
    if request.top_p is not None:
        kwargs["top_p"] = request.top_p

    try:
        response, prompt_tokens, completion_tokens = await LiteLLMWrapper.chat_completion(
            provider=provider,
            api_key=provider_key.key,
            model=request.model,
            messages=messages,
            stream=request.stream,
            **kwargs
        )

        if request.stream:
            return StreamingResponse(
                LiteLLMWrapper.stream_response(
                    response, client.id, provider.id, request.model, start_time
                ),
                media_type="text/event-stream"
            )
        else:
            # Log successful request
            latency_ms = int((time.time() - start_time) * 1000)
            await async_logger.log(LogEntry(
                client_key_id=client.id,
                provider_id=provider.id,
                model=request.model,
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                latency_ms=latency_ms,
                success=True
            ))

            # Return response
            return response

    except HTTPException:
        raise
    except Exception as e:
        # Log failed request
        latency_ms = int((time.time() - start_time) * 1000)
        await async_logger.log(LogEntry(
            client_key_id=client.id,
            provider_id=provider.id,
            model=request.model,
            latency_ms=latency_ms,
            success=False,
            error_message=str(e)
        ))

        # Handle specific errors
        error_str = str(e).lower()
        if "rate" in error_str or "limit" in error_str:
            raise HTTPException(status_code=429, detail={"error": str(e)})
        raise HTTPException(status_code=500, detail={"error": str(e)})


@router.get("/models")
async def list_models(client: ClientKey = Depends(get_current_client)):
    """List available models."""
    return {
        "object": "list",
        "data": [
            {"id": "gpt-4o", "object": "model", "owned_by": "openai"},
            {"id": "gpt-4o-mini", "object": "model", "owned_by": "openai"},
            {"id": "o1", "object": "model", "owned_by": "openai"},
            {"id": "o1-mini", "object": "model", "owned_by": "openai"},
            {"id": "claude-3-opus", "object": "model", "owned_by": "anthropic"},
            {"id": "claude-3-sonnet", "object": "model", "owned_by": "anthropic"},
            {"id": "claude-3-haiku", "object": "model", "owned_by": "anthropic"},
        ]
    }