backend/api/routes_chat.py

"""
Chat endpoints — streaming EA Advisor with session persistence and DRL enrichment.
Powered by AMD MI300X + Knowledge Graph RAG.
"""

import asyncio
import json
import logging
import uuid
from typing import Annotated

from fastapi import APIRouter, Depends, Query, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from backend.config import Settings, get_settings
from backend.dependencies import get_neo4j_client, get_llm_client
from backend.graph.neo4j_client import Neo4jClient
from backend.graph.cypher_queries import (
    ENSURE_CHAT_SESSION,
    APPEND_CHAT_EXCHANGE,
    GET_SESSION_MESSAGES,
    GET_RECENT_CHAT_SESSIONS,
    DELETE_CHAT_SESSION,
    CHECK_DOMAIN_DRL_STATUS,
)
from backend.llm.client import LLMClient

log = logging.getLogger(__name__)
router = APIRouter()

CHAT_SYSTEM = """You are an Enterprise Architecture Advisor powered by AMD MI300X and Qwen-72B.
You have access to a knowledge graph containing 1,416 enterprise capabilities across 44 strategic domains,
44 governance standards, and 44 digital transformation trends.
Answer questions about enterprise architecture, digital transformation, governance, compliance, and technology strategy.
Be specific and actionable. Cite specific capabilities, governance standards, and trends from the context provided.
When relevant, mention which domain or subdomain a capability belongs to.
If asked to generate a roadmap, suggest the user use the Strategic Roadmap tab for a full structured output."""


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _classify_intent(message: str) -> str:
    msg = message.lower()
    if any(kw in msg for kw in ["roadmap", "transformation plan", "implementation plan"]):
        return "generate_roadmap"
    if any(kw in msg for kw in ["show me", "list", "what domains", "explore"]):
        return "explore_domain"
    return "general_qa"


async def _build_rag_context(
    message: str,
    neo4j: Neo4jClient,
    llm: LLMClient,
) -> tuple[str, list[dict], list[str]]:
    """
    Run graph RAG.
    Returns (context_str, sources_list, domain_names_retrieved).
    """
    from backend.agents.retriever import RetrieverAgent
    retriever = RetrieverAgent(neo4j=neo4j, llm=llm)
    try:
        caps = await retriever.retrieve(org_type=message[:80], goals=[message], sectors=[])
    except Exception as exc:
        log.warning(f"RAG retrieval failed: {exc}")
        caps = []

    sources: list[dict] = []
    context_lines: list[str] = []
    domain_names: list[str] = []

    for c in caps[:6]:
        cap = c.capability
        std = c.standard or {}
        trend = c.trend or {}
        name = cap.get("name", "")
        domain = c.domain.get("name", "")
        outcomes = (cap.get("business_outcomes") or [])[:2]

        context_lines.append(
            f"- Capability: {name} (Domain: {domain})\n"
            f"  Outcomes: {'; '.join(outcomes)}\n"
            f"  Governance: {std.get('name', '')}\n"
            f"  Trend: {trend.get('name', '')}"
        )
        sources.append({
            "name": name,
            "domain": domain,
            "standard": std.get("name", ""),
            "trend": trend.get("name", ""),
        })
        if domain and domain not in domain_names:
            domain_names.append(domain)

    context = "\n".join(context_lines) if context_lines else "No specific capabilities matched."
    return context, sources, domain_names


def _build_messages(history: list[dict], user_message: str, context: str) -> list[dict]:
    messages = []
    for msg in history[-8:]:
        if msg.get("role") in ("user", "assistant") and msg.get("content"):
            # Skip the initial welcome assistant message to avoid invalid conversation order
            if msg["role"] == "assistant" and "Enterprise Architecture Advisor" in msg["content"]:
                continue
            messages.append({"role": msg["role"], "content": msg["content"]})
    messages.append({
        "role": "user",
        "content": f"KNOWLEDGE GRAPH CONTEXT:\n{context}\n\nQUESTION: {user_message}",
    })
    return messages


