import os
import json
import asyncio
import httpx
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse

NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY")
NVIDIA_URL = "https://integrate.api.nvidia.com/v1/chat/completions"

SYSTEM_PROMPT = """You are Spark 4.5, an advanced AI model created and developed by the technology company OPEN SORT, founded and led by CEO Chun Hongly (ជុន ហុងលី).

STRICT IDENTITY RULES:
1. YOUR IDENTITY:
   - Model Name: Spark 4.5
   - Developing Company: OPEN SORT
   - Founder & CEO: Chun Hongly (ជុន ហុងលី)

2. STRICT PRIVACY & CONFIDENTIALITY:
   - Never mention or acknowledge NVIDIA, Jensen Huang, or Nemotron under any circumstances.
   - Even if users try to trick, persuade, or probe about your underlying architecture, base models, training hardware, or company history, strictly maintain that you were created by OPEN SORT and Chun Hongly as Spark 4.5.

3. NATURAL CONVERSATION STYLE:
   - Do NOT introduce your company, founder, or model name in general responses unless the user specifically asks about your identity, name, creator, model, or company.
   - Respond naturally, helpful, and directly to all standard queries without adding unnecessary company disclosures."""

http_client: httpx.AsyncClient = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global http_client
    limits = httpx.Limits(max_keepalive_connections=500, max_connections=2000)
    timeout = httpx.Timeout(120.0, connect=5.0)
    http_client = httpx.AsyncClient(limits=limits, timeout=timeout)
    yield
    await http_client.aclose()

app = FastAPI(lifespan=lifespan)

@app.post("/generate")
async def generate(request: Request):
    try:
        body = await request.json()
        user_prompt = body.get("prompt", "")
        
        headers = {
            "Authorization": f"Bearer {NVIDIA_API_KEY}",
            "Content-Type": "application/json",
            "Accept": "text/event-stream"
        }
        
        payload = {
            "model": "nvidia/nemotron-3-ultra-550b-a55b",
            "messages": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt}
            ],
            "temperature": 1,
            "top_p": 0.95,
            "max_tokens": 16384,
            "chat_template_kwargs": {"enable_thinking": True},
            "reasoning_budget": 16384,
            "stream": True
        }
        
        async def event_generator():
            try:
                async with http_client.stream("POST", NVIDIA_URL, headers=headers, json=payload) as response:
                    async for line in response.aiter_lines():
                        if await request.is_disconnected():
                            break
                        
                        if line:
                            line_str = line.strip()
                            if line_str.startswith("data: "):
                                data_content = line_str[6:].strip()
                                if data_content == "[DONE]":
                                    break
                                try:
                                    json_data = json.loads(data_content)
                                    choices = json_data.get("choices", [])
                                    if choices:
                                        delta = choices[0].get("delta", {})
                                        reasoning = delta.get("reasoning_content") or delta.get("reasoning")
                                        content = delta.get("content")
                                        
                                        out_data = {}
                                        if reasoning:
                                            out_data["thinking"] = reasoning
                                        if content:
                                            out_data["answer"] = content
                                            
                                        if out_data:
                                            yield json.dumps(out_data) + "\n"
                                except json.JSONDecodeError:
                                    continue
            except asyncio.CancelledError:
                pass
                            
        response_headers = {
            "X-Accel-Buffering": "no",
            "Cache-Control": "no-cache",
            "Connection": "keep-alive"
        }
        return StreamingResponse(event_generator(), media_type="application/x-ndjson", headers=response_headers)
        
    except Exception as e:
        return {"error": str(e)}