from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from main import telecom_agent, telecom_agent_stream
import asyncio
import json

app = FastAPI(title="Telecom GenAI Agent")

# CORS - adjust origins in production
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # restrict in production
    allow_methods=["*"],
    allow_headers=["*"],
)

def extract_user_input(raw_query) -> str:
    """
    Normalizes all incoming query formats into a clean string.
    """
    if isinstance(raw_query, list) and raw_query:
        last_msg = raw_query[-1]
        if isinstance(last_msg, dict):
            return (
                last_msg.get("content")
                or (last_msg.get("parts", [{}])[0].get("text") if isinstance(last_msg.get("parts"), list) else "")
                or str(last_msg)
            ).strip()
        return str(last_msg).strip()
    return str(raw_query).strip()     
# Non-streaming endpoint
@app.post("/query")
async def query_agent(request: Request):
    try:
        data = await request.json()
    except Exception:
        return JSONResponse({"error": "Invalid JSON"}, status_code=400)
    session = data.get("session", "default")
    raw_query = data.get("query", "")
    user_input = extract_user_input(raw_query)
    if not user_input:
        return JSONResponse({"error": "query is required"}, status_code=400)
    if len(user_input) > 2000:
        return JSONResponse({"error": "query too long"}, status_code=413)
    try:
        answer, evaluation, analytics = await asyncio.to_thread(
            telecom_agent,
            user_input,
            session=session
        )
        return JSONResponse({
            "answer": answer,
            "evaluation": evaluation,
            "analytics": analytics
        })
    except Exception as e:
        return JSONResponse({"error": "Agent failed"}, status_code=500)         
# Streaming endpoint (NDJSON)
@app.post("/stream")
async def stream_agent(request: Request):
    try:
        data = await request.json()
    except Exception:
        return JSONResponse({"error": "Invalid JSON"}, status_code=400)
    session = data.get("session", "default")
    raw_query = data.get("query", "")
    user_input = extract_user_input(raw_query)
    if not user_input:
        return JSONResponse({"error": "query is required"}, status_code=400)
    if len(user_input) > 2000:
        return JSONResponse({"error": "query too long"}, status_code=413)
    try:
        gen = telecom_agent_stream(user_input, session=session)
    except Exception:
        return JSONResponse({"error": "Agent failed to start stream"}, status_code=500)
    async def event_generator():
        try:           
            for chunk in gen:
                if isinstance(chunk, dict):
                    payload = chunk.get("text", chunk)
                else:
                    payload = chunk            
                yield json.dumps({"text": str(payload)}) + "\n"
                await asyncio.sleep(0)
        except Exception:
            pass 
        finally:
            if hasattr(gen, "close"):
                gen.close()
    return StreamingResponse(
        event_generator(),
        media_type="application/x-ndjson",
        headers={"Cache-Control": "no-store"}
    )
    
@app.get("/")
async def root():
    return {"message": "Telecom AI running"} 
    
# Health check
@app.get("/health")
async def health():
    return {"status": "ok", "ram_safe": "8GB ready"}

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