"""Document text extraction — runs on the StartOS box (CPU only).

Confidential documents are dropped into /data/inbox. Before anything is shipped
to the Sparks, we extract plain text here so that only normalized text (never the
original binaries) crosses to the review containers. Supported formats:

    .pdf            -> pypdf
    .docx           -> python-docx
    .txt .md .text  -> read as UTF-8

Anything else is skipped with a note. Each extracted document becomes a single
UTF-8 .txt file in the per-job staging directory, which is rsynced to the Spark
and mounted read-only into every reviewer container at /docs.
"""
from __future__ import annotations

import os

TEXT_EXTS = {".txt", ".md", ".text", ".markdown"}
SUPPORTED = TEXT_EXTS | {".pdf", ".docx"}


def _extract_pdf(path: str) -> str:
    from pypdf import PdfReader

    reader = PdfReader(path)
    parts = []
    for i, page in enumerate(reader.pages, 1):
        try:
            txt = page.extract_text() or ""
        except Exception as e:
            txt = f"(page {i}: extraction error: {e})"
        parts.append(f"\n\n===== page {i} =====\n{txt}")
    return "".join(parts).strip()


def _extract_docx(path: str) -> str:
    import docx

    doc = docx.Document(path)
    lines = [p.text for p in doc.paragraphs]
    # Include table cell text too — contracts/specs often hide content in tables.
    for table in doc.tables:
        for row in table.rows:
            cells = [c.text.strip() for c in row.cells]
            if any(cells):
                lines.append(" | ".join(cells))
    return "\n".join(lines).strip()


def _extract_text(path: str) -> str:
    with open(path, errors="replace") as f:
        return f.read().strip()


def extract_file(path: str) -> str:
    ext = os.path.splitext(path)[1].lower()
    if ext == ".pdf":
        return _extract_pdf(path)
    if ext == ".docx":
        return _extract_docx(path)
    if ext in TEXT_EXTS:
        return _extract_text(path)
    raise ValueError(f"unsupported file type: {ext or '(none)'}")


def _safe_name(name: str) -> str:
    base = os.path.splitext(os.path.basename(name))[0]
    keep = "".join(c if (c.isalnum() or c in "-_ ") else "_" for c in base).strip()
    return (keep or "document").replace(" ", "_")


def extract_inbox(inbox: str, out_dir: str, log=print) -> list[dict]:
    """Extract every supported file in `inbox` to a .txt in `out_dir`.

    Returns a manifest: [{source, out, chars, ok, error}]. Skips unsupported
    files (recorded with ok=False) rather than failing the whole job."""
    os.makedirs(out_dir, exist_ok=True)
    manifest: list[dict] = []
    if not os.path.isdir(inbox):
        return manifest
    seen: dict[str, int] = {}
    for fn in sorted(os.listdir(inbox)):
        src = os.path.join(inbox, fn)
        if not os.path.isfile(src):
            continue
        ext = os.path.splitext(fn)[1].lower()
        rec = {"source": fn, "out": None, "chars": 0, "ok": False, "error": ""}
        if ext not in SUPPORTED:
            rec["error"] = f"unsupported type {ext or '(none)'}"
            log(f"[extract] skip {fn}: {rec['error']}")
            manifest.append(rec)
            continue
        try:
            text = extract_file(src)
        except Exception as e:
            rec["error"] = str(e)[:300]
            log(f"[extract] FAILED {fn}: {rec['error']}")
            manifest.append(rec)
            continue
        stem = _safe_name(fn)
        if stem in seen:
            seen[stem] += 1
            stem = f"{stem}-{seen[stem]}"
        else:
            seen[stem] = 1
        out_name = f"{stem}.txt"
        out_path = os.path.join(out_dir, out_name)
        header = f"# Source document: {fn}\n\n"
        with open(out_path, "w") as f:
            f.write(header + text + "\n")
        rec.update({"out": out_name, "chars": len(text), "ok": True})
        log(f"[extract] {fn} -> {out_name} ({len(text)} chars)")
        manifest.append(rec)
    return manifest