"""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
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
from langchain_core.messages import 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")

# 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)}


@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", "/docs"],
    }