api/routes.py

import asyncio
import importlib.util
import json
import logging
import sys
import uuid
from datetime import datetime, timezone
from pathlib import Path

logger = logging.getLogger(__name__)

# ── SDK shadowing fix ────────────────────────────────────────────────────────
# Load rag_agent.py directly via importlib so its own SDK-fix code runs first,
# setting sys.modules['agents'] to the openai-agents SDK.  After exec_module()
# returns, `from agents import Runner` resolves to the SDK, not the local package.
_rag_agent_path = Path(__file__).parent.parent / "agents" / "rag_agent.py"
_spec = importlib.util.spec_from_file_location("agents.rag_agent", str(_rag_agent_path))
_rag_mod = importlib.util.module_from_spec(_spec)
sys.modules["agents.rag_agent"] = _rag_mod
_spec.loader.exec_module(_rag_mod)  # runs rag_agent.py → SDK is now sys.modules['agents']
create_rag_agent = _rag_mod.create_rag_agent
get_run_config = _rag_mod.get_run_config

# Now safe to import SDK symbols — agents points to the SDK in sys.modules
from agents import Runner  # noqa: E402
from agents.stream_events import RawResponsesStreamEvent  # noqa: E402
from fastapi import APIRouter, HTTPException  # noqa: E402
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse  # noqa: E402
from openai.types.chat import ChatCompletionChunk  # noqa: E402
from openai.types.responses import ResponseTextDeltaEvent  # noqa: E402
from pydantic import BaseModel  # noqa: E402

from cache.query_cache import _cache as _disk_cache  # noqa: E402
from cache.query_cache import find_similar_cached, get_cache_stats, save_to_cache  # noqa: E402
from db.history import list_sessions, load_history, save_turn  # noqa: E402
from ingestion import ingest_all_documents  # noqa: E402
from scraper.content_store import is_content_stale, load_latest_content  # noqa: E402
from tools.list_documents_tool import list_indexed_documents  # noqa: E402

router = APIRouter()


# ── request / response models ────────────────────────────────────────────────

class HistoryMessage(BaseModel):
    role: str        # "user" | "assistant"
    content: str


class ChatRequest(BaseModel):
    message: str
    session_id: str = ""
    history: list[HistoryMessage] = []


class FeedbackRequest(BaseModel):
    session_id: str = ""
    message: str          # the user question that was answered
    rating: int           # 1 (bad) – 5 (excellent)
    comment: str = ""     # optional free-text comment


# ── helpers ──────────────────────────────────────────────────────────────────

def _build_input(db_history: list[dict], message: str) -> list[dict]:
    """Merge DB history with the current message into Agents SDK input format."""
    messages = [{"role": t["role"], "content": t["content"]} for t in db_history]
    messages.append({"role": "user", "content": message})
    return messages


async def _stream_and_cache(request: ChatRequest, db_history: list[dict]):
    """Stream token deltas as SSE, then persist the turn and save to cache."""
    # Signal immediately that the agent is working so the frontend can show a spinner
    yield f"data: {json.dumps({'status': 'thinking'})}\n\n"

    full_text: list[str] = []
    try:
        agent = create_rag_agent()
        result = Runner.run_streamed(
            agent,
            _build_input(db_history, request.message),
            run_config=get_run_config(),
        )

        async for event in result.stream_events():
            if not isinstance(event, RawResponsesStreamEvent):
                continue
            data = event.data
            delta: str | None = None
            if isinstance(data, ResponseTextDeltaEvent):
                # Responses API (OpenAI native)
                delta = data.delta
            elif isinstance(data, ChatCompletionChunk):
                # Chat Completions API (Gemini compatible)
                delta = data.choices[0].delta.content if data.choices else None
            if delta:
                full_text.append(delta)
                yield f"data: {json.dumps({'delta': delta})}\n\n"

    except Exception as exc:
        error_msg = (
            "I'm sorry, I ran into a technical issue while processing your request. "
            "Please try again, or contact AgentRax support if the problem persists."
        )
        yield f"data: {json.dumps({'error': error_msg, 'detail': str(exc)})}\n\n"

    finally:
        yield "data: [DONE]\n\n"
        if full_text:
            response_text = "".join(full_text)
            save_to_cache(request.message, response_text, [])
            if request.session_id:
                await asyncio.to_thread(save_turn, request.session_id, request.message, response_text)


# ── routes ───────────────────────────────────────────────────────────────────

