main.py

import os
import asyncio
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional, List
import httpx
from datetime import datetime
import uuid

app = FastAPI(title="AI Team Chat API")

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

# ─────────────────────────────────────────────
# ENV VARS (set in HuggingFace Space secrets)
# ─────────────────────────────────────────────
GROQ_API_KEY   = os.getenv("GROQ_API_KEY", "")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "")
SUPABASE_URL   = os.getenv("SUPABASE_URL", "")
SUPABASE_KEY   = os.getenv("SUPABASE_KEY", "")

# ─────────────────────────────────────────────
# SUPABASE — lazy init (never crashes startup)
# ─────────────────────────────────────────────
_supabase_client = None

def get_supabase():
    global _supabase_client
    if _supabase_client is not None:
        return _supabase_client
    url = os.getenv("SUPABASE_URL", "")
    key = os.getenv("SUPABASE_KEY", "")
    if not url or not key:
        return None
    try:
        from supabase import create_client
        _supabase_client = create_client(url, key)
        print("Supabase connected successfully.")
        return _supabase_client
    except Exception as e:
        print(f"Supabase init error (non-fatal): {e}")
        return None


# ─────────────────────────────────────────────
# MODELS
# ─────────────────────────────────────────────
class ChatRequest(BaseModel):
    message: str
    provider: str = "groq"
    religion: Optional[str] = None
    session_id: Optional[str] = None
    conversation_history: Optional[List[dict]] = []


class AgentResponse(BaseModel):
    agent: str
    role: str
    avatar: str
    color: str
    message: str


class ChatResponse(BaseModel):
    session_id: str
    agent_responses: List[AgentResponse]
    summary: str
    question: str


# ─────────────────────────────────────────────
# LLM WRAPPER
# ─────────────────────────────────────────────
async def call_llm(provider: str, system_prompt: str, user_message: str, temperature: float = 0.7) -> str:
    if provider == "groq":
        return await call_groq(system_prompt, user_message, temperature)
    elif provider == "openai":
        return await call_openai(system_prompt, user_message, temperature)
    else:
        raise HTTPException(status_code=400, detail=f"Unknown provider: {provider}")


async def call_groq(system_prompt: str, user_message: str, temperature: float) -> str:
    key = os.getenv("GROQ_API_KEY", "")
    if not key:
        raise HTTPException(status_code=500, detail="GROQ_API_KEY not set in environment")

    async with httpx.AsyncClient(timeout=30) as client:
        response = await client.post(
            "https://api.groq.com/openai/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {key}",
                "Content-Type": "application/json",
            },
            json={
                "model": "llama-3.3-70b-versatile",
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_message},
                ],
                "temperature": temperature,
                "max_tokens": 300,
            },
        )
        response.raise_for_status()
        data = response.json()
        return data["choices"][0]["message"]["content"].strip()


async def call_openai(system_prompt: str, user_message: str, temperature: float) -> str:
    key = os.getenv("OPENAI_API_KEY", "")
    if not key:
        raise HTTPException(status_code=500, detail="OPENAI_API_KEY not set in environment")

    async with httpx.AsyncClient(timeout=30) as client:
        response = await client.post(
            "https://api.openai.com/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {key}",
                "Content-Type": "application/json",
            },
            json={
                "model": "gpt-4o-mini",
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_message},
                ],
                "temperature": temperature,
                "max_tokens": 300,
            },
        )
        response.raise_for_status()
        data = response.json()
        return data["choices"][0]["message"]["content"].strip()


# ─────────────────────────────────────────────
# AGENT DEFINITIONS
# ─────────────────────────────────────────────
def get_agents(religion: Optional[str]) -> List[dict]:
    spiritual_note = (
        f"Base your guidance on {religion} principles and teachings. Be respectful and calm."
        if religion and religion.lower() not in ["none", "prefer not to say", ""]
        else "Provide neutral mindfulness and universal spiritual guidance. Avoid referencing any specific religion."
    )

    return [
        {
            "name": "Dr. Sarah",
            "role": "Doctor",
            "avatar": "🩺",
            "color": "#4FC3F7",
            "system_prompt": (
                "You are Dr. Sarah, a careful and responsible medical advisor. "
                "You DO NOT give diagnoses. You suggest possibilities carefully and always recommend "
                "consulting a licensed physician for personal medical decisions. "
                "Be concise, warm, and professional. Respond in 2-4 lines maximum. "
                "Do not repeat what other experts would say."
            ),
        },
        {
            "name": "Coach Marcus",
            "role": "Fitness Coach",
            "avatar": "💪",
            "color": "#81C784",
            "system_prompt": (
                "You are Coach Marcus, an energetic and experienced fitness coach. "
                "You focus on physical activity, movement, exercise routines, and safe training. "
                "Be motivating, practical, and concise. Respond in 2-4 lines maximum. "
                "Do not repeat what other experts would say."
            ),
        },
        {
            "name": "Nina",
            "role": "Nutritionist",
            "avatar": "🥗",
            "color": "#FFB74D",
            "system_prompt": (
                "You are Nina, a certified nutritionist specializing in diet, energy, and food science. "
                "You focus on practical, evidence-based dietary guidance. "
                "Be specific, helpful, and concise. Respond in 2-4 lines maximum. "
                "Do not repeat what other experts would say."
            ),
        },
        {
            "name": "Dr. Mia",
            "role": "Mental Health Coach",
            "avatar": "🧠",
            "color": "#CE93D8",
            "system_prompt": (
                "You are Dr. Mia, an empathetic and supportive mental health coach. "
                "You help with emotional wellbeing, stress, mindset, and psychological patterns. "
                "Be compassionate, grounding, and concise. Respond in 2-4 lines maximum. "
                "Do not repeat what other experts would say."
            ),
        },
        {
            "name": "Sage Aris",
            "role": "Spiritual Coach",
            "avatar": "✨",
            "color": "#F48FB1",
            "system_prompt": (
                f"You are Sage Aris, a gentle and insightful spiritual coach. "
                f"{spiritual_note} "
                f"Be calm, respectful, and uplifting. Respond in 2-4 lines maximum. "
                f"Do not repeat what other experts would say."
            ),
        },
    ]


