"""FastAPI endpoints for document chunking + demo frontend."""

from __future__ import annotations

import asyncio
import json
import mimetypes
import re
import shutil
import tempfile
import time
from datetime import datetime, timezone
from pathlib import Path

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

from rag_cut.models import ChunkResult, SplitConfig, SplitMode
from rag_cut.pipeline import chunk_document
from rag_cut.retrieval import recall_chunks

ROOT = Path(__file__).resolve().parent.parent.parent
FRONTEND_DIR = ROOT / "frontend"
STORAGE_DIR = ROOT / "storage"
RESULTS_DIR = STORAGE_DIR / "results"
UPLOADS_DIR = STORAGE_DIR / "uploads"
DOC_ID_RE = re.compile(r"^[a-f0-9]{12}$")
WORD_EXTS = {".doc", ".docx"}

app = FastAPI(title="RAG-cut Chunking API", version="0.1.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)

(STORAGE_DIR / "assets").mkdir(parents=True, exist_ok=True)
RESULTS_DIR.mkdir(parents=True, exist_ok=True)
UPLOADS_DIR.mkdir(parents=True, exist_ok=True)
app.mount("/assets", StaticFiles(directory=STORAGE_DIR / "assets"), name="assets")


class RecallRequest(BaseModel):
    doc_id: str
    query: str = Field(min_length=1, max_length=500)
    top_k: int = Field(default=5, ge=1, le=20)


def _result_path(doc_id: str) -> Path:
    if not DOC_ID_RE.fullmatch(doc_id):
        raise HTTPException(status_code=400, detail="Invalid doc_id")
    return RESULTS_DIR / f"{doc_id}.json"


def _save_result(result: ChunkResult) -> None:
    _result_path(result.doc_id).write_text(result.model_dump_json(indent=2), encoding="utf-8")


def _find_original(doc_id: str) -> Path | None:
    uploads_dir = UPLOADS_DIR / doc_id
    if not uploads_dir.is_dir():
        return None
    files = sorted(p for p in uploads_dir.iterdir() if p.is_file())
    return files[0] if files else None


def _find_preview(doc_id: str) -> Path | None:
    """Return a browser-previewable source, preferring converted Word PDFs."""
    original = _find_original(doc_id)
    if original is None:
        return None
    if original.suffix.lower() == ".pdf":
        return original
    if original.suffix.lower() not in WORD_EXTS:
        return None

    conversion_dir = STORAGE_DIR / "assets" / doc_id / "_conversion"
    exact = conversion_dir / f"{original.stem}.pdf"
    if exact.is_file():
        return exact
    if conversion_dir.is_dir():
        return next(iter(sorted(conversion_dir.glob("*.pdf"))), None)
    return None


def _load_result(doc_id: str) -> ChunkResult:
    result_path = _result_path(doc_id)
    if not result_path.exists():
        raise HTTPException(status_code=404, detail="Chunk result not found")
    try:
        return ChunkResult.model_validate_json(result_path.read_text(encoding="utf-8"))
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Invalid result file: {exc}") from exc


def _result_summary(path: Path) -> dict | None:
    """Build list summary without loading the full chunks array."""
    try:
        with path.open("r", encoding="utf-8") as fh:
            head = fh.read(32768)
        idx = head.find('\n  "chunks"')
        if idx == -1:
            idx = head.find('"chunks"')
        if idx != -1:
            data = json.loads(head[:idx].rstrip().rstrip(",") + "\n}")
        else:
            data = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None
    doc_id = data.get("doc_id") or path.stem
    if not DOC_ID_RE.fullmatch(str(doc_id)):
        return None
    mtime = path.stat().st_mtime
    return {
        "doc_id": doc_id,
        "filename": data.get("filename") or path.name,
        "split_mode": data.get("split_mode"),
        "chunk_count": data.get("chunk_count", 0),
        "block_count": data.get("block_count", 0),
        "has_original": _find_original(doc_id) is not None,
        "saved_at": datetime.fromtimestamp(mtime, tz=timezone.utc).isoformat(),
        "saved_ts": mtime,
    }


@app.get("/health")
def health():
    return {"status": "ok"}


@app.get("/api/results")
def list_results():
    """List persisted chunk results (newest first)."""
    items: list[dict] = []
    for path in RESULTS_DIR.glob("*.json"):
        summary = _result_summary(path)
        if summary:
            items.append(summary)
    items.sort(key=lambda item: item["saved_ts"], reverse=True)
    for item in items:
        item.pop("saved_ts", None)
    return {"count": len(items), "results": items}


@app.get("/api/results/{doc_id}")
def get_result(doc_id: str):
    """Load a full persisted chunk result by doc_id."""
    chunk_result = _load_result(doc_id)
    payload = json.loads(chunk_result.model_dump_json())
    payload["has_original"] = _find_original(doc_id) is not None
    return JSONResponse(content=payload)


@app.get("/api/results/{doc_id}/original")
def get_original(doc_id: str):
    """Serve the original uploaded file for a historical result (if still on disk)."""
    if not DOC_ID_RE.fullmatch(doc_id):
        raise HTTPException(status_code=400, detail="Invalid doc_id")
    original = _find_original(doc_id)
    if original is None:
        raise HTTPException(status_code=404, detail="Original file not found")
    media_type, _ = mimetypes.guess_type(original.name)
    return FileResponse(
        path=original,
        media_type=media_type or "application/octet-stream",
        filename=original.name,
    )


@app.get("/api/results/{doc_id}/preview")
def get_preview(doc_id: str):
    """Serve an inline PDF preview for PDF and converted Word documents."""
    if not DOC_ID_RE.fullmatch(doc_id):
        raise HTTPException(status_code=400, detail="Invalid doc_id")
    preview = _find_preview(doc_id)
    if preview is None:
        raise HTTPException(status_code=404, detail="PDF preview not found")
    return FileResponse(path=preview, media_type="application/pdf")


def _delete_tree(path: Path) -> bool:
    if not path.exists():
        return False
    if path.is_dir():
        shutil.rmtree(path)
    else:
        path.unlink()
    return True


@app.delete("/api/results/{doc_id}")
def delete_result(doc_id: str):
    """Delete a persisted chunk result and its uploads/assets."""
    result_path = _result_path(doc_id)
    if not result_path.exists():
        raise HTTPException(status_code=404, detail="Chunk result not found")

    removed = {
        "result": _delete_tree(result_path),
        "uploads": _delete_tree(UPLOADS_DIR / doc_id),
        "assets": _delete_tree(STORAGE_DIR / "assets" / doc_id),
    }
    return {"doc_id": doc_id, "deleted": True, "removed": removed}


@app.post("/api/chunk")
async def chunk_file(
    file: UploadFile = File(...),
    mode: str | None = Form(None),
    delimiter: str | None = Form(None),
    parent_delimiter: str | None = Form(None),
    child_delimiter: str | None = Form(None),
    max_chunk_size: int | None = Form(None),
    child_max_size: int | None = Form(None),
    overlap: int | None = Form(None),
    header_row_start: int | None = Form(None),
    header_row_end: int | None = Form(None),
    start_row: int | None = Form(None),
    rows_per_chunk: int | None = Form(None),
):
    config = None
    if mode is not None:
        try:
            split_mode = SplitMode(mode)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=f"Invalid mode: {mode}") from exc

        config = SplitConfig(
            mode=split_mode,
            delimiter=delimiter,
            parent_delimiter=parent_delimiter,
            child_delimiter=child_delimiter,
            max_chunk_size=1500 if max_chunk_size is None else max_chunk_size,
            child_max_size=512 if child_max_size is None else child_max_size,
            overlap=150 if overlap is None else overlap,
            header_row_start=1 if header_row_start is None else header_row_start,
            header_row_end=1 if header_row_end is None else header_row_end,
            start_row=2 if start_row is None else start_row,
            rows_per_chunk=1 if rows_per_chunk is None else rows_per_chunk,
        )

    filename = Path(file.filename or "upload").name
    payload = await file.read()
    with tempfile.TemporaryDirectory() as tmp:
        tmp_path = Path(tmp) / filename
        tmp_path.write_bytes(payload)
        try:
            # Heavy parsing (MinerU / LibreOffice / OCR) must not block the event loop,
            # otherwise /health and other uploads hang while one document is processing.
            result = await asyncio.to_thread(chunk_document, tmp_path, config)
            _save_result(result)
            return JSONResponse(content=json.loads(result.model_dump_json()))
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=str(exc)) from exc
        except FileNotFoundError as exc:
            raise HTTPException(status_code=404, detail=str(exc)) from exc
        except RuntimeError as exc:
            raise HTTPException(status_code=500, detail=str(exc)) from exc


@app.post("/api/recall")
def recall(request: RecallRequest):
    result_path = _result_path(request.doc_id)
    if not result_path.exists():
        raise HTTPException(status_code=404, detail="Chunk result not found; upload and split the document again")

    started = time.perf_counter()
    chunk_result = ChunkResult.model_validate_json(result_path.read_text(encoding="utf-8"))
    results, candidate_count = recall_chunks(request.query, chunk_result.chunks, request.top_k)
    return {
        "doc_id": request.doc_id,
        "query": request.query,
        "top_k": request.top_k,
        "candidate_count": candidate_count,
        "returned_count": len(results),
        "elapsed_ms": round((time.perf_counter() - started) * 1000, 3),
        "results": results,
    }


if FRONTEND_DIR.exists():
    app.mount("/", StaticFiles(directory=FRONTEND_DIR, html=True), name="frontend")