#!/usr/bin/env python3
"""Batch reprocess all health documents into the JARVIS health DB.

Design:
- Never deletes documents.
- Makes a DB backup before writing.
- Resolves stale paths between /home/agent/Gesundheit and ~/.hermes/assets/Gesundheit.
- Extracts searchable full text from PDFs via PyMuPDF, with optional Docling fallback for weak extraction.
- Extracts XLSX sheet text via openpyxl.
- Stores full text in dokumente.extrahierte_inhalte and lightweight insights in document_insights.
- Marks duplicate candidates conservatively via review_status/processing_quality only.
- Does NOT auto-promote OCR/PDF lab values to confirmed laborwerte. Those belong in staging/4-eyes review.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shutil
import sqlite3
import sys
import time
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from typing import Iterable

BASE = Path.home() / ".hermes" / "assets" / "Gesundheit"
DB = BASE / "health_data.db"
REPORTS = BASE / "reports"
BACKUPS = BASE / "backups"
EXTRACT_DIR = BASE / "extracted_text"
SEARCH_ROOTS = [
    BASE,
    Path.home() / "Gesundheit",
    Path("/home/agent/Gesundheit"),
]

DATE_RE = re.compile(r"(?:(20\d{2})[-./](\d{1,2})[-./](\d{1,2})|(?<!\d)(\d{1,2})[.](\d{1,2})[.](20\d{2})(?!\d))")
DIAG_RE = re.compile(r"\b(Diagnose(?:n)?|Beurteilung|Befund|Anamnese|Therapie|Medikation|Labor|CRP|D-Dimer|Faktor\s*VIII|Behçet|Thrombose)\b", re.I)
LAB_PARAM_RE = re.compile(r"\b(CRP|C-Reaktives Protein|D-?Dimer|Faktor\s*VIII|Leukozyten|Thrombozyten|Hämoglobin|Haemoglobin|Ferritin|Vitamin\s*D|LDL|HDL|Cholesterin|Triglyceride|Kreatinin|ALAT|ASAT|GGT|Glukose)\b", re.I)


def connect() -> sqlite3.Connection:
    con = sqlite3.connect(DB)
    con.row_factory = sqlite3.Row
    return con


def sha256(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def path_candidates(raw: str | None, name: str) -> Iterable[Path]:
    seen: set[str] = set()
    vals = [raw, name]
    for v in vals:
        if not v:
            continue
        v = os.path.expanduser(str(v))
        candidates = [Path(v)]
        if not Path(v).is_absolute():
            candidates += [BASE / v, Path.home() / v, Path("/home/agent") / v]
        # old path variants
        candidates.append(Path(v.replace("/home/agent/Gesundheit", str(Path.home() / "Gesundheit"))))
        candidates.append(Path(v.replace("Gesundheit/Archiv", "/home/agent/Gesundheit/Archiv")))
        candidates.append(Path(v.replace("Gesundheit/archiv", "/home/agent/Gesundheit/Archiv")))
        for p in candidates:
            key = str(p)
            if key not in seen:
                seen.add(key); yield p
    # filename search fallback
    roots = [r for r in SEARCH_ROOTS if r.exists()]
    for root in roots:
        try:
            yield from root.rglob(name)
        except Exception:
            pass
    # fuzzy search fallback: date-prefix and normalized stem. Covers old archives where
    # spaces/underscores/umlauts were changed after DB insertion.
    date_prefix = re.match(r"(\d{6})", name or "")
    norm_name = re.sub(r"[^a-z0-9]+", "", Path(name).stem.lower())
    for root in roots:
        try:
            patterns = []
            if date_prefix:
                patterns += [f"{date_prefix.group(1)}*", f"{date_prefix.group(1)[:2]}{date_prefix.group(1)[2:4]}{date_prefix.group(1)[4:]}*"]
            for pat in patterns:
                for cand in root.rglob(pat):
                    if cand.is_file():
                        yield cand
            if norm_name:
                for cand in root.rglob("*"):
                    if cand.is_file() and re.sub(r"[^a-z0-9]+", "", cand.stem.lower()) == norm_name:
                        yield cand
        except Exception:
            pass


def resolve_path(row: sqlite3.Row) -> Path | None:
    for p in path_candidates(row["dateipfad"], row["datei_name"]):
        try:
            if p.exists() and p.is_file():
                return p.resolve()
        except OSError:
            continue
    return None


def extract_pdf_pymupdf(path: Path) -> tuple[str, dict]:
    import fitz
    parts = []
    meta = {"method": "pymupdf", "pages": 0}
    with fitz.open(path) as doc:
        meta["pages"] = doc.page_count
        for i, page in enumerate(doc, start=1):
            txt = page.get_text("text") or ""
            parts.append(f"\n\n--- PAGE {i} ---\n{txt.strip()}")
    return "\n".join(parts).strip(), meta


def extract_pdf_docling(path: Path, timeout_s: int = 240) -> tuple[str, dict]:
    # Import lazily; Docling is excellent but heavy. Use only if PyMuPDF text is weak.
    from multiprocessing import Process, Queue

    def worker(q: Queue, p: str) -> None:
        try:
            from docling.document_converter import DocumentConverter
            conv = DocumentConverter()
            result = conv.convert(p)
            q.put((result.document.export_to_markdown(), None))
        except Exception as e:
            q.put(("", repr(e)))

    q: Queue = Queue()
    proc = Process(target=worker, args=(q, str(path)))
    proc.start(); proc.join(timeout_s)
    if proc.is_alive():
        proc.terminate(); proc.join(10)
        return "", {"method": "docling", "error": "timeout"}
    try:
        text, err = q.get_nowait()
    except Exception:
        return "", {"method": "docling", "error": "no_result"}
    return text or "", {"method": "docling", "error": err}


def extract_xlsx(path: Path) -> tuple[str, dict]:
    import openpyxl
    wb = openpyxl.load_workbook(path, read_only=True, data_only=True)
    parts = []
    cells = 0
    for ws in wb.worksheets:
        parts.append(f"\n\n--- SHEET {ws.title} ---")
        max_r = min(ws.max_row or 0, 500)
        max_c = min(ws.max_column or 0, 80)
        for row in ws.iter_rows(min_row=1, max_row=max_r, max_col=max_c, values_only=True):
            vals = [str(v).strip() for v in row if v is not None and str(v).strip() != ""]
            if vals:
                cells += len(vals)
                parts.append("\t".join(vals))
    return "\n".join(parts).strip(), {"method": "openpyxl", "sheets": wb.sheetnames, "cells": cells}


def normalize_date(match: re.Match) -> str:
    if match.group(1):
        y, m, d = match.group(1), match.group(2), match.group(3)
    else:
        d, m, y = match.group(4), match.group(5), match.group(6)
    try:
        return f"{int(y):04d}-{int(m):02d}-{int(d):02d}"
    except Exception:
        return ""


def extract_insights(text: str) -> dict:
    dates = sorted({normalize_date(m) for m in DATE_RE.finditer(text) if normalize_date(m)})[:20]
    lab_params = sorted({m.group(0) for m in LAB_PARAM_RE.finditer(text)})[:30]
    keywords = sorted({m.group(0) for m in DIAG_RE.finditer(text)})[:30]
    snippets = []
    for m in DIAG_RE.finditer(text):
        start = max(0, m.start() - 120); end = min(len(text), m.end() + 220)
        snip = re.sub(r"\s+", " ", text[start:end]).strip()
        if snip and snip not in snippets:
            snippets.append(snip)
        if len(snippets) >= 8:
            break
    return {"dates": dates, "lab_params": lab_params, "keywords": keywords, "snippets": snippets}


def quality(text: str, meta: dict) -> str:
    n = len(text or "")
    if n >= 1500:
        return "good"
    if n >= 300:
        return "partial"
    if meta.get("method") == "docling" and meta.get("error"):
        return "failed"
    return "weak"


def upsert_insight(con: sqlite3.Connection, doc_id: int, insight: dict, q: str) -> None:
    summary = []
    if insight["dates"]:
        summary.append("Datums: " + ", ".join(insight["dates"][:8]))
    if insight["lab_params"]:
        summary.append("Laborparameter: " + ", ".join(insight["lab_params"][:12]))
    if insight["keywords"]:
        summary.append("Keywords: " + ", ".join(insight["keywords"][:12]))
    if insight["snippets"]:
        summary.append("Snippets: " + " | ".join(insight["snippets"][:3]))
    title = "Volltext-Reprocessing"
    content = "\n".join(summary) + "\n\nJSON:\n" + json.dumps(insight, ensure_ascii=False)
    datum = insight["dates"][0] if insight["dates"] else None
    severity = "info" if q in ("good", "partial") else "warning"
    con.execute(
        """
        INSERT INTO document_insights(dokument_id, insight_type, datum, titel, inhalt, severity, source, created_at)
        VALUES(?,?,?,?,?,?,?,CURRENT_TIMESTAMP)
        """,
        (doc_id, "fulltext_reprocessing", datum, title, content, severity, "process_all_health_documents.py"),
    )


def process_one(con: sqlite3.Connection, row: sqlite3.Row, force: bool = False, use_docling: bool = True) -> dict:
    doc_id = row["id"]
    name = row["datei_name"]
    p = resolve_path(row)
    result = {"id": doc_id, "name": name, "path": str(p) if p else None, "ok": False, "quality": "missing", "chars": 0, "method": None, "error": None}
    if not p:
        con.execute("UPDATE dokumente SET processing_quality=?, review_status=COALESCE(review_status,?) WHERE id=?", ("missing_file", "needs_path_resolution", doc_id))
        return result
    ext = p.suffix.lower()
    text = ""; meta = {}
    try:
        if ext == ".pdf":
            text, meta = extract_pdf_pymupdf(p)
            if use_docling and len(text) < 500:
                dtext, dmeta = extract_pdf_docling(p)
                if len(dtext) > len(text):
                    text, meta = dtext, dmeta
                else:
                    meta = {"method": "pymupdf", "docling_fallback": dmeta}
        elif ext in (".xlsx", ".xlsm"):
            text, meta = extract_xlsx(p)
        elif ext in (".jpg", ".jpeg", ".png", ".tif", ".tiff"):
            # Prefer existing sidecar OCR if available; fall back to pytesseract if installed.
            sidecars = [p.with_name(p.stem + "_ocr.txt"), p.with_suffix(".txt")]
            sidecar = next((s for s in sidecars if s.exists() and s.stat().st_size > 0), None)
            if sidecar:
                text = sidecar.read_text(encoding="utf-8", errors="ignore")
                meta = {"method": "sidecar_ocr", "sidecar": str(sidecar)}
            else:
                try:
                    import pytesseract
                    from PIL import Image
                    text = pytesseract.image_to_string(Image.open(p), lang="deu+eng")
                    meta = {"method": "tesseract"}
                except Exception as e:
                    text = ""
                    meta = {"method": "image", "error": repr(e)}
        else:
            try:
                text = p.read_text(encoding="utf-8", errors="ignore")
                meta = {"method": "text"}
            except Exception:
                text = ""; meta = {"method": "unknown"}
        q = quality(text, meta)
        h = sha256(p)
        EXTRACT_DIR.mkdir(exist_ok=True)
        (EXTRACT_DIR / f"document_{doc_id}_{re.sub(r'[^A-Za-z0-9_.-]+','_',name)}.txt").write_text(text, encoding="utf-8")
        insight = extract_insights(text)
        docdate = insight["dates"][0] if insight["dates"] else row["document_date"]
        institution = None
        for inst in ["USZ", "KSB", "Viollier", "Rothen", "Vista", "NFPraxis", "Dermatologische Klinik", "Klinische Chemie", "Immunologie"]:
            if inst.lower() in (text + " " + name).lower():
                institution = inst; break
        con.execute(
            """
            UPDATE dokumente
               SET extrahierte_inhalte=?, verarbeite_datum=CURRENT_TIMESTAMP, status='eingearbeitet',
                   processing_quality=?, review_status=CASE WHEN ? IN ('good','partial') THEN COALESCE(review_status,'auto_extracted') ELSE COALESCE(review_status,'needs_review') END,
                   datei_hash=?, local_original_path=?, dateipfad=?, document_date=COALESCE(document_date,?), institution=COALESCE(institution,?)
             WHERE id=?
            """,
            (text, q, q, h, str(p), str(p), docdate, institution, doc_id),
        )
        upsert_insight(con, doc_id, insight, q)
        result.update(ok=True, quality=q, chars=len(text), method=meta.get("method"), dates=insight["dates"][:5], lab_params=insight["lab_params"][:8])
    except Exception as e:
        result["error"] = repr(e)
        con.execute("UPDATE dokumente SET processing_quality=?, review_status=COALESCE(review_status,?) WHERE id=?", ("extract_error", "needs_review", doc_id))
    return result


def mark_duplicates(con: sqlite3.Connection) -> list[dict]:
    rows = list(con.execute("SELECT id,datei_name,datei_hash,local_original_path,dateipfad FROM dokumente"))
    by_hash = defaultdict(list)
    by_norm = defaultdict(list)
    for r in rows:
        if r["datei_hash"]:
            by_hash[r["datei_hash"]].append(r)
        norm = re.sub(r"[^a-z0-9]+", "", Path(r["datei_name"]).stem.lower())
        by_norm[norm].append(r)
    dups = []
    for key, group in by_hash.items():
        if len(group) > 1:
            keep = min(g["id"] for g in group)
            for g in group:
                if g["id"] != keep:
                    con.execute("UPDATE dokumente SET review_status='duplicate_candidate', processing_quality=COALESCE(processing_quality,'duplicate_hash') WHERE id=?", (g["id"],))
            dups.append({"type": "hash", "key": key[:12], "ids": [g["id"] for g in group], "keep": keep})
    for key, group in by_norm.items():
        if key and len(group) > 1:
            dups.append({"type": "name", "key": key[:40], "ids": [g["id"] for g in group]})
    return dups


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--force", action="store_true")
    ap.add_argument("--no-docling", action="store_true")
    ap.add_argument("--limit", type=int, default=0)
    args = ap.parse_args()
    BACKUPS.mkdir(parents=True, exist_ok=True); REPORTS.mkdir(parents=True, exist_ok=True)
    backup = BACKUPS / f"health_data_before_full_reprocess_{datetime.now():%Y%m%d_%H%M%S}.db"
    shutil.copy2(DB, backup)
    con = connect()
    before = dict(total=con.execute("SELECT count(*) FROM dokumente").fetchone()[0], without_text=con.execute("SELECT count(*) FROM dokumente WHERE length(COALESCE(extrahierte_inhalte,''))<100").fetchone()[0])
    rows = list(con.execute("SELECT * FROM dokumente WHERE ? OR length(COALESCE(extrahierte_inhalte,''))<100 ORDER BY CASE status WHEN 'neu' THEN 0 WHEN 'eingearbeitet' THEN 1 ELSE 2 END, id DESC", (1 if args.force else 0,)))
    if args.limit:
        rows = rows[:args.limit]
    results = []
    t0 = time.time()
    batch_id = f"full_reprocess_{datetime.now():%Y%m%d_%H%M%S}"
    con.execute("INSERT INTO health_processing_runs(batch_id,run_type,status,started_at,summary) VALUES(?,?,?,CURRENT_TIMESTAMP,?)", (batch_id, "full_document_reprocess", "running", json.dumps(before)))
    run_id = con.execute("SELECT last_insert_rowid()").fetchone()[0]
    con.commit()
    for i, row in enumerate(rows, start=1):
        res = process_one(con, row, force=args.force, use_docling=not args.no_docling)
        results.append(res)
        con.commit()
        print(json.dumps({"progress": f"{i}/{len(rows)}", **res}, ensure_ascii=False), flush=True)
    dups = mark_duplicates(con)
    after = dict(total=con.execute("SELECT count(*) FROM dokumente").fetchone()[0], without_text=con.execute("SELECT count(*) FROM dokumente WHERE length(COALESCE(extrahierte_inhalte,''))<100").fetchone()[0], good=con.execute("SELECT count(*) FROM dokumente WHERE processing_quality='good'").fetchone()[0], partial=con.execute("SELECT count(*) FROM dokumente WHERE processing_quality='partial'").fetchone()[0], missing=con.execute("SELECT count(*) FROM dokumente WHERE processing_quality='missing_file'").fetchone()[0])
    summary = {"backup": str(backup), "before": before, "after": after, "processed": len(results), "ok": sum(1 for r in results if r["ok"]), "duplicates": dups[:50], "elapsed_s": round(time.time()-t0, 1)}
    con.execute("UPDATE health_processing_runs SET status=?, finished_at=CURRENT_TIMESTAMP, summary=? WHERE id=?", ("completed", json.dumps(summary, ensure_ascii=False), run_id))
    con.commit(); con.close()
    out = REPORTS / f"full_document_reprocess_{datetime.now():%Y%m%d_%H%M%S}.json"
    out.write_text(json.dumps({"summary": summary, "results": results}, indent=2, ensure_ascii=False), encoding="utf-8")
    print("SUMMARY", json.dumps(summary, ensure_ascii=False), flush=True)
    print(out)

if __name__ == "__main__":
    main()