from fastapi import FastAPI, HTTPException # type: ignore
from fastapi.responses import StreamingResponse # type: ignore
from pydantic import BaseModel # type: ignore
from typing import List, Optional, Dict, Any
import json
import asyncio
from llm_proxy import get_llm, LLMProvider, ChatMessage

app = FastAPI(title="Unified LLM API")

# Pydantic models for request/response
class ChatRequest(BaseModel):
    model: str
    messages: List[Dict[str, str]]
    provider: Optional[str] = None
    stream: bool = False
    max_tokens: Optional[int] = None
    temperature: Optional[float] = None

class GenerateRequest(BaseModel):
    model: str
    prompt: str
    provider: Optional[str] = None
    stream: bool = False
    max_tokens: Optional[int] = None
    temperature: Optional[float] = None

class ChatResponseModel(BaseModel):
    content: str
    model: str
    provider: str
    finish_reason: Optional[str] = None
    usage: Optional[Dict[str, int]] = None

@app.post("/chat")
async def chat_endpoint(request: ChatRequest):
    """Chat endpoint that works with any configured LLM provider"""
    try:
        llm = get_llm()
        
        # Convert provider string to enum if provided
        provider = None
        if request.provider:
            try:
                provider = LLMProvider(request.provider.lower())
            except ValueError:
                raise HTTPException(
                    status_code=400, 
                    detail=f"Unsupported provider: {request.provider}"
                )
        
        # Prepare kwargs
        kwargs = {}
        if request.max_tokens:
            kwargs['max_tokens'] = request.max_tokens
        if request.temperature:
            kwargs['temperature'] = request.temperature
        
        if request.stream:
            async def generate_response():
                try:
                    # Use the type-safe streaming method
                    async for chunk in llm.chat_stream(
                        model=request.model,
                        messages=request.messages,
                        provider=provider,
                        **kwargs
                    ):
                        # Format as Server-Sent Events
                        data = {
                            "content": chunk.content,
                            "model": chunk.model,
                            "provider": provider.value if provider else llm.default_provider.value,
                            "finish_reason": chunk.finish_reason
                        }
                        yield f"data: {json.dumps(data)}\n\n"
                except Exception as e:
                    error_data = {"error": str(e)}
                    yield f"data: {json.dumps(error_data)}\n\n"
                finally:
                    yield "data: [DONE]\n\n"
            
            return StreamingResponse(
                generate_response(),
                media_type="text/event-stream",
                headers={
                    "Cache-Control": "no-cache",
                    "Connection": "keep-alive"
                }
            )
        else:
            # Use the type-safe single response method
            response = await llm.chat_single(
                model=request.model,
                messages=request.messages,
                provider=provider,
                **kwargs
            )
            
            return ChatResponseModel(
                content=response.content,
                model=response.model,
                provider=provider.value if provider else llm.default_provider.value,
                finish_reason=response.finish_reason,
                usage=response.usage
            )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/generate")
async def generate_endpoint(request: GenerateRequest):
    """Generate endpoint for simple text generation"""
    try:
        llm = get_llm()
        
        provider = None
        if request.provider:
            try:
                provider = LLMProvider(request.provider.lower())
            except ValueError:
                raise HTTPException(
                    status_code=400, 
                    detail=f"Unsupported provider: {request.provider}"
                )
        
        kwargs = {}
        if request.max_tokens:
            kwargs['max_tokens'] = request.max_tokens
        if request.temperature:
            kwargs['temperature'] = request.temperature
        
        if request.stream:
            async def generate_response():
                try:
                    # Use the type-safe streaming method
                    async for chunk in llm.generate_stream(
                        model=request.model,
                        prompt=request.prompt,
                        provider=provider,
                        **kwargs
                    ):
                        data = {
                            "content": chunk.content,
                            "model": chunk.model,
                            "provider": provider.value if provider else llm.default_provider.value,
                            "finish_reason": chunk.finish_reason
                        }
                        yield f"data: {json.dumps(data)}\n\n"
                except Exception as e:
                    error_data = {"error": str(e)}
                    yield f"data: {json.dumps(error_data)}\n\n"
                finally:
                    yield "data: [DONE]\n\n"
            
            return StreamingResponse(
                generate_response(),
                media_type="text/event-stream"
            )
        else:
            # Use the type-safe single response method
            response = await llm.generate_single(
                model=request.model,
                prompt=request.prompt,
                provider=provider,
                **kwargs
            )
            
            return ChatResponseModel(
                content=response.content,
                model=response.model,
                provider=provider.value if provider else llm.default_provider.value,
                finish_reason=response.finish_reason,
                usage=response.usage
            )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/models")
async def list_models(provider: Optional[str] = None):
    """List available models for a provider"""
    try:
        llm = get_llm()
        
        provider_enum = None
        if provider:
            try:
                provider_enum = LLMProvider(provider.lower())
            except ValueError:
                raise HTTPException(
                    status_code=400, 
                    detail=f"Unsupported provider: {provider}"
                )
        
        models = await llm.list_models(provider_enum)
        return {
            "provider": provider_enum.value if provider_enum else llm.default_provider.value,
            "models": models
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/providers")
async def list_providers():
    """List all configured providers"""
    llm = get_llm()
    return {
        "providers": [provider.value for provider in llm._initialized_providers],
        "default": llm.default_provider.value
    }

@app.post("/providers/{provider}/set-default")
async def set_default_provider(provider: str):
    """Set the default provider"""
    try:
        llm = get_llm()
        provider_enum = LLMProvider(provider.lower())
        llm.set_default_provider(provider_enum)
        return {"message": f"Default provider set to {provider}", "default": provider}
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))

# Health check endpoint
@app.get("/health")
async def health_check():
    """Health check endpoint"""
    llm = get_llm()
    return {
        "status": "healthy",
        "providers_configured": len(llm._initialized_providers),
        "default_provider": llm.default_provider.value
    }

if __name__ == "__main__":
    import uvicorn # type: ignore
    uvicorn.run(app, host="0.0.0.0", port=8000)