def _get_untrained_domains(domain_names: list[str], neo4j: Neo4jClient) -> list[str]:
    """Return subset of domain_names that have not yet been DRL-trained."""
    if not domain_names:
        return []
    try:
        rows = neo4j.run_query(CHECK_DOMAIN_DRL_STATUS, domain_names=domain_names)
        return [r["name"] for r in rows if not r.get("trained")]
    except Exception as exc:
        log.warning(f"DRL status check failed: {exc}")
        return []


async def _trigger_drl_enrichment(domain_names: list[str]) -> str:
    """Fire-and-forget background DRL training for untrained domains."""
    run_id = uuid.uuid4().hex[:10]

    async def _run():
        try:
            import os
            from dotenv import load_dotenv
            load_dotenv()
            from neo4j import GraphDatabase
            from pipeline.train_on_graph import GraphTrainer

            driver = GraphDatabase.driver(
                os.getenv("NEO4J_URI"),
                auth=(os.getenv("NEO4J_USERNAME"), os.getenv("NEO4J_PASSWORD")),
            )
            try:
                trainer = GraphTrainer(driver, os.getenv("NEO4J_DATABASE", "neo4j"))
                loop = asyncio.get_event_loop()
                for domain in domain_names:
                    await loop.run_in_executor(
                        None,
                        lambda d=domain: trainer.run(episodes_per_domain=50, target_domain=d),
                    )
                log.info(f"DRL enrichment {run_id} complete for: {domain_names}")
            finally:
                driver.close()
        except Exception as exc:
            log.warning(f"DRL enrichment {run_id} failed: {exc}")

    asyncio.create_task(_run())
    return run_id


def _store_exchange(
    neo4j: Neo4jClient,
    session_id: str,
    user_content: str,
    asst_content: str,
    sources: list[dict],
    title: str,
) -> None:
    """Persist a user↔assistant exchange to Neo4j."""
    try:
        neo4j.run_query(ENSURE_CHAT_SESSION, session_id=session_id, title=title)
        neo4j.run_query(
            APPEND_CHAT_EXCHANGE,
            session_id=session_id,
            user_msg_id=uuid.uuid4().hex,
            user_content=user_content,
            asst_msg_id=uuid.uuid4().hex,
            asst_content=asst_content,
            sources_json=json.dumps(sources),
        )
    except Exception as exc:
        log.warning(f"Failed to store chat exchange: {exc}")


# ---------------------------------------------------------------------------
# Session management endpoints
# ---------------------------------------------------------------------------

@router.post("/chat/sessions")
async def create_session(
    neo4j: Annotated[Neo4jClient, Depends(get_neo4j_client)],
    session_id: str | None = None,
    title: str = "",
):
    """Create or touch a chat session. Returns session_id."""
    sid = session_id or uuid.uuid4().hex
    try:
        neo4j.run_query(ENSURE_CHAT_SESSION, session_id=sid, title=title or "New Conversation")
    except Exception as exc:
        log.warning(f"Session create failed: {exc}")
    return {"session_id": sid}


@router.get("/chat/sessions")
async def list_sessions(neo4j: Annotated[Neo4jClient, Depends(get_neo4j_client)]):
    """Return the 15 most recent chat sessions."""
    try:
        rows = neo4j.run_query(GET_RECENT_CHAT_SESSIONS)
        return [
            {
                "session_id": r["session_id"],
                "title": r.get("title", "Untitled"),
                "last_active": r.get("last_active", ""),
                "message_count": r.get("message_count", 0),
            }
            for r in rows
        ]
    except Exception as exc:
        log.warning(f"List sessions failed: {exc}")
        return []


@router.get("/chat/sessions/{session_id}/messages")
async def get_session(
    session_id: str,
    neo4j: Annotated[Neo4jClient, Depends(get_neo4j_client)],
):
    """Return all messages for a session."""
    try:
        rows = neo4j.run_query(GET_SESSION_MESSAGES, session_id=session_id)
        messages = []
        for r in rows:
            srcs = []
            try:
                srcs = json.loads(r.get("sources_json", "[]"))
            except Exception:
                pass
            messages.append({
                "role": r["role"],
                "content": r["content"],
                "sources": srcs,
                "created_at": r.get("created_at", ""),
            })
        return {"session_id": session_id, "messages": messages}
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))


