backend/server/endpoints.py

import asyncio
import os
import time

from fastapi import APIRouter, Header, HTTPException
from typing import Optional
import datetime
from backend.adk.agents.cfe.agent import get_cfe_agent
from backend.adk.agents.intent_classification.agent import get_ic_agent
from backend.adk.agents.recsys.agent import get_recsys_agent
from backend.adk.assembly.pipeline import get_root_agent
from backend.adk.assembly.run import _call_agent_async, get_model_response
from backend.schema.cfe import CFEOutput, CFEContext
from backend.schema.recSys import RecsysOutput, RecommendationContext
from backend.schema.intentClassifier import IntentClassificationOutput
from backend.schema.cqGen import CQOutput
from backend.adk.agents.clar_q_gen.cq_generator import generate_clarifying_questions
import json
from utils.firestore_utils import ingest_response_firestore, get_firestore_client

router = APIRouter(tags=["ADK Endpoints"])

# Serialize concurrent requests that swap GOOGLE_API_KEY — acceptable for a demo
_gemini_key_lock = asyncio.Lock()


_VERTEXAI_ENV_VARS = (
    "GOOGLE_GENAI_USE_VERTEXAI",
    "GOOGLE_GENAI_USE_ENTERPRISE",
    "GOOGLE_CLOUD_PROJECT",
    "GOOGLE_CLOUD_LOCATION",
    "GOOGLE_APPLICATION_CREDENTIALS",
)


async def _with_gemini_key(coro, api_key: Optional[str]):
    """Run *coro* using a user-supplied Gemini API key.

    Temporarily sets GOOGLE_API_KEY and clears any Vertex AI env vars so that
    google-genai uses Gemini API key auth instead of Vertex AI credentials.
    """
    if not api_key:
        return await coro
    async with _gemini_key_lock:
        orig_api_key = os.environ.get("GOOGLE_API_KEY")
        orig_vertexai = {k: os.environ.pop(k, None) for k in _VERTEXAI_ENV_VARS}
        os.environ["GOOGLE_API_KEY"] = api_key
        try:
            return await coro
        finally:
            if orig_api_key is not None:
                os.environ["GOOGLE_API_KEY"] = orig_api_key
            else:
                os.environ.pop("GOOGLE_API_KEY", None)
            for k, v in orig_vertexai.items():
                if v is not None:
                    os.environ[k] = v


# ---------------------------------------------------------------------------
# Clarifying questions
# ---------------------------------------------------------------------------

@router.post("/generate-clarifying-questions", response_model=CQOutput)
async def get_clarifying_questions(
    user_input: str,
    x_model_provider: str = Header(default="gemma", alias="X-Model-Provider"),
    x_api_key: Optional[str] = Header(default=None, alias="X-Api-Key"),
):
    """Step 1: generate clarifying questions for a user query."""
    try:
        if x_model_provider == "gemma":
            from backend.llm.gemma_pipeline import generate_cq
            return await generate_cq(user_input)

        # Gemini path
        async def _gemini():
            return await generate_clarifying_questions(user_input)

        return await _with_gemini_key(_gemini(), x_api_key)

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ---------------------------------------------------------------------------
# Intent classifier (standalone endpoint — called only from Gemini path)
# ---------------------------------------------------------------------------

@router.get("/intent-classifier", response_model=IntentClassificationOutput)
async def get_intent_classifier_response(
    session_id: str,
    x_model_provider: str = Header(default="gemini", alias="X-Model-Provider"),
    x_api_key: Optional[str] = Header(default=None, alias="X-Api-Key"),
):
    try:
        async def _gemini():
            model_init = await get_ic_agent()
            agent_name, response_text = None, None
            async for name, text in _call_agent_async(
                query=f"[SESSION_ID:{session_id}]",
                root_agent=model_init,
                session_id=session_id,
            ):
                agent_name, response_text = name, text

            response = json.loads(response_text)
            response["session_id"] = session_id
            ingestion_success = await ingest_response_firestore(
                "intent_classifier_responses", session_id, response
            )
            response["db_ingestion_status"] = ingestion_success
            return IntentClassificationOutput(**response)

        return await _with_gemini_key(_gemini(), x_api_key)

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ---------------------------------------------------------------------------
# Recommender (standalone — called only from Gemini path)
# ---------------------------------------------------------------------------

