atc-data-quality/rag-api/main.py

"""RAG Knowledge API — LangChain + ChromaDB + Docling + LLM with persistent document registry."""

from __future__ import annotations

import hashlib
import json
import os
import re
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import httpx
from fastapi import FastAPI, File, Form, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langchain_text_splitters import RecursiveCharacterTextSplitter
import chromadb
from pydantic import BaseModel

CHROMA_HOST = os.getenv("CHROMA_HOST", "chromadb")
CHROMA_PORT = int(os.getenv("CHROMA_PORT", "8000"))
CHROMA_URL = f"http://{CHROMA_HOST}:{CHROMA_PORT}"
DOCLING_URL = os.getenv("DOCLING_URL", "http://docling-serve:5001").rstrip("/")
LLM_URL = os.getenv("LLM_URL", "http://10.0.20.106:8001/v1")
LLM_MODEL = os.getenv("LLM_MODEL", "gpt-4o")
LLM_API_KEY = os.getenv("LLM_API_KEY", "sk-local")
DATA_DIR = Path(os.getenv("RAG_DATA_DIR", "/data"))
UPLOADS_DIR = DATA_DIR / "uploads"
REGISTRY_PATH = DATA_DIR / "document_registry.json"
EMBED_MODEL = os.getenv("EMBED_MODEL", "sentence-transformers/all-MiniLM-L6-v2")

# Governance / agent integration
OPENMETADATA_URL = os.getenv("OPENMETADATA_URL", "http://10.0.21.47:8585").rstrip("/")
OPENMETADATA_TOKEN = os.getenv("OPENMETADATA_TOKEN", "")
COMMAND_CENTER_URL = os.getenv("COMMAND_CENTER_URL", "http://api:3201").rstrip("/")
CATALOG_COLLECTION = os.getenv("CATALOG_COLLECTION", "catalog")
MASK_TOKEN_HINT = "🔒 MASKED (masking policy ON)"

# Use placeholder images in markdown — embedded base64 destroys RAG quality.
DOCLING_IMAGE_MODE = os.getenv("DOCLING_IMAGE_MODE", "placeholder")

app = FastAPI(title="ATC RAG Knowledge API", version="1.2.0")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=120)
_embeddings: HuggingFaceEmbeddings | None = None

_BASE64_BLOB = re.compile(r"[A-Za-z0-9+/]{120,}={0,2}")
_BASE64_IMG = re.compile(r"!\[[^\]]*\]\(data:image/[^)]+\)", re.IGNORECASE)
_IMAGE_REF = re.compile(r"!\[Image\]\([^)]+\)")


def _ensure_dirs() -> None:
    UPLOADS_DIR.mkdir(parents=True, exist_ok=True)
    if not REGISTRY_PATH.exists():
        REGISTRY_PATH.write_text(json.dumps({"documents": []}, indent=2))


def _load_registry() -> list[dict[str, Any]]:
    _ensure_dirs()
    try:
        data = json.loads(REGISTRY_PATH.read_text())
        return data.get("documents", [])
    except Exception:
        return []


def _save_registry(docs: list[dict[str, Any]]) -> None:
    _ensure_dirs()
    REGISTRY_PATH.write_text(
        json.dumps({"documents": docs, "updated_at": datetime.now(timezone.utc).isoformat()}, indent=2, default=str)
    )


def _file_hash(content: bytes) -> str:
    return hashlib.sha256(content).hexdigest()


def _safe_collection(name: str) -> str:
    return "".join(c if c.isalnum() or c in "-_" else "_" for c in name.strip())[:64] or "default"


