from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import httpx
import logging
import sys
import traceback

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", handlers=[logging.StreamHandler(sys.stdout)])
logger = logging.getLogger(__name__)

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

class MessageRequest(BaseModel):
    message: str

BRAIN_URL = "http://opencode-brain:5000"
KNOWLEDGE_URL = "http://knowledge-service:8080/query"
AUTH = httpx.BasicAuth("opencode", "sam4jo")

@app.post("/chat")
async def chat(request: MessageRequest):
    user_msg = request.message.lower()
    timeout_long = httpx.Timeout(180.0, connect=10.0)
    timeout_short = httpx.Timeout(5.0, connect=2.0)
    
    context = ""
    # Check for keywords to trigger Librarian (DB) lookup
    if any(kw in user_msg for kw in ["sam", "hobby", "music", "guitar", "skiing", "experience"]):
        logger.info("Gateway: Consulting Librarian (DB)...")
        async with httpx.AsyncClient(timeout=timeout_short) as client:
            try:
                k_res = await client.post(KNOWLEDGE_URL, json={"question": request.message})
                if k_res.status_code == 200:
                    context = k_res.json().get("context", "")
            except Exception as e:
                logger.warning(f"Gateway: Librarian offline/slow: {str(e)}")

    # Forward to Brain (LLM)
    async with httpx.AsyncClient(auth=AUTH, timeout=timeout_long) as brain_client:
        try:
            session_res = await brain_client.post(f"{BRAIN_URL}/session", json={"title": "Demo"})
            session_id = session_res.json()["id"]
            final_prompt = f"CONTEXT:\n{context}\n\nUSER: {request.message}" if context else request.message
            response = await brain_client.post(f"{BRAIN_URL}/session/{session_id}/message", json={"parts": [{"type": "text", "text": final_prompt}]})
            
            # FIX: Iterate through parts array to find text response
            data = response.json()
            if "parts" in data:
                for part in data["parts"]:
                    if part.get("type") == "text" and "text" in part:
                        return {"response": part["text"]}
            
            return {"response": "AI responded but no text found in expected format."}
        except Exception:
            logger.error(f"Gateway: Brain failure: {traceback.format_exc()}")
            return {"response": "Error: The Brain is taking too long or is disconnected."}