@router.post("/chat")
async def chat(request: ChatRequest):
    """Return AgentRax answer, served from semantic cache when possible.

    Flow:
    1. Load conversation history from PostgreSQL by session_id.
    2. Check semantic cache (threshold 0.92). On HIT: return immediately, save turn.
    3. On MISS: stream agent response, save to cache + DB on completion.

    Cache HIT  → JSON  { answer, sources }  with X-Cache: HIT  (zero LLM cost)
    Cache MISS → SSE token stream            with X-Cache: MISS
    """
    # Load DB history first — used by both HIT and MISS paths for context
    db_history = await asyncio.to_thread(load_history, request.session_id)

    cached = find_similar_cached(request.message)
    if cached:
        # Still persist the cached answer so history is complete
        if request.session_id:
            await asyncio.to_thread(
                save_turn, request.session_id, request.message, cached["answer"]
            )
        return JSONResponse(
            content={"answer": cached["answer"], "sources": cached.get("sources", [])},
            headers={"X-Cache": "HIT"},
        )

    return StreamingResponse(
        _stream_and_cache(request, db_history),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
            "X-Cache": "MISS",
        },
    )


@router.get("/sessions")
async def get_sessions():
    """Return all sessions ordered by last activity, with title and message count."""
    try:
        sessions = await asyncio.to_thread(list_sessions)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    # Serialize datetime objects to ISO strings
    return {
        "sessions": [
            {**s, "last_active": s["last_active"].isoformat() if s.get("last_active") else None}
            for s in sessions
        ]
    }


@router.get("/history/{session_id}")
async def get_history(session_id: str):
    """Return full conversation history for a session from PostgreSQL."""
    try:
        turns = await asyncio.to_thread(load_history, session_id, 100)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    return {"session_id": session_id, "history": turns}


@router.get("/cache/stats")
async def cache_stats():
    """Return hit counts and top queries from the semantic cache."""
    return get_cache_stats()


@router.post("/cache/clear")
async def cache_clear():
    """Wipe all entries from the semantic cache (admin use only)."""
    _disk_cache.clear()
    return {"status": "ok", "message": "Cache cleared."}


@router.get("/site/status")
async def site_status():
    """Return metadata about the last scraped AgentRax snapshot."""
    content = load_latest_content()
    if not content:
        return {"status": "no_data", "is_stale": True}
    return {
        "title": content.get("title"),
        "url": content.get("url"),
        "scraped_at": content.get("scraped_at") or content.get("saved_at"),
        "is_stale": is_content_stale(),
    }


@router.post("/ingest")
async def ingest():
    """Scan documents/ and index all .pdf, .docx, and .md files into ChromaDB."""
    try:
        result = await asyncio.to_thread(ingest_all_documents)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc

    return {
        "status": "ok",
        "documents_indexed": result.get("ingested", 0),
        "total_chunks": result.get("total_chunks", 0),
    }


@router.post("/ingest-frontend-docs")
async def ingest_frontend_docs():
    """Index .md documentation files from the AgentraX frontend docs/ folder.

    Automatically locates D:/Agentrax-Updated-Backend/AgentraX_Frontend/docs/
    and indexes all Markdown files (ERD, FRD, implementation notes, etc.)
    into ChromaDB so the help agent can answer detailed technical questions.
    """
    # Path: api/routes.py → api/ → AgentraXhelpAgent/ → Agentrax-Updated-Backend/ → AgentraX_Frontend/docs/
    frontend_docs = Path(__file__).parent.parent.parent / "AgentraX_Frontend" / "docs"

    if not frontend_docs.exists():
        raise HTTPException(
            status_code=404,
            detail=f"Frontend docs directory not found at: {frontend_docs}",
        )

    md_files = list(frontend_docs.glob("*.md"))
    if not md_files:
        return {"status": "ok", "message": "No .md files found in frontend docs.", "documents_indexed": 0}

    try:
        result = await asyncio.to_thread(ingest_all_documents, [frontend_docs])
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc

    return {
        "status": "ok",
        "source": str(frontend_docs),
        "documents_indexed": result.get("ingested", 0),
        "total_chunks": result.get("total_chunks", 0),
        "files": result.get("files", []),
    }


@router.get("/documents")
async def documents():
    """Return all source file names currently indexed in the vector store."""
    try:
        docs = await asyncio.to_thread(list_indexed_documents)
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc

    return {"documents": docs}


@router.post("/feedback")
async def feedback(request: FeedbackRequest):
    """Accept user feedback (1–5 star rating) on a chat response.

    Feedback is appended to feedback_log.jsonl in the project root.
    Useful for identifying gaps in the knowledge base or poor responses.
    """
    if not (1 <= request.rating <= 5):
        raise HTTPException(status_code=422, detail="Rating must be between 1 and 5.")

    entry = {
        "id": str(uuid.uuid4()),
        "session_id": request.session_id,
        "message": request.message,
        "rating": request.rating,
        "comment": request.comment,
        "submitted_at": datetime.now(timezone.utc).isoformat(),
    }

    log_path = Path(__file__).parent.parent / "feedback_log.jsonl"
    try:
        with log_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(entry, ensure_ascii=False) + "\n")
    except Exception as exc:
        logger.exception("Failed to write feedback entry.")
        raise HTTPException(status_code=500, detail="Could not save feedback.") from exc

    return {"status": "ok", "message": "Thank you for your feedback!"}


@router.get("/health")
async def health():
    """Liveness check."""
    return {"status": "ok"}