def clean_text_for_rag(text: str) -> str:
    """Strip embedded images and base64 blobs that pollute vector search."""
    text = _BASE64_IMG.sub("<!-- image -->", text)
    text = _IMAGE_REF.sub("<!-- image -->", text)
    text = _BASE64_BLOB.sub("", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def is_garbage_chunk(text: str) -> bool:
    """Detect chunks that are mostly binary/base64 noise."""
    if not text or len(text) < 20:
        return True
    if "data:image" in text:
        return True
    if _BASE64_BLOB.search(text):
        return True
    alpha = sum(1 for c in text if c.isalpha() or c.isspace())
    if alpha / max(len(text), 1) < 0.35:
        return True
    return False


def get_embeddings() -> HuggingFaceEmbeddings:
    global _embeddings
    if _embeddings is None:
        _embeddings = HuggingFaceEmbeddings(model_name=EMBED_MODEL)
    return _embeddings


def get_chroma_client() -> chromadb.HttpClient:
    return chromadb.HttpClient(host=CHROMA_HOST, port=CHROMA_PORT)


def get_vectorstore(collection: str) -> Chroma:
    return Chroma(
        client=get_chroma_client(),
        collection_name=collection,
        embedding_function=get_embeddings(),
    )


def get_llm(temperature: float = 0.2) -> ChatOpenAI:
    return ChatOpenAI(
        base_url=LLM_URL,
        api_key=LLM_API_KEY,
        model=LLM_MODEL,
        temperature=temperature,
    )


def _delete_doc_vectors(collection: str, doc_id: str) -> None:
    try:
        col = get_chroma_client().get_collection(collection)
        col.delete(where={"doc_id": doc_id})
    except Exception:
        pass


async def extract_text(content: bytes, filename: str) -> str:
    ext = Path(filename).suffix.lower()
    if ext in {".txt", ".md", ".csv", ".json"}:
        try:
            raw = content.decode("utf-8")
        except UnicodeDecodeError:
            raw = content.decode("latin-1", errors="replace")
        return clean_text_for_rag(raw)

    async with httpx.AsyncClient(timeout=300.0) as client:
        r = await client.post(
            f"{DOCLING_URL}/v1/convert/file",
            files={"files": (filename, content, "application/octet-stream")},
            data={
                "to_formats": ["md"],
                "image_export_mode": DOCLING_IMAGE_MODE,
                "do_ocr": "true",
                "table_mode": "accurate",
            },
        )
        if r.status_code >= 400:
            raise ValueError(f"Docling failed: {r.text[:300]}")
        doc = r.json().get("document") or {}
        md = doc.get("md_content") or doc.get("text_content") or ""
        md = clean_text_for_rag(md)
        if len(md) < 50:
            raise ValueError("No readable text extracted from document")
        return md


def _find_duplicate(content_hash: str, collection: str) -> dict[str, Any] | None:
    for d in _load_registry():
        if d.get("content_hash") == content_hash and d.get("collection") == collection:
            return d
    return None


def _find_by_id(doc_id: str) -> dict[str, Any] | None:
    for d in _load_registry():
        if d.get("id") == doc_id:
            return d
    return None


async def _ingest_bytes(
    content: bytes,
    filename: str,
    collection: str,
    source: str = "upload",
    *,
    force_reindex: bool = False,
) -> dict[str, Any]:
    col = _safe_collection(collection)
    content_hash = _file_hash(content)
    existing = _find_duplicate(content_hash, col)

    if existing and not force_reindex:
        return {
            "ok": True,
            "duplicate": True,
            "skipped": True,
            "message": f"Document already indexed as '{existing['filename']}' — chat immediately, no re-upload needed.",
            **{k: existing[k] for k in ("id", "filename", "collection", "chunks", "characters", "ingested_at") if k in existing},
        }

    if existing and force_reindex:
        doc_id = existing["id"]
        _delete_doc_vectors(col, doc_id)
        stored_path = Path(existing.get("stored_path", ""))
        if stored_path.exists():
            stored_path.write_bytes(content)
        else:
            stored_name = f"{doc_id}_{Path(filename).name}"
            stored_path = UPLOADS_DIR / stored_name
            stored_path.write_bytes(content)
    else:
        doc_id = uuid.uuid4().hex[:12]
        stored_name = f"{doc_id}_{Path(filename).name}"
        stored_path = UPLOADS_DIR / stored_name
        stored_path.write_bytes(content)

    text = await extract_text(content, filename)
    chunks = [c for c in _splitter.split_text(text) if not is_garbage_chunk(c)]
    if not chunks:
        raise ValueError("No usable text chunks after cleaning — document may be image-only")

    docs = [
        Document(
            page_content=chunk,
            metadata={
                "source": filename,
                "doc_id": doc_id,
                "chunk": i,
                "content_hash": content_hash,
                "ingested_at": datetime.now(timezone.utc).isoformat(),
            },
        )
        for i, chunk in enumerate(chunks)
    ]
    vs = get_vectorstore(col)
    vs.add_documents(docs)

    record = {
        "id": doc_id,
        "filename": filename,
        "collection": col,
        "content_hash": content_hash,
        "stored_path": str(stored_path),
        "chunks": len(chunks),
        "characters": len(text),
        "bytes": len(content),
        "source": source,
        "ingested_at": datetime.now(timezone.utc).isoformat(),
    }
    registry = _load_registry()
    registry = [d for d in registry if not (d.get("id") == doc_id and d.get("collection") == col)]
    registry.insert(0, record)
    _save_registry(registry)
    return {"ok": True, "duplicate": False, "reindexed": force_reindex, **record}


async def _summarize_text(text: str, filename: str) -> str:
    """Summarize document text using LLM with map-reduce for long docs."""
    llm = get_llm(temperature=0.1)
    max_chunk = 12000
    if len(text) <= max_chunk:
        prompt = (
            f"Summarize this document ({filename}) clearly in English. "
            "Include: main topic, key sections, important technologies/products mentioned, and target audience. "
            "Use bullet points and short paragraphs.\n\nDocument:\n{text}"
        )
        resp = llm.invoke([HumanMessage(content=prompt.format(text=text[:max_chunk]))])
        return resp.content if hasattr(resp, "content") else str(resp)

    # Map-reduce for long documents
    parts = [text[i : i + max_chunk] for i in range(0, min(len(text), 60000), max_chunk)]
    partials: list[str] = []
    for i, part in enumerate(parts[:5]):
        resp = llm.invoke([
            HumanMessage(content=(
                f"Summarize part {i + 1}/{min(len(parts), 5)} of '{filename}'. "
                f"List key topics, products, and technical points:\n\n{part}"
            ))
        ])
        partials.append(resp.content if hasattr(resp, "content") else str(resp))

    combined = "\n\n".join(partials)
    final = llm.invoke([
        HumanMessage(content=(
            f"Create a clear executive summary of '{filename}' from these section summaries. "
            "Structure: Overview, Main Topics, Key Technologies, Audience. Use bullet points.\n\n"
            f"{combined}"
        ))
    ])
    return final.content if hasattr(final, "content") else str(final)


class ChatRequest(BaseModel):
    message: str
    collection: str = "default"
    top_k: int = 5


class IngestTextRequest(BaseModel):
    text: str
    collection: str = "default"
    source: str = "manual"


class SummarizeRequest(BaseModel):
    collection: str = "default"
    doc_id: str | None = None
    filename: str | None = None


@app.get("/health")
async def health():
    chroma_ok = docling_ok = llm_ok = False
    doc_count = len(_load_registry())
    try:
        async with httpx.AsyncClient(timeout=5.0) as c:
            cr = await c.get(f"{CHROMA_URL}/api/v1/heartbeat")
            chroma_ok = cr.status_code < 400
            dr = await c.get(f"{DOCLING_URL}/health")
            docling_ok = dr.status_code < 400 and dr.json().get("status") == "ok"
            lr = await c.get(f"{LLM_URL.rstrip('/')}/models")
            llm_ok = lr.status_code < 400
    except Exception:
        pass
    return {
        "ok": chroma_ok,
        "chroma": chroma_ok,
        "docling": docling_ok,
        "llm": llm_ok,
        "embed_model": EMBED_MODEL,
        "stored_documents": doc_count,
        "persistent": True,
        "docling_image_mode": DOCLING_IMAGE_MODE,
    }


@app.get("/documents")
async def list_documents(collection: str | None = None):
    docs = _load_registry()
    if collection:
        col = _safe_collection(collection)
        docs = [d for d in docs if d.get("collection") == col]
    return {"documents": docs, "total": len(docs)}


@app.get("/documents/{doc_id}")
async def get_document(doc_id: str):
    doc = _find_by_id(doc_id)
    if doc:
        return doc
    return JSONResponse({"error": "not found"}, status_code=404)


@app.get("/documents/{doc_id}/file")
async def download_document(doc_id: str):
    doc = _find_by_id(doc_id)
    if doc:
        path = Path(doc.get("stored_path", ""))
        if path.exists():
            return FileResponse(path, filename=doc.get("filename", path.name))
    return JSONResponse({"error": "not found"}, status_code=404)


@app.post("/documents/{doc_id}/reindex")
async def reindex_document(doc_id: str):
    doc = _find_by_id(doc_id)
    if not doc:
        return JSONResponse({"error": "not found"}, status_code=404)
    path = Path(doc.get("stored_path", ""))
    if not path.exists():
        return JSONResponse({"error": "stored file missing"}, status_code=404)
    try:
        content = path.read_bytes()
        result = await _ingest_bytes(
            content,
            doc["filename"],
            doc["collection"],
            source=doc.get("source", "reindex"),
            force_reindex=True,
        )
        return result
    except ValueError as exc:
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=422)


