endpoints.py

from fastapi import APIRouter, HTTPException, Request, Header, Depends
from fastapi.responses import JSONResponse, StreamingResponse
from datetime import datetime
import uuid
from typing import Optional

from src.core.config import config
from src.core.logging import logger
from src.core.client import OpenAIClient
from src.models.claude import ClaudeMessagesRequest, ClaudeTokenCountRequest
from src.conversion.request_converter import convert_claude_to_openai
from src.conversion.response_converter import (
    convert_openai_to_claude_response,
    convert_openai_streaming_to_claude_with_cancellation,
)
from src.core.model_manager import model_manager

router = APIRouter()

# Create OpenAI client - use env key if available, otherwise create a dummy one
if config.openai_api_key:
    openai_client = OpenAIClient(
        config.openai_api_key,
        config.openai_base_url,
        config.request_timeout,
        api_version=config.azure_api_version,
    )
else:
    # In passthrough mode, create a client with a dummy key for compatibility
    # The actual client will be created per-request with user's API key
    openai_client = OpenAIClient(
        "sk-dummy",  # This won't be used for actual requests
        config.openai_base_url,
        config.request_timeout,
        api_version=config.azure_api_version,
    )

def get_openai_client(user_api_key: Optional[str] = None) -> OpenAIClient:
    """Get OpenAI client - use env key if available, otherwise create with user-provided key."""
    if config.openai_api_key:
        # Proxy mode: use pre-configured client
        return openai_client
    elif user_api_key:
        # Passthrough mode: create client with user-provided key
        if not config.validate_api_key(user_api_key):
            raise HTTPException(
                status_code=401,
                detail="Invalid OpenAI API key format. Please provide a valid API key starting with 'sk-'."
            )
        return OpenAIClient(
            user_api_key,
            config.openai_base_url,
            config.request_timeout,
            api_version=config.azure_api_version,
        )
    else:
        raise HTTPException(
            status_code=401,
            detail="No OpenAI API key available. Please set OPENAI_API_KEY environment variable or provide your OpenAI API key in Authorization header."
        )

async def validate_api_key_and_extract(x_api_key: Optional[str] = Header(None), authorization: Optional[str] = Header(None)):
    """Extract API key and handle validation based on configuration."""
    client_api_key = None
    
    # Extract API key from headers
    if x_api_key:
        client_api_key = x_api_key
    elif authorization and authorization.startswith("Bearer "):
        client_api_key = authorization.replace("Bearer ", "")
    
    # If OPENAI_API_KEY is configured in environment, use proxy validation
    if config.openai_api_key:
        # Proxy mode: validate against ANTHROPIC_API_KEY if configured
        if config.anthropic_api_key:
            if not client_api_key or not config.validate_client_api_key(client_api_key):
                logger.warning(f"Invalid API key provided by client")
                raise HTTPException(
                    status_code=401,
                    detail="Invalid API key. Please provide a valid Anthropic API key."
                )
        # Return None since we'll use env OPENAI_API_KEY
        return None
    else:
        # Passthrough mode: no ANTHROPIC validation, use client key as OpenAI key
        if not client_api_key:
            raise HTTPException(
                status_code=401,
                detail="No API key provided. Please provide your OpenAI API key in Authorization header."
            )
        # Return the client API key to be used as OpenAI key
        return client_api_key

@router.post("/v1/messages")
async def create_message(request: ClaudeMessagesRequest, http_request: Request, user_api_key: str = Depends(validate_api_key_and_extract)):
    try:
        logger.debug(
            f"Processing Claude request: model={request.model}, stream={request.stream}"
        )

        # Generate unique request ID for cancellation tracking
        request_id = str(uuid.uuid4())
        
        # Get OpenAI client (either default or created with user's key)
        openai_client = get_openai_client(user_api_key)

        # Convert Claude request to OpenAI format
        openai_request = convert_claude_to_openai(request, model_manager)

        # Check if client disconnected before processing
        if await http_request.is_disconnected():
            raise HTTPException(status_code=499, detail="Client disconnected")

        if request.stream:
            # Streaming response - wrap in error handling
            try:
                openai_stream = openai_client.create_chat_completion_stream(
                    openai_request, request_id
                )
                return StreamingResponse(
                    convert_openai_streaming_to_claude_with_cancellation(
                        openai_stream,
                        request,
                        logger,
                        http_request,
                        openai_client,
                        request_id,
                    ),
                    media_type="text/event-stream",
                    headers={
                        "Cache-Control": "no-cache",
                        "Connection": "keep-alive",
                        "Access-Control-Allow-Origin": "*",
                        "Access-Control-Allow-Headers": "*",
                    },
                )
            except HTTPException as e:
                # Convert to proper error response for streaming
                logger.error(f"Streaming error: {e.detail}")
                import traceback

                logger.error(traceback.format_exc())
                error_message = openai_client.classify_openai_error(e.detail)
                error_response = {
                    "type": "error",
                    "error": {"type": "api_error", "message": error_message},
                }
                return JSONResponse(status_code=e.status_code, content=error_response)
        else:
            # Non-streaming response
            openai_response = await openai_client.create_chat_completion(
                openai_request, request_id
            )
            claude_response = convert_openai_to_claude_response(
                openai_response, request
            )
            return claude_response
    except HTTPException:
        raise
    except Exception as e:
        import traceback

        logger.error(f"Unexpected error processing request: {e}")
        logger.error(traceback.format_exc())
        error_message = openai_client.classify_openai_error(str(e))
        raise HTTPException(status_code=500, detail=error_message)