@router.delete("/chat/sessions/{session_id}")
async def delete_session(
    session_id: str,
    neo4j: Annotated[Neo4jClient, Depends(get_neo4j_client)],
):
    """Delete a chat session and all its messages."""
    try:
        neo4j.run_query(DELETE_CHAT_SESSION, session_id=session_id)
        return {"deleted": session_id}
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))


# ---------------------------------------------------------------------------
# Chat endpoints
# ---------------------------------------------------------------------------

class ChatRequest(BaseModel):
    message: str
    history: list[dict] = []
    domain_filter: str | None = None
    session_id: str | None = None


@router.post("/chat")
async def chat(
    request: ChatRequest,
    settings: Annotated[Settings, Depends(get_settings)],
    neo4j: Annotated[Neo4jClient, Depends(get_neo4j_client)],
    llm: Annotated[LLMClient, Depends(get_llm_client)],
):
    intent = _classify_intent(request.message)
    context, sources, domain_names = await _build_rag_context(request.message, neo4j, llm)
    messages = _build_messages(request.history, request.message, context)

    reply = await llm.chat(
        messages=messages,
        system=CHAT_SYSTEM,
        max_tokens=settings.llm_max_tokens,
        temperature=0.4,
    )

    # Persist if session provided
    if request.session_id:
        title = request.message[:60]
        _store_exchange(neo4j, request.session_id, request.message, reply, sources, title)

    # Check DRL enrichment
    untrained = _get_untrained_domains(domain_names, neo4j)
    enrich_run_id = None
    if untrained:
        enrich_run_id = await _trigger_drl_enrichment(untrained)
        log.info(f"DRL enrichment triggered for: {untrained} (run {enrich_run_id})")

    return {
        "reply": reply,
        "sources": sources,
        "suggested_action": intent if intent != "general_qa" else None,
        "enrich_triggered": bool(untrained),
        "enrich_domains": untrained,
    }


@router.get("/chat/stream")
async def chat_stream(
    message: str = Query(...),
    history: str = Query(default="[]"),
    session_id: str = Query(default=""),
    settings: Annotated[Settings, Depends(get_settings)] = None,
    neo4j: Annotated[Neo4jClient, Depends(get_neo4j_client)] = None,
    llm: Annotated[LLMClient, Depends(get_llm_client)] = None,
):
    history_list: list[dict] = []
    try:
        history_list = json.loads(history)
    except Exception:
        pass

    context, sources, domain_names = await _build_rag_context(message, neo4j, llm)
    messages = _build_messages(history_list, message, context)

    # Check untrained domains before streaming
    untrained = _get_untrained_domains(domain_names, neo4j)

    async def generate():
        full_chunks: list[str] = []
        try:
            async for chunk in llm.chat_stream(
                messages=messages,
                system=CHAT_SYSTEM,
                max_tokens=settings.llm_max_tokens,
            ):
                full_chunks.append(chunk)
                yield f"data: {json.dumps({'text': chunk})}\n\n"
        except Exception as exc:
            log.error(f"Stream error: {exc}")
            yield f"data: {json.dumps({'text': f' [Stream error: {exc}]'})}\n\n"
        finally:
            full_response = "".join(full_chunks)

            # Persist exchange
            if session_id:
                title = message[:60]
                _store_exchange(neo4j, session_id, message, full_response, sources, title)

            # Trigger DRL enrichment in background
            enrich_run_id = None
            if untrained:
                try:
                    enrich_run_id = await _trigger_drl_enrichment(untrained)
                    log.info(f"DRL enrichment triggered for {untrained} (run {enrich_run_id})")
                except Exception as exc:
                    log.warning(f"DRL trigger failed: {exc}")

            yield f"data: {json.dumps({'sources': sources, 'enrich_triggered': bool(untrained), 'enrich_domains': untrained, 'enrich_run_id': enrich_run_id, 'done': True})}\n\n"

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