@app.get("/collections")
async def list_collections():
    try:
        client = get_chroma_client()
        cols = client.list_collections()
        registry = _load_registry()
        items = []
        for col in cols:
            files = {d["filename"] for d in registry if d.get("collection") == col.name}
            items.append({
                "name": col.name,
                "documents": col.count(),
                "files": len(files),
                "filenames": sorted(files)[:20],
            })
        return {"collections": items}
    except Exception as exc:
        return JSONResponse({"error": str(exc), "collections": []}, status_code=502)


@app.post("/collections")
async def create_collection(name: str = Form(...)):
    safe = _safe_collection(name)
    get_vectorstore(safe)
    return {"ok": True, "collection": safe}


@app.post("/ingest")
async def ingest_file(
    file: UploadFile = File(...),
    collection: str = Form("default"),
    force_reindex: bool = Form(False),
):
    _ensure_dirs()
    safe_name = file.filename or "upload.txt"
    content = await file.read()
    try:
        result = await _ingest_bytes(content, safe_name, collection, force_reindex=force_reindex)
        if not result.get("ok"):
            return JSONResponse(result, status_code=422)
        return result
    except ValueError as exc:
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=422)


@app.post("/ingest/text")
async def ingest_text(body: IngestTextRequest):
    col = _safe_collection(body.collection)
    content = body.text.encode("utf-8")
    filename = f"{body.source}.txt"
    return await _ingest_bytes(content, filename, col, source=body.source)


