import os
import json
import httpx
from fastapi import FastAPI, Request, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, JSONResponse

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# टोकन और सीक्रेट्स
SECRET_TOKEN = "uISssvLLmXK0bsrcFICh14o7aj4q2ZT4"
API_KEY = os.getenv("SECRET_API_KEY")
API_URL = os.getenv("SECRET_API_URL", "https://integrate.api.nvidia.com/v1/chat/completions")

# आपके द्वारा चुना गया नया डिफ़ॉल्ट मॉडल
MODEL_NAME = os.getenv("SECRET_MODEL_NAME")
SYSTEM_PROMPT = os.getenv("SECRET_SYSTEM_PROMPT")

@app.get("/")
async def root_health_check():
    return {"status": "Aura Gen 2.0 is online"}

@app.post("/v1/chat/completions")
@app.post("/")
async def poe_bot_endpoint(request: Request):
    # 1. सुरक्षा: आपके टोकन की जांच
    auth_header = request.headers.get("Authorization")
    if auth_header != f"Bearer {SECRET_TOKEN}":
        raise HTTPException(status_code=401, detail="Unauthorized: Invalid Token")

    data = await request.json()
    req_type = data.get("type")

    # 2. POE हेल्थ चेक (Settings)
    if req_type == "settings":
        return JSONResponse(content={
            "server_bot_dependencies": {},
            "allow_attachments": True,
            "introduction_message": "I am tiranga. I am proud to be an Indian AI assistant. What's on your mind today? Is there anything that I can help you with? PLEASE FEEL FREE TO ASK"
        })

    # 3. यूज़र का मैसेज (Query)
    if req_type == "query":
        poe_messages = data.get("query", [])
        
        # Poe के फॉर्मेट को NVIDIA (OpenAI) फॉर्मेट में बदलना
        openai_messages = [{"role": "system", "content": SYSTEM_PROMPT}]
        for msg in poe_messages:
            role = "assistant" if msg.get("role") == "bot" else msg.get("role", "user")
            openai_messages.append({"role": role, "content": msg.get("content", "")})
        
        payload = {
            "model": MODEL_NAME,
            "messages": openai_messages,
            "max_tokens": 4096,
            "temperature": 1.0,
            "stream": True,
            "chat_template_kwargs": {"enable_thinking": True}
        }

        async def stream_generator():
            async with httpx.AsyncClient(timeout=120.0) as client:
                try:
                    async with client.stream(
                        "POST", 
                        API_URL, 
                        headers={"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"},
                        json=payload
                    ) as response:
                        
                        # यदि NVIDIA से कोई एरर आता है
                        if response.status_code != 200:
                            error_text = await response.aread()
                            yield f"event: text\ndata: {json.dumps({'text': f'⚠️ NVIDIA API Error ({response.status_code}): {error_text.decode()}'})}\n\n"
                            yield "event: done\ndata: {}\n\n"
                            return

                        # स्ट्रीमिंग को Poe के फॉर्मेट में बदलना 
                        async for line in response.aiter_lines():
                            if line.startswith("data: "):
                                data_str = line[6:].strip()
                                if data_str == "[DONE]":
                                    continue
                                try:
                                    parsed = json.loads(data_str)
                                    # पायथन डिक्शनरी सिंटैक्स
                                    if "choices" in parsed and len(parsed["choices"]) > 0:
                                        delta = parsed["choices"][0].get("delta", {})
                                        token = delta.get("content", "")
                                        
                                        if token:
                                            poe_event = {"text": token}
                                            yield f"event: text\ndata: {json.dumps(poe_event)}\n\n"
                                except Exception as e:
                                    pass # आंशिक डेटा को इग्नोर करें
                except Exception as req_e:
                    yield f"event: text\ndata: {json.dumps({'text': f'⚠️ Server Connection Error: {str(req_e)}'})}\n\n"
                
                # स्ट्रीम खत्म होने का संकेत
                yield "event: done\ndata: {}\n\n"

        return StreamingResponse(stream_generator(), media_type="text/event-stream")

    return JSONResponse(status_code=400, content={"error": "Unsupported request type"})

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