@router.post("/v1/messages/count_tokens")
async def count_tokens(request: ClaudeTokenCountRequest, user_api_key: str = Depends(validate_api_key_and_extract)):
    try:
        # For token counting, we'll use a simple estimation
        # In a real implementation, you might want to use tiktoken or similar

        total_chars = 0

        # Count system message characters
        if request.system:
            if isinstance(request.system, str):
                total_chars += len(request.system)
            elif isinstance(request.system, list):
                for block in request.system:
                    if hasattr(block, "text"):
                        total_chars += len(block.text)

        # Count message characters
        for msg in request.messages:
            if msg.content is None:
                continue
            elif isinstance(msg.content, str):
                total_chars += len(msg.content)
            elif isinstance(msg.content, list):
                for block in msg.content:
                    if hasattr(block, "text") and block.text is not None:
                        total_chars += len(block.text)

        # Rough estimation: 4 characters per token
        estimated_tokens = max(1, total_chars // 4)

        return {"input_tokens": estimated_tokens}

    except Exception as e:
        logger.error(f"Error counting tokens: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/health")
async def health_check():
    """Health check endpoint"""
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "openai_api_configured": bool(config.openai_api_key),
        "api_key_valid": config.validate_api_key() if config.openai_api_key else "per_request",
        "client_api_key_validation": bool(config.anthropic_api_key),
    }


@router.get("/test-connection")
async def test_connection():
    """Test API connectivity to OpenAI"""
    try:
        # Use default client if available, otherwise require user API key
        if not config.openai_api_key:
            return JSONResponse(
                status_code=400,
                content={
                    "status": "failed",
                    "message": "No default OpenAI API key configured. Test connection requires OPENAI_API_KEY environment variable.",
                    "timestamp": datetime.now().isoformat(),
                }
            )
        
        # Simple test request to verify API connectivity
        test_response = await openai_client.create_chat_completion(
            {
                "model": config.small_model,
                "messages": [{"role": "user", "content": "Hello"}],
                "max_tokens": 5,
            }
        )

        return {
            "status": "success",
            "message": "Successfully connected to OpenAI API",
            "model_used": config.small_model,
            "timestamp": datetime.now().isoformat(),
            "response_id": test_response.get("id", "unknown"),
        }

    except Exception as e:
        logger.error(f"API connectivity test failed: {e}")
        return JSONResponse(
            status_code=503,
            content={
                "status": "failed",
                "error_type": "API Error",
                "message": str(e),
                "timestamp": datetime.now().isoformat(),
                "suggestions": [
                    "Check your OPENAI_API_KEY is valid",
                    "Verify your API key has the necessary permissions",
                    "Check if you have reached rate limits",
                ],
            },
        )


@router.get("/")
async def root():
    """Root endpoint"""
    return {
        "message": "Claude-to-OpenAI API Proxy v1.0.0",
        "status": "running",
        "config": {
            "openai_base_url": config.openai_base_url,
            "max_tokens_limit": config.max_tokens_limit,
            "api_key_configured": bool(config.openai_api_key),
            "client_api_key_validation": bool(config.anthropic_api_key),
            "big_model": config.big_model,
            "small_model": config.small_model,
        },
        "endpoints": {
            "messages": "/v1/messages",
            "count_tokens": "/v1/messages/count_tokens",
            "health": "/health",
            "test_connection": "/test-connection",
        },
    }