@app.post("/summarize")
async def summarize(body: SummarizeRequest):
    col = _safe_collection(body.collection)
    registry = _load_registry()
    doc: dict[str, Any] | None = None
    if body.doc_id:
        doc = _find_by_id(body.doc_id)
    elif body.filename:
        for d in registry:
            if d.get("filename") == body.filename and d.get("collection") == col:
                doc = d
                break
    else:
        docs_in_col = [d for d in registry if d.get("collection") == col]
        if len(docs_in_col) == 1:
            doc = docs_in_col[0]

    if not doc:
        return JSONResponse({"ok": False, "error": "Document not found — specify doc_id or filename"}, status_code=404)

    path = Path(doc.get("stored_path", ""))
    if not path.exists():
        return JSONResponse({"ok": False, "error": "Stored file missing"}, status_code=404)

    try:
        content = path.read_bytes()
        text = await extract_text(content, doc["filename"])
        summary = await _summarize_text(text, doc["filename"])
        return {
            "ok": True,
            "summary": summary,
            "filename": doc["filename"],
            "doc_id": doc["id"],
            "collection": col,
            "characters": len(text),
        }
    except ValueError as exc:
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=422)
    except Exception as exc:
        return JSONResponse({"ok": False, "error": f"Summarize failed: {exc}"}, status_code=502)


