import os
from fastapi import FastAPI, HTTPException, Query
from fastapi.responses import StreamingResponse
from openai import AsyncOpenAI

app = FastAPI()

# Define available models (you can expand this list)
AVAILABLE_MODELS = {
    "openai/gpt-4.1-mini": "OpenAI GPT-4 Mini",
    "deepseek/DeepSeek-V3-0324": "DeepSeek V3 0324",
    # Add more models as needed
}

async def generate_ai_response(prompt: str, model: str):
    # Configuration for unofficial GitHub AI endpoint
    token = os.getenv("GITHUB_TOKEN")
    if not token:
        raise HTTPException(status_code=500, detail="GitHub token not configured")
    
    endpoint = "https://models.github.ai/inference"
    
    # Validate the model
    if model not in AVAILABLE_MODELS:
        raise HTTPException(status_code=400, detail=f"Model not available. Choose from: {', '.join(AVAILABLE_MODELS.keys())}")

    client = AsyncOpenAI(base_url=endpoint, api_key=token)

    try:
        stream = await client.chat.completions.create(
            messages=[
                
                {"role": "user", "content": prompt}
            ],
            model=model,
            temperature=1.0,
            top_p=1.0,
            stream=True
        )

        async for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                yield chunk.choices[0].delta.content

    except Exception as err:
        yield f"Error: {str(err)}"
        raise HTTPException(status_code=500, detail="AI generation failed")

@app.post("/generate")
async def generate_response(
    prompt: str = Query(..., description="The prompt for the AI"),
    model: str = Query("openai/gpt-4.1-mini", description="The model to use for generation")
):
    if not prompt:
        raise HTTPException(status_code=400, detail="Prompt cannot be empty")
    
    return StreamingResponse(
        generate_ai_response(prompt, model),
        media_type="text/event-stream"
    )

def get_app():
    return app