import os
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from openai import AsyncOpenAI
from pydantic import BaseModel
import httpx

# Initialize FastAPI app
app = FastAPI()

# Define request body model for the prompt
class PromptRequest(BaseModel):
    prompt: str

# Initialize OpenAI client
token = os.getenv("GITHUB_TOKEN")
if not token:
    raise ValueError("GITHUB_TOKEN environment variable not set")

# Use environment variables for endpoint and model, with fallbacks
endpoint = os.getenv("API_ENDPOINT", "https://api.openai.com/v1")  # Fallback to OpenAI-compatible endpoint
model = os.getenv("MODEL_NAME", "gpt-4o-mini")  # Default to a known model

# Initialize AsyncOpenAI with a custom HTTP client to avoid proxies issue
client = AsyncOpenAI(
    base_url=endpoint,
    api_key=token,
    http_client=httpx.AsyncClient()  # Explicitly use httpx.AsyncClient without proxies
)

# Async generator to stream chunks
async def stream_response(prompt: str):
    try:
        # Create streaming chat completion
        stream = await client.chat.completions.create(
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": prompt}
            ],
            temperature=1.0,
            top_p=1.0,
            model=model,
            stream=True
        )

        # Yield each chunk as it arrives
        async for chunk in stream:
            if chunk.choices and len(chunk.choices) > 0:
                content = chunk.choices[0].delta.content or ""
                if content:
                    yield content

    except Exception as err:
        yield f"Error: {str(err)}"

# Endpoint to handle prompt and stream response
@app.post("/generate")
async def generate_response(request: PromptRequest):
    try:
        return StreamingResponse(
            stream_response(request.prompt),
            media_type="text/event-stream"
        )
    except Exception as err:
        raise HTTPException(status_code=500, detail=f"Server error: {str(err)}")

# Health check endpoint for Hugging Face Spaces
@app.get("/")
async def health_check():
    return {"status": "healthy"}