@app.post("/chat")
async def chat(body: ChatRequest):
    col = _safe_collection(body.collection)
    try:
        vs = get_vectorstore(col)
        count = get_chroma_client().get_collection(col).count()
    except Exception as exc:
        return JSONResponse({"ok": False, "error": f"Collection unavailable: {exc}"}, status_code=404)

    if count == 0:
        stored = [d for d in _load_registry() if d.get("collection") == col]
        if stored:
            return JSONResponse({
                "ok": False,
                "error": "Vectors missing but files exist — click Re-index on the document in the library.",
                "stored_documents": len(stored),
            }, status_code=400)
        return JSONResponse({"ok": False, "error": "Collection is empty — upload documents first."}, status_code=400)

    retriever = vs.as_retriever(search_kwargs={"k": min(body.top_k * 3, 20)})
    raw_docs = retriever.invoke(body.message)
    docs = [d for d in raw_docs if not is_garbage_chunk(d.page_content)][: body.top_k]

    if not docs:
        return JSONResponse({
            "ok": False,
            "error": "Retrieved chunks are corrupted (old base64 index). Click Re-index on the document.",
        }, status_code=400)

    context = "\n\n---\n\n".join(
        f"[Source: {d.metadata.get('source', '?')} | chunk {d.metadata.get('chunk', '?')}]\n{d.page_content}"
        for d in docs
    )

    system = (
        "You are a helpful data assistant for the Dell ATC platform. "
        "Answer ONLY based on the provided context. If the context does not contain the answer, say so clearly. "
        "Cite sources by filename when relevant. Be concise and technical."
    )
    user = f"Context:\n{context}\n\nQuestion: {body.message}"

    try:
        llm = get_llm()
        resp = llm.invoke([SystemMessage(content=system), HumanMessage(content=user)])
        answer = resp.content if hasattr(resp, "content") else str(resp)
    except Exception as exc:
        return JSONResponse({"ok": False, "error": f"LLM error: {exc}"}, status_code=502)

    sources = [
        {"source": d.metadata.get("source"), "chunk": d.metadata.get("chunk"), "preview": d.page_content[:200]}
        for d in docs
    ]
    return {"ok": True, "answer": answer, "sources": sources, "collection": col, "context_chunks": len(docs)}


# ─────────────────────────────────────────────────────────────────────────────
# Data-catalog sync — pull OpenMetadata + Command Center governance into ChromaDB
# so the assistant can answer questions about tables, columns, PII and lineage.
# ─────────────────────────────────────────────────────────────────────────────

async def _om_get(client: httpx.AsyncClient, path: str, params: dict | None = None) -> dict[str, Any]:
    headers = {"Authorization": f"Bearer {OPENMETADATA_TOKEN}"} if OPENMETADATA_TOKEN else {}
    r = await client.get(f"{OPENMETADATA_URL}{path}", headers=headers, params=params)
    r.raise_for_status()
    return r.json()


async def _cc_get(client: httpx.AsyncClient, path: str) -> dict[str, Any]:
    r = await client.get(f"{COMMAND_CENTER_URL}{path}")
    r.raise_for_status()
    return r.json()


async def _cc_post(client: httpx.AsyncClient, path: str, body: dict) -> dict[str, Any]:
    r = await client.post(f"{COMMAND_CENTER_URL}{path}", json=body)
    r.raise_for_status()
    return r.json()


def _table_to_doc(t: dict[str, Any]) -> tuple[str, str]:
    """Render an OpenMetadata table entity as a catalog document."""
    fqn = t.get("fullyQualifiedName") or t.get("name", "?")
    lines = [f"# Table: {fqn}"]
    if t.get("description"):
        lines.append(t["description"].strip())
    cols = t.get("columns") or []
    if cols:
        lines.append("\nColumns:")
        for c in cols:
            tags = [tg.get("tagFQN", "") for tg in (c.get("tags") or [])]
            pii = " [PII]" if any("PII" in tg for tg in tags) else ""
            tag_str = f" — tags: {', '.join(t for t in tags if t)}" if tags else ""
            lines.append(f"  - {c.get('name')} ({c.get('dataType', '?')}){pii}{tag_str}")
    return fqn, "\n".join(lines)