@router.get("/recommender-output", response_model=RecsysOutput)
async def get_recommender_response(
    session_id: str,
    has_context: bool = True,
    x_model_provider: str = Header(default="gemini", alias="X-Model-Provider"),
    x_api_key: Optional[str] = Header(default=None, alias="X-Api-Key"),
):
    try:
        async def _gemini():
            if has_context:
                model_init = await get_recsys_agent(has_context=True)
                collection_name = "context_aware_recommendations"
            else:
                model_init = await get_recsys_agent(has_context=False)
                collection_name = "baseline_recommendations"

            agent_name, response_text = None, None
            async for name, text in _call_agent_async(
                query=f"[SESSION_ID:{session_id}]",
                root_agent=model_init,
                session_id=session_id,
            ):
                agent_name, response_text = name, text

            if not response_text or not response_text.strip():
                raise HTTPException(status_code=502, detail="Empty response from recommendation agent")

            response = json.loads(response_text)
            ingestion_success = await ingest_response_firestore(collection_name, session_id, response)
            response["db_ingestion_status"] = ingestion_success
            return RecsysOutput(**response)

        return await _with_gemini_key(_gemini(), x_api_key)

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ---------------------------------------------------------------------------
# CFE (standalone — called only from Gemini path)
# ---------------------------------------------------------------------------

@router.get("/cfe-output", response_model=CFEOutput)
async def get_cfe_response(
    session_id: str,
    x_model_provider: str = Header(default="gemini", alias="X-Model-Provider"),
    x_api_key: Optional[str] = Header(default=None, alias="X-Api-Key"),
):
    try:
        async def _gemini():
            model_init = await get_cfe_agent()
            agent_name, response_text = None, None
            async for name, text in _call_agent_async(
                query=f"[SESSION_ID:{session_id}]",
                root_agent=model_init,
                session_id=session_id,
            ):
                agent_name, response_text = name, text

            response = json.loads(response_text)
            ingestion_success = await ingest_response_firestore("cfe_responses", session_id, response)
            response["db_ingestion_status"] = ingestion_success
            return CFEOutput(**response)

        return await _with_gemini_key(_gemini(), x_api_key)

    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=str(e))


# ---------------------------------------------------------------------------
# Full pipeline (main entry point called by the frontend)
# ---------------------------------------------------------------------------

@router.get("/run-pipeline", response_model=CFEOutput)
async def run_pipeline(
    session_id: str,
    x_model_provider: str = Header(default="gemma", alias="X-Model-Provider"),
    x_api_key: Optional[str] = Header(default=None, alias="X-Api-Key"),
):
    try:
        start_time = time.time()

        if x_model_provider == "gemma":
            from backend.llm.gemma_pipeline import run_full_pipeline as gemma_run
            cfe_output = await gemma_run(session_id)
            cfe_output.time_taken_seconds = time.time() - start_time
            return cfe_output

        # Gemini path
        async def _gemini():
            model_init = await get_root_agent()
            result = await get_model_response(
                query=f"[USER QUERY]: {session_id}]",
                root_agent=model_init,
                session_id=session_id,
                return_cfe_only=True,
            )
            if result is None:
                raise HTTPException(status_code=404, detail="CFE response not found in pipeline")
            result.time_taken_seconds = time.time() - start_time

            response_dict = result.model_dump()
            ingestion_success = await ingest_response_firestore(
                "cfe_pipeline_responses", session_id, response_dict
            )
            result.db_ingestion_status = ingestion_success
            return result

        return await _with_gemini_key(_gemini(), x_api_key)

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/run-pipeline-all-responses")
async def run_pipeline_all_responses(session_id: str):
    try:
        model_init = await get_root_agent()
        all_responses = await get_model_response(
            query=f"[SESSION_ID:{session_id}]",
            root_agent=model_init,
            session_id=session_id,
            return_cfe_only=False,
        )
        return {"session_id": session_id, "responses": all_responses}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))