from fastapi import FastAPI, HTTPException, Depends, Header, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from typing import List, Optional, Literal
import json
import g4f
from g4f.Provider import OpenaiAccount, RetryProvider
from g4f.models import ModelUtils

app = FastAPI()

# Complete list of available models from G4F
models = [
    # OpenAI models
    "gpt-4", "gpt-4-turbo", "gpt-4o", "gpt-3.5-turbo",
    
    # Anthropic models
    "claude-3-opus", "claude-3-sonnet", "claude-3-haiku", "claude-2.1",
    
    # Google models
    "gemini-pro", "gemini-1.5-pro", "gemini-1.5-flash",
    
    # Meta models
    "llama-2-70b", "llama-2-13b", "llama-2-7b", "llama-3-70b", "llama-3-8b",
    
    # Other providers
    "mistral-7b", "mixtral-8x7b", "command-r-plus", "cohere-command-r",
    "deepseek-chat", "deepseek-coder", "code-llama-34b", "code-llama-70b",
    
    # Specialized models
    "grok-1", "grok-1.5", "grok-2", "o1", "o3-mini", "flux", "flux-pro"
]

# Configure G4F backend
class CustomBackend(g4f.Provider.BackendApi):
    working = True
    ssl = False
    url = "https://ahe.hopto.org"
    headers = {"Authorization": "Basic Z2dnOmc0Zl8="}
    
    @classmethod
    async def create_async_generator(
        cls,
        model: str,
        messages: g4f.typing.Messages,
        **kwargs
    ) -> g4f.typing.AsyncResult:
        if model in OpenaiAccount.get_models():
            kwargs["provider"] = OpenaiAccount
        async for chunk in super().create_async_generator(model, messages, **kwargs):
            yield chunk

# Pydantic models
class Message(BaseModel):
    role: Literal["system", "user", "assistant"]
    content: str

class ChatRequest(BaseModel):
    model: str
    messages: List[Message]
    temperature: Optional[float] = 0.7
    max_tokens: Optional[int] = None
    top_p: Optional[float] = 0.9
    streaming: bool = True

class ChatResponse(BaseModel):
    role: str = "assistant"
    content: str
    model: Optional[str] = None

class ModelListResponse(BaseModel):
    models: List[str]

# API Key Verification
async def verify_api_key(x_api_key: str = Header(...)):
    if x_api_key != "fb207532285886a5568298b4b4e61124":
        raise HTTPException(status_code=403, detail="Invalid API key")

@app.get("/v1/models", response_model=ModelListResponse, tags=["Models"])
async def get_models():
    """Get list of all available models"""
    return ModelListResponse(models=models)

@app.post("/v1/chat/completions", response_model=Optional[ChatResponse], tags=["Chat"])
async def chat_completion(
    request: ChatRequest,
    api_key: str = Depends(verify_api_key)
):
    """
    Handle chat completion requests with streaming support.
    
    Args:
        request: ChatRequest containing model, messages and parameters
        api_key: Verified API key
        
    Returns:
        Either a StreamingResponse or direct ChatResponse
    """
    # Validate model
    if request.model not in models:
        raise HTTPException(
            status_code=400,
            detail=f"Invalid model. Available models: {', '.join(models)}"
        )

    # Prepare messages
    messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
    
    try:
        if request.streaming:
            async def stream_generator():
                response = await g4f.ChatCompletion.create_async(
                    model=request.model,
                    messages=messages,
                    temperature=request.temperature,
                    top_p=request.top_p,
                    max_tokens=request.max_tokens,
                    provider=RetryProvider([CustomBackend])
                )
                
                async for chunk in response:
                    if isinstance(chunk, dict):
                        yield f"data: {json.dumps(chunk)}\n\n"
                    else:
                        yield f"data: {json.dumps({'content': str(chunk)})}\n\n"
                yield "data: [DONE]\n\n"

            return StreamingResponse(
                stream_generator(),
                media_type="text/event-stream"
            )
        else:
            response = await g4f.ChatCompletion.create_async(
                model=request.model,
                messages=messages,
                temperature=request.temperature,
                top_p=request.top_p,
                max_tokens=request.max_tokens,
                provider=RetryProvider([CustomBackend])
            )
            
            if isinstance(response, str):
                return ChatResponse(content=response, model=request.model)
            elif isinstance(response, dict):
                return ChatResponse(
                    content=response.get("choices", [{}])[0].get("message", {}).get("content", ""),
                    model=request.model
                )
            return ChatResponse(content=str(response), model=request.model

    except Exception as e:
        raise HTTPException(
            status_code=500,
            detail=f"Error processing request: {str(e)}"
        )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)