async def _sync_catalog() -> dict[str, Any]:
    col = _safe_collection(CATALOG_COLLECTION)
    docs: list[Document] = []
    counts = {"tables": 0, "pii_datasets": 0, "flows": 0, "movements": 0}
    now = datetime.now(timezone.utc).isoformat()

    async with httpx.AsyncClient(timeout=60.0, verify=False) as client:
        # 1) OpenMetadata tables (name, columns, types, PII/classification tags)
        try:
            after: str | None = None
            seen = 0
            while seen < 2000:
                params = {"fields": "columns,tags,description", "limit": 100}
                if after:
                    params["after"] = after
                page = await _om_get(client, "/api/v1/tables", params)
                for t in page.get("data", []):
                    fqn, text = _table_to_doc(t)
                    docs.append(Document(page_content=text, metadata={
                        "source": f"openmetadata:{fqn}", "kind": "table", "fqn": fqn,
                        "doc_id": f"om_{_safe_collection(fqn)}", "ingested_at": now,
                    }))
                    seen += 1
                counts["tables"] = seen
                after = (page.get("paging") or {}).get("after")
                if not after:
                    break
        except Exception as exc:  # noqa: BLE001
            counts["openmetadata_error"] = str(exc)[:200]

        # 2) Command Center PII catalog (masked/unmasked, source per dataset)
        try:
            pii = await _cc_get(client, "/api/pii")
            for ds in pii.get("datasets", []):
                cols = ds.get("pii_columns", [])
                key = ds.get("key") or ds.get("node_id") or "?"
                lines = [f"# PII dataset: {ds.get('label') or key} (table: {ds.get('table', '?')})"]
                for c in cols:
                    status = "masked" if c.get("masked") else "unmasked"
                    via = c.get("source", "heuristic")
                    lines.append(f"  - {c.get('name')}: {c.get('category', 'PII')} [{status}, via {via}]")
                docs.append(Document(page_content="\n".join(lines), metadata={
                    "source": f"pii:{key}", "kind": "pii", "ingested_at": now,
                    "doc_id": f"pii_{_safe_collection(str(key))}",
                }))
            counts["pii_datasets"] = len(pii.get("datasets", []))
        except Exception as exc:  # noqa: BLE001
            counts["pii_error"] = str(exc)[:200]

        # 3) Data-flow graph → lineage / movement edges
        try:
            flow = await _cc_get(client, "/api/dataflow")
            nodes = {n["id"]: n for n in flow.get("nodes", [])}
            edge_lines = ["# Data flow / lineage (Command Center)"]
            for e in flow.get("edges", []):
                src = nodes.get(e.get("from"), {}).get("label", e.get("from"))
                dst = nodes.get(e.get("to"), {}).get("label", e.get("to"))
                mid = f", movement_id={e['movement_id']}" if e.get("movement_id") else ""
                edge_lines.append(f"  - {src} → {dst} ({e.get('kind', '')}{mid})")
            docs.append(Document(page_content="\n".join(edge_lines), metadata={
                "source": "dataflow:lineage", "kind": "lineage", "ingested_at": now,
                "doc_id": "cc_lineage",
            }))
            counts["flows"] = len(flow.get("edges", []))
        except Exception as exc:  # noqa: BLE001
            counts["dataflow_error"] = str(exc)[:200]

        # 4) Movements / pipelines
        try:
            mv = await _cc_get(client, "/api/movements")
            mv_lines = ["# Data movements / pipelines (Command Center)"]
            for m in mv.get("movements", []):
                mv_lines.append(
                    f"  - id={m.get('id')} · {m.get('label')} ({m.get('kind')}): "
                    f"{m.get('from', '?')} → {m.get('to', '?')} · agent={m.get('agent', '—')}"
                )
            docs.append(Document(page_content="\n".join(mv_lines), metadata={
                "source": "movements", "kind": "movements", "ingested_at": now,
                "doc_id": "cc_movements",
            }))
            counts["movements"] = len(mv.get("movements", []))
        except Exception as exc:  # noqa: BLE001
            counts["movements_error"] = str(exc)[:200]

    if not docs:
        return {"ok": False, "error": "No catalog data could be fetched", "counts": counts}

    # Rebuild the catalog collection cleanly each sync.
    try:
        get_chroma_client().delete_collection(col)
    except Exception:
        pass
    get_vectorstore(col).add_documents(docs)
    return {"ok": True, "collection": col, "documents": len(docs), "counts": counts, "synced_at": now}


@app.post("/catalog/sync")
async def catalog_sync():
    try:
        return await _sync_catalog()
    except Exception as exc:  # noqa: BLE001
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=502)


@app.get("/catalog/status")
async def catalog_status():
    col = _safe_collection(CATALOG_COLLECTION)
    try:
        count = get_chroma_client().get_collection(col).count()
        return {"ok": True, "collection": col, "documents": count}
    except Exception:
        return {"ok": True, "collection": col, "documents": 0}


# ─────────────────────────────────────────────────────────────────────────────
# Agentic chat — a model-agnostic JSON-action tool loop over the live platform,
# with an approval-gated executor for write actions and SSE token streaming.
# ─────────────────────────────────────────────────────────────────────────────