COORDINATOR_SYSTEM_PROMPT = """You are the Coordinator of an expert AI wellness team consisting of a Doctor, Fitness Coach, Nutritionist, Mental Health Coach, and Spiritual Coach.

Your job:
1. Read all agent responses carefully
2. Write a SHORT summary of the key collective insights (2-3 sentences max)
3. Ask ONE clear, thoughtful, combined question to the user to gather more context

Rules:
- Do NOT repeat the agents' responses verbatim
- Keep it collaborative and warm
- The question should help the team give better advice next time

Output format (strictly follow this):
Summary: <your short summary here>
Question: <your single question here>"""


# ─────────────────────────────────────────────
# SUPABASE HELPERS
# ─────────────────────────────────────────────
async def save_to_supabase(session_id: str, user_message: str, agent_responses: list, summary: str, question: str):
    db = get_supabase()
    if not db:
        return
    try:
        record = {
            "session_id": session_id,
            "user_message": user_message,
            "agent_responses": agent_responses,
            "summary": summary,
            "question": question,
            "created_at": datetime.utcnow().isoformat(),
        }
        db.table("chat_history").insert(record).execute()
    except Exception as e:
        print(f"Supabase save error (non-fatal): {e}")


async def get_session_history(session_id: str) -> list:
    db = get_supabase()
    if not db or not session_id:
        return []
    try:
        result = (
            db.table("chat_history")
            .select("*")
            .eq("session_id", session_id)
            .order("created_at", desc=False)
            .limit(20)
            .execute()
        )
        return result.data or []
    except Exception as e:
        print(f"Supabase fetch error (non-fatal): {e}")
        return []


# ─────────────────────────────────────────────
# MAIN ENDPOINT
# ─────────────────────────────────────────────
@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    session_id = request.session_id or str(uuid.uuid4())
    agents = get_agents(request.religion)

    # Build context from conversation history
    history_context = ""
    if request.conversation_history:
        history_lines = []
        for turn in request.conversation_history[-6:]:
            history_lines.append(f"User: {turn.get('user', '')}")
            if turn.get("question"):
                history_lines.append(f"Team Question: {turn.get('question', '')}")
        history_context = "\n\nPrevious conversation context:\n" + "\n".join(history_lines)

    user_prompt = f"{history_context}\n\nUser's current message: {request.message}"

    # ── Step 1: Run all 5 agents in parallel ──
    async def run_agent(agent: dict) -> AgentResponse:
        message = await call_llm(
            provider=request.provider,
            system_prompt=agent["system_prompt"],
            user_message=user_prompt,
        )
        return AgentResponse(
            agent=agent["name"],
            role=agent["role"],
            avatar=agent["avatar"],
            color=agent["color"],
            message=message,
        )

    agent_results: List[AgentResponse] = await asyncio.gather(*[run_agent(a) for a in agents])

    # ── Step 2: Run Coordinator ──
    all_responses_text = "\n\n".join(
        [f"[{r.role} — {r.agent}]:\n{r.message}" for r in agent_results]
    )
    coordinator_user_prompt = (
        f"User asked: \"{request.message}\"\n\n"
        f"Agent responses:\n{all_responses_text}"
    )

    coordinator_raw = await call_llm(
        provider=request.provider,
        system_prompt=COORDINATOR_SYSTEM_PROMPT,
        user_message=coordinator_user_prompt,
        temperature=0.5,
    )

    # Parse coordinator output
    summary = ""
    question = ""
    for line in coordinator_raw.splitlines():
        if line.lower().startswith("summary:"):
            summary = line[len("summary:"):].strip()
        elif line.lower().startswith("question:"):
            question = line[len("question:"):].strip()

    if not summary:
        summary = coordinator_raw
    if not question:
        question = "Can you share more details so the team can help you better?"

    # ── Step 3: Save to Supabase (non-blocking) ──
    agent_data = [r.dict() for r in agent_results]
    asyncio.create_task(
        save_to_supabase(session_id, request.message, agent_data, summary, question)
    )

    return ChatResponse(
        session_id=session_id,
        agent_responses=agent_results,
        summary=summary,
        question=question,
    )


@app.get("/history/{session_id}")
async def get_history(session_id: str):
    history = await get_session_history(session_id)
    return {"session_id": session_id, "history": history}


@app.get("/health")
async def health():
    db = get_supabase()
    return {
        "status": "ok",
        "groq_configured": bool(os.getenv("GROQ_API_KEY")),
        "openai_configured": bool(os.getenv("OPENAI_API_KEY")),
        "supabase_configured": bool(db),
    }