def _catalog_search(query: str, k: int = 6) -> str:
    col = _safe_collection(CATALOG_COLLECTION)
    try:
        vs = get_vectorstore(col)
        if get_chroma_client().get_collection(col).count() == 0:
            return "Catalog is empty — run /catalog/sync first."
        hits = vs.as_retriever(search_kwargs={"k": k}).invoke(query)
        return "\n\n---\n\n".join(h.page_content for h in hits) or "No matches."
    except Exception as exc:  # noqa: BLE001
        return f"catalog search failed: {exc}"


async def _tool_search_catalog(client: httpx.AsyncClient, args: dict) -> str:
    return _catalog_search(str(args.get("query", "")))


async def _tool_get_pii(client: httpx.AsyncClient, args: dict) -> str:
    return json.dumps((await _cc_get(client, "/api/pii")).get("datasets", []))[:3000]


async def _tool_get_cdc(client: httpx.AsyncClient, args: dict) -> str:
    return json.dumps(await _cc_get(client, "/api/changes/stats"))[:2000]


async def _tool_list_movements(client: httpx.AsyncClient, args: dict) -> str:
    return json.dumps((await _cc_get(client, "/api/movements")).get("movements", []))[:3000]


async def _tool_lookup_records(client: httpx.AsyncClient, args: dict) -> str:
    """Look up actual records in a source dataset. Masked columns come back redacted —
    the real values are enforced server-side and never reach the model."""
    key = str(args.get("key") or args.get("dataset") or "").strip()
    if not key:
        return "error: 'key' is required (one of: postgres, mysql, mongodb, curated)"
    body = {"key": key, "search": args.get("search"), "limit": int(args.get("limit", 5) or 5)}
    try:
        res = await _cc_post(client, "/api/pii/lookup", body)
    except Exception as exc:  # noqa: BLE001
        return f"lookup failed: {exc}"
    return json.dumps(res)[:3500]


async def _tool_propose_movement(client: httpx.AsyncClient, args: dict) -> str:
    """WRITE action — never executes directly; files an approval request (Mo & Bart gate)."""
    mid = str(args.get("movement_id", "")).strip()
    reason = str(args.get("reason", "Requested via assistant"))[:300]
    if not mid:
        return "error: movement_id is required"
    body = {
        "agent_id": "etl-guardian",
        "action": f"Trigger data movement '{mid}'",
        "reason": reason,
        "action_type": "etl.restart",
        "target": mid,
        "payload": {"executor": "movement", "movement_id": mid},
        "priority": "normal",
    }
    res = await _cc_post(client, "/api/approvals", body)
    appr = res.get("approval", {})
    return (
        f"Approval filed (id={appr.get('id', '?')}, status={appr.get('status', 'pending')}). "
        "The movement runs automatically once Mo or Bart approves it in the Approvals queue."
    )


_TOOLS: dict[str, Any] = {
    "search_catalog": _tool_search_catalog,
    "get_pii": _tool_get_pii,
    "get_cdc": _tool_get_cdc,
    "list_movements": _tool_list_movements,
    "lookup_records": _tool_lookup_records,
    "propose_movement": _tool_propose_movement,
}

_AGENT_SYSTEM = (
    "You are the Dell ATC data-platform assistant. You can inspect a live lakehouse "
    "(PostgreSQL, MySQL, MongoDB sources → Debezium CDC → Kafka → Iceberg/Trino, governed "
    "by OpenMetadata) and propose data movements.\n\n"
    "You work in a strict loop. On every turn reply with EXACTLY ONE single-line JSON object, nothing else:\n"
    '  To use a tool:  {"action": "<tool>", "input": {<args>}}\n'
    '  To finish:      {"final": true}\n\n'
    "Available tools:\n"
    '  - search_catalog{"query": str}  → tables, columns, PII tags, lineage from the catalog\n'
    '  - get_pii{}                     → live PII columns per dataset (masked/unmasked status)\n'
    '  - get_cdc{}                     → live CDC change volume per source\n'
    '  - list_movements{}              → available pipelines/movements with their ids\n'
    '  - lookup_records{"key": str, "search": str}  → actual records in a source dataset\n'
    "        (key is one of: postgres, mysql, mongodb, curated; search matches the name column)\n"
    '  - propose_movement{"movement_id": str, "reason": str}  → WRITE: files an approval (human-gated)\n\n'
    "DATA MASKING POLICY — STRICT:\n"
    "Some columns are masked by the operator's policy. In any tool result, a masked value appears "
    f'literally as "{MASK_TOKEN_HINT}". You must NEVER reveal, guess, infer, reconstruct, or work '
    "around a masked value. If the user asks for a field that comes back masked, reply that the value "
    "is withheld for privacy/security because masking is enabled for that column, and tell them an "
    "operator can disable masking per column in the Command Center: Data Flow → click the source node "
    "→ PII panel.\n"
    "If a requested value is returned in CLEAR TEXT (not the mask token), you MUST state that exact "
    "value verbatim in your answer — do not withhold or paraphrase it; masking is off for that column.\n\n"
    "Do NOT call propose_movement unless the user explicitly asks to run, trigger or mask a pipeline. "
    "Answering a question about data is NOT a reason to propose a movement.\n\n"
    'When you have enough information, reply {"final": true} and you will then be asked to write the answer.'
)

_JSON_RE = re.compile(r"\{.*\}", re.DOTALL)


def _parse_action(text: str) -> dict[str, Any] | None:
    m = _JSON_RE.search(text or "")
    if not m:
        return None
    try:
        obj = json.loads(m.group(0))
        return obj if isinstance(obj, dict) else None
    except Exception:
        return None


class AgentRequest(BaseModel):
    message: str
    max_steps: int = 5


def _sse(event: str, data: dict) -> str:
    return f"event: {event}\ndata: {json.dumps(data, default=str)}\n\n"


async def _agent_stream(message: str, max_steps: int):
    llm = get_llm(temperature=0.1)
    transcript: list[Any] = [SystemMessage(content=_AGENT_SYSTEM), HumanMessage(content=message)]
    steps: list[dict[str, Any]] = []

    async with httpx.AsyncClient(timeout=60.0, verify=False) as client:
        for _ in range(max(1, min(max_steps, 8))):
            try:
                resp = llm.invoke(transcript + [HumanMessage(content="Next action as a single JSON object:")])
                raw = resp.content if hasattr(resp, "content") else str(resp)
            except Exception as exc:  # noqa: BLE001
                yield _sse("error", {"error": f"LLM error: {exc}"})
                return

            action = _parse_action(raw)
            if not action or action.get("final"):
                break

            tool = action.get("action")
            args = action.get("input") or {}
            fn = _TOOLS.get(tool)
            if not fn:
                transcript.append(HumanMessage(content=f"Observation: unknown tool '{tool}'."))
                continue

            yield _sse("step", {"tool": tool, "input": args})
            try:
                observation = await fn(client, args)
            except Exception as exc:  # noqa: BLE001
                observation = f"tool error: {exc}"
            steps.append({"tool": tool, "input": args, "observation": observation[:500]})
            yield _sse("observation", {"tool": tool, "observation": observation[:1200]})
            transcript.append(AIMessage(content=raw))
            transcript.append(HumanMessage(content=f"Observation from {tool}: {observation[:4000]}"))

    # Final answer — stream tokens to the client.
    final_prompt = transcript + [HumanMessage(content=(
        "Using the observations above, write the final answer for the user. "
        "Be concise and technical, cite concrete table/column/movement names, and if you filed an "
        "approval say so explicitly. Plain prose only — no JSON."
    ))]
    yield _sse("answer_start", {"steps": len(steps)})
    try:
        for chunk in llm.stream(final_prompt):
            tok = chunk.content if hasattr(chunk, "content") else str(chunk)
            if tok:
                yield _sse("token", {"t": tok})
    except Exception as exc:  # noqa: BLE001
        yield _sse("error", {"error": f"LLM stream error: {exc}"})
        return
    yield _sse("done", {"steps": steps})


@app.post("/agent")
async def agent(body: AgentRequest):
    return StreamingResponse(
        _agent_stream(body.message, body.max_steps),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )


@app.get("/")
async def root():
    return {
        "service": "ATC RAG Knowledge API",
        "persistent_storage": "ChromaDB + document registry on disk",
        "endpoints": ["/health", "/documents", "/collections", "/ingest", "/chat", "/summarize",
                      "/catalog/sync", "/catalog/status", "/agent", "/docs"],
    }