document_sorter/scope_triage.py

#!/usr/bin/env python3
"""
Scope triage v2 for the staged document-recovery project.

Purpose:
- Read-only top-level classification before detailed sorting.
- Exactly three outcomes: 04_ARBEIT, 03_PRIVAT, 01_INBOX.
- One file per LLM request.
- Curated files under ARCHIVE_EXTRACTED/WORK are deterministically 04_ARBEIT.
- Daily model first for ordinary SOURCE files; Architect verifies uncertain or policy-fragile cases.
- Architect receives a larger reasoning/output budget.
- If Daily and Architect disagree, route conservatively to 01_INBOX.
- Report keeps both model votes and review flags for auditability.
- Resume-safe TSV report written after every completed file.

This script does NOT move, rename, delete, or modify source/target files.
It only reads files and writes a TSV report under the staging MANIFESTS folder.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import os
import tempfile
import urllib.request
from collections import Counter
from pathlib import Path
from typing import Optional

import quick_sorter as q

READABLE = {
    ".pdf", ".docx", ".xlsx", ".pptx",
    ".odt", ".txt", ".md", ".csv",
}

MIXED_SOURCE_NAME = "Arbeits_Rechner_BackUP_Dokumente_Duplikate_vorhanden"

SCOPE_PROMPT = """You are a conservative top-level file-ownership classifier.
Your only job is to decide the PRIMARY ownership/life area of ONE file.

Return one JSON OBJECT only with exactly these keys:
scope, confidence, reason

Allowed scope values:
- 04_ARBEIT = primarily an operational/professional record used for the user's employer,
  job role, work projects, work teams, professional events, or work administration.
- 03_PRIVAT = primarily the user's own personal/private record as an individual.
- 01_INBOX = ownership/purpose is genuinely unclear, mixed, generic, or unsafe to decide.

CRITICAL OWNERSHIP RULE:
Do NOT ask merely whether the subject is 'related to work'. Ask who primarily owns/uses
the record and why it exists. A document can mention work but still be a private record.

Decision rules:
- Source folder is provenance/evidence, NOT authority.
- A file found on a work computer can still be private.
- A file found in a private folder can still be work material.
- Choose 04_ARBEIT for correspondence written/sent in a professional role, employer
  projects, Checkpoint/Aidshilfe/Cafe-Bach operations, work-team records, professional
  prevention work, work-event planning, employer administration, and work expenses or
  reimbursements created for professional activity.
- Friendly or personal wording inside professional correspondence does NOT make it private.
- Choose 03_PRIVAT for the user's own employment contract, payslip, personal bank/account
  documents, personal insurance/tax records, personal travel/event tickets, personal daily
  planning, and personal health/vaccination records. The user's own employment contract is
  PRIVATE documentation even though employment is work-related.
- Generic sample files, blank import templates, font/software licences, generic reference
  material, or context-free technical assets are NOT automatically work. If their concrete
  use/ownership is unclear, choose 01_INBOX.
- Generic medical/scientific guidance is not automatically a personal health record. If no
  personal patient information and no clear work use is present, choose 01_INBOX.
- Product comparisons/offers are 04_ARBEIT only when the excerpt shows a concrete employer
  project or professional purchase context; otherwise choose 01_INBOX.
- Do NOT route to DEV, VAULT, ARCHIV, or any other area in this recovery pass.
- If work/private ownership cannot be determined safely, choose 01_INBOX.

Anchor examples:
- A New-Year greeting to a venue that says the sender works for Checkpoint/Aidshilfe and
  discusses prevention materials/collaboration -> 04_ARBEIT.
- The user's own Arbeitsvertrag -> 03_PRIVAT.
- A blank/sample contact-import CSV -> 01_INBOX unless actual work records are present.
- A generic Roboto/font LICENSE -> 01_INBOX unless tied to a concrete work deliverable.

Confidence calibration:
- 0.90-0.95 only with direct ownership/purpose evidence.
- 0.75-0.89 when likely but some interpretation is required.
- <=0.70 when evidence is thin or mixed; prefer 01_INBOX.
- Never use high confidence merely because of the source folder.
- reason: maximum 14 words, factual and concise.
"""


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 build_queue(stage: Path, policy_path: Path) -> list[dict]:
    source = stage / "SOURCE"
    target = stage / "TARGET"
    extracted = stage / "ARCHIVE_EXTRACTED" / "WORK"

    policy = q.load_policy(policy_path)
    destinations = q.allowed_destinations(source, policy, "04_ARBEIT")
    aliases = policy.get("_strong_aliases", {}).get("04_ARBEIT", {})

    records = q.scan(source, target, destinations, aliases=aliases)
    q.global_duplicate_prepass(records, source, target)

    items: list[dict] = []

    for rec in records:
        if rec.status != "pending":
            continue
        if q.is_archive_file(rec.path):
            continue
        if rec.ext.lower() not in READABLE:
            continue

        items.append({
            "key": "SOURCE\t" + rec.rel_path,
            "origin": "SOURCE",
            "path": rec.path,
            "rel_path": rec.rel_path,
            "source_context": rec.source_context,
            "date_guess": rec.date_guess,
            "date_source": rec.date_source,
        })

    source_hashes = {
        sha256(p)
        for p in source.rglob("*")
        if p.is_file()
    }
    target_hashes = {
        sha256(p)
        for p in target.rglob("*")
        if p.is_file()
    }

    seen: set[str] = set()

    for p in sorted(extracted.rglob("*")):
        if not p.is_file() or p.suffix.lower() not in READABLE:
            continue

        digest = sha256(p)
        if digest in source_hashes or digest in target_hashes or digest in seen:
            continue
        seen.add(digest)

        rel = str(p.relative_to(extracted))
        # Preserve meaningful archive provenance at least at recovery-group level.
        group = Path(rel).parts[0] if Path(rel).parts else "ARCHIVE"
        source_context = f"archive-recovered/{group}/{Path(rel).parent}"

        date_guess, date_source = q.best_date(p, root=extracted)

        items.append({
            "key": "ARCHIVE\t" + rel,
            "origin": "ARCHIVE",
            "path": p,
            "rel_path": rel,
            "source_context": source_context,
            "date_guess": date_guess,
            "date_source": date_source,
        })

    return items


def model_limits(model: str) -> tuple[int, int]:
    # Reasoning models may spend substantial tokens before emitting their tiny JSON answer.
    name = model.lower()
    if "gpt-oss" in name or "architect" in name:
        return 2500, 240
    if "35b" in name or "debug" in name:
        return 6000, 300
    return 500, 120


def extract_json_object(text: str) -> dict:
    text = text.strip()
    if text.startswith("```"):
        text = text.strip("`")
        if text.lower().startswith("json"):
            text = text[4:].lstrip()

    try:
        value = json.loads(text)
    except json.JSONDecodeError:
        start = text.find("{")
        end = text.rfind("}")
        if start < 0 or end <= start:
            raise
        value = json.loads(text[start:end + 1])

    if not isinstance(value, dict):
        raise ValueError("model did not return a JSON object")
    return value


def call_scope_model(item: dict, endpoint: str, model: str) -> dict:
    text, content_status = q.extract_excerpt(item["path"])

    payload_item = {
        "filename": item["path"].name,
        "relative_path": item["rel_path"],
        "source_context": item["source_context"],
        "date_guess": item["date_guess"],
        "date_source": item["date_source"],
        "content_status": content_status,
        "excerpt": text[:q.MAX_EXCERPT],
    }

    prompt = (
        SCOPE_PROMPT
        + "\nINPUT:\n"
        + json.dumps(payload_item, ensure_ascii=False)
    )

    max_tokens, timeout = model_limits(model)
    request_payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0.0,
        "max_tokens": max_tokens,
    }

    req = urllib.request.Request(
        endpoint,
        data=json.dumps(request_payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )

    with urllib.request.urlopen(req, timeout=timeout) as resp:
        data = json.loads(resp.read().decode("utf-8"))

    choice = data["choices"][0]
    answer = choice.get("message", {}).get("content") or ""
    result = extract_json_object(answer)

    scope = str(result.get("scope", "")).strip()
    try:
        confidence = float(result.get("confidence", 0))
    except (TypeError, ValueError):
        confidence = 0.0

    reason = str(result.get("reason", "")).strip()[:160]

    if scope not in {"04_ARBEIT", "03_PRIVAT", "01_INBOX"}:
        raise ValueError(f"invalid scope: {scope!r}")

    confidence = max(0.0, min(1.0, confidence))

    return {
        "scope": scope,
        "confidence": confidence,
        "reason": reason,
        "content_status": content_status,
    }


def scope_review_flags(item: dict, result: Optional[dict]) -> list[str]:
    flags: list[str] = []
    if not isinstance(result, dict):
        return ["no-primary-result"]

    scope = str(result.get("scope", ""))
    filename = item["path"].name.lower()
    rel = str(item.get("rel_path", "")).lower()

    try:
        excerpt, _ = q.extract_excerpt(item["path"])
    except Exception:
        excerpt = ""
    content = excerpt.lower()
    combined = " ".join((filename, rel, content))

    # Classes where broad 'work-related' reasoning often confuses ownership.
    always_fragile = (
        "anschreiben", "arbeitsvertrag", "ticket", "konto",
        "kundeninformation", "sample", "testsample", "license", "lizenz",
        "preisvergleich", "impf", "tagesplan",
    )
    if any(term in filename for term in always_fragile):
        flags.append("ownership-fragile-filename")

    strong_work = (
        "aidshilfe", "checkpoint", "cafe bach", "café bach", "herzenslust",
        "prävention", "praevention", "vor-ort", "fak", "dienstplan",
    )
    if scope == "03_PRIVAT" and any(term in combined for term in strong_work):
        flags.append("private-vs-strong-work-evidence")

    personal_record = (
        "arbeitsvertrag", "konto", "ticket", "versicherung", "steuer",
        "gehalts", "lohnabrechnung", "impfpass",
    )
    if scope == "04_ARBEIT" and any(term in combined for term in personal_record):
        flags.append("work-vs-personal-record-class")

    generic_asset = ("sample", "testsample", "license", "lizenz")
    if scope == "04_ARBEIT" and any(term in filename for term in generic_asset):
        flags.append("work-vs-generic-asset")

    return flags


def needs_fallback(item: dict, result: Optional[dict]) -> tuple[bool, list[str]]:
    flags = scope_review_flags(item, result)
    if not isinstance(result, dict):
        return True, flags
    if result.get("scope") not in {"04_ARBEIT", "03_PRIVAT", "01_INBOX"}:
        flags.append("invalid-primary-scope")
        return True, flags
    try:
        if float(result.get("confidence", 0)) < 0.85:
            flags.append("primary-confidence-below-0.85")
    except (TypeError, ValueError):
        flags.append("invalid-primary-confidence")

    status = str(result.get("content_status", ""))
    if "no-text" in status or status.startswith("error:") or status.startswith("missing:"):
        flags.append("limited-content")

    return bool(flags), flags


def classify(item: dict, endpoint: str, model: str, fallback_model: str) -> dict:
    # ARCHIVE items in this queue come only from ARCHIVE_EXTRACTED/WORK.
    # That directory was already curated and verified as work material during
    # the archive-recovery phase; calendar/private archive payload is held
    # separately under HOLD_CALENDAR and never enters build_queue().
    if item.get("origin") == "ARCHIVE":
        return {
            "scope": "04_ARBEIT",
            "confidence": 0.99,
            "reason": "Curated work archive provenance verified before scope triage.",
            "content_status": "provenance:curated-work-archive",
            "model_path": "deterministic-provenance",
            "review_flags": "",
            "primary_scope": "",
            "primary_confidence": "",
            "fallback_scope": "",
            "fallback_confidence": "",
        }

    primary = None
    primary_error = ""

    try:
        primary = call_scope_model(item, endpoint, model)
    except Exception as exc:
        primary_error = f"{type(exc).__name__}: {exc}"

    use_fallback, review_flags = needs_fallback(item, primary)
    use_fallback = bool(fallback_model) and use_fallback

    def decorate(result: dict, model_path: str, fallback: Optional[dict] = None) -> dict:
        out = dict(result)
        out["model_path"] = model_path
        out["review_flags"] = ",".join(review_flags)
        out["primary_scope"] = primary.get("scope", "") if primary else ""
        out["primary_confidence"] = primary.get("confidence", "") if primary else ""
        out["fallback_scope"] = fallback.get("scope", "") if fallback else ""
        out["fallback_confidence"] = fallback.get("confidence", "") if fallback else ""
        return out

    if not use_fallback:
        return decorate(primary, "primary")

    fallback = None
    fallback_error = ""
    try:
        fallback = call_scope_model(item, endpoint, fallback_model)
    except Exception as exc:
        fallback_error = f"{type(exc).__name__}: {exc}"

    if primary is None and fallback is None:
        result = {
            "scope": "01_INBOX",
            "confidence": 0.0,
            "reason": f"both models failed: {primary_error}; {fallback_error}"[:160],
            "content_status": "model-error",
        }
        return decorate(result, "failure")

    if primary is None:
        return decorate(fallback, "fallback-only", fallback)

    if fallback is None:
        result = dict(primary)
        result["confidence"] = min(float(result.get("confidence", 0)), 0.69)
        result["reason"] = (
            f"fallback failed; {result.get('reason', '')}"
        )[:160]
        return decorate(result, "primary-fallback-failed")

    if primary["scope"] == fallback["scope"]:
        # Agreement is useful, but keep confidence conservative: never higher than the weaker vote.
        result = dict(fallback)
        result["confidence"] = min(
            float(primary.get("confidence", 0)),
            float(fallback.get("confidence", 0)),
        )
        return decorate(result, "models-agree", fallback)

    result = {
        "scope": "01_INBOX",
        "confidence": 0.50,
        "reason": (
            f"model disagreement: {model}={primary['scope']}; "
            f"{fallback_model}={fallback['scope']}"
        )[:160],
        "content_status": primary.get("content_status", ""),
    }
    return decorate(result, "models-disagree", fallback)


FIELDS = [
    "key", "origin", "relative_path", "source_context",
    "filename", "extension", "date_guess", "date_source",
    "scope", "confidence", "reason", "content_status", "model_path",
    "review_flags", "primary_scope", "primary_confidence",
    "fallback_scope", "fallback_confidence",
]


def load_existing(report: Path) -> dict[str, dict]:
    if not report.is_file():
        return {}
    rows = {}
    with report.open(newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f, delimiter="\t")
        for row in reader:
            key = row.get("key", "")
            if key:
                rows[key] = row
    return rows


def write_report_atomic(report: Path, rows: dict[str, dict], order: list[str]) -> None:
    report.parent.mkdir(parents=True, exist_ok=True)
    fd, temp_name = tempfile.mkstemp(
        prefix=report.name + ".",
        suffix=".tmp",
        dir=report.parent,
        text=True,
    )
    os.close(fd)
    temp = Path(temp_name)
    try:
        with temp.open("w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=FIELDS, delimiter="\t")
            writer.writeheader()
            for key in order:
                if key in rows:
                    writer.writerow({name: rows[key].get(name, "") for name in FIELDS})
        temp.replace(report)
    finally:
        if temp.exists():
            temp.unlink()


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--stage",
        default="/data/90_SORT_TEST_ARBEIT_IMPORT/FULL_STAGING_2026-09-14",
    )
    ap.add_argument(
        "--policy",
        default=str(Path(__file__).with_name("sorter_policy.json")),
    )
    ap.add_argument("--endpoint", default=q.DEFAULT_ENDPOINT)
    ap.add_argument("--model", default=q.DEFAULT_MODEL)
    ap.add_argument("--fallback-model", default=q.DEFAULT_FALLBACK_MODEL)
    ap.add_argument("--no-fallback", action="store_true")
    ap.add_argument(
        "--report",
        default="",
        help="TSV report path; defaults under STAGE/MANIFESTS",
    )
    ap.add_argument(
        "--pilot-mixed",
        action="store_true",
        help="Only the original readable files from the known mixed backup",
    )
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    stage = Path(args.stage).expanduser().resolve()
    policy_path = Path(args.policy).expanduser().resolve()

    if not stage.is_dir():
        raise SystemExit(f"Not a staging directory: {stage}")
    if not policy_path.is_file():
        raise SystemExit(f"Policy not found: {policy_path}")

    items = build_queue(stage, policy_path)

    if args.pilot_mixed:
        prefix = MIXED_SOURCE_NAME + "/"
        items = [
            item for item in items
            if item["origin"] == "SOURCE"
            and item["rel_path"].startswith(prefix)
        ]

    if args.limit > 0:
        items = items[:args.limit]

    if args.report:
        report = Path(args.report).expanduser().resolve()
    else:
        name = (
            "scope_triage_pilot_mixed_v3.tsv"
            if args.pilot_mixed
            else "scope_triage_readable_v3.tsv"
        )
        report = stage / "MANIFESTS" / name

    print("=== SCOPE TRIAGE ===")
    print("Queue:", len(items))
    print("Mode:", "pilot-mixed" if args.pilot_mixed else "full-readable")
    print("Allowed outcomes: 04_ARBEIT, 03_PRIVAT, 01_INBOX")
    print("Primary model:", args.model)
    fallback_model = "" if args.no_fallback else args.fallback_model.strip()
    print("Fallback model:", fallback_model or "disabled")
    print("Report:", report)

    if args.dry_run:
        print("DRY RUN: no LLM requests made")
        for item in items[:20]:
            print(item["origin"], item["rel_path"])
        return

    existing = load_existing(report)
    order = [item["key"] for item in items]
    rows = {k: v for k, v in existing.items() if k in set(order)}

    completed = len(rows)
    if completed:
        print(f"Resume: {completed} already completed")

    for index, item in enumerate(items, 1):
        key = item["key"]
        if key in rows:
            continue

        result = classify(
            item,
            args.endpoint,
            args.model,
            fallback_model,
        )

        row = {
            "key": key,
            "origin": item["origin"],
            "relative_path": item["rel_path"],
            "source_context": item["source_context"],
            "filename": item["path"].name,
            "extension": item["path"].suffix.lower(),
            "date_guess": item["date_guess"],
            "date_source": item["date_source"],
            "scope": result["scope"],
            "confidence": f"{float(result['confidence']):.2f}",
            "reason": result["reason"],
            "content_status": result["content_status"],
            "model_path": result["model_path"],
            "review_flags": result.get("review_flags", ""),
            "primary_scope": result.get("primary_scope", ""),
            "primary_confidence": (
                f"{float(result['primary_confidence']):.2f}"
                if result.get("primary_confidence", "") != "" else ""
            ),
            "fallback_scope": result.get("fallback_scope", ""),
            "fallback_confidence": (
                f"{float(result['fallback_confidence']):.2f}"
                if result.get("fallback_confidence", "") != "" else ""
            ),
        }

        rows[key] = row
        write_report_atomic(report, rows, order)

        print(
            f"[{index:3}/{len(items)}] "
            f"{result['scope']:9} "
            f"{float(result['confidence']):.2f}  "
            f"{item['path'].name}"
        )

    counts = Counter(row["scope"] for row in rows.values())
    low = sum(float(row["confidence"]) < 0.70 for row in rows.values())
    disagreements = sum(
        row.get("model_path") == "models-disagree"
        for row in rows.values()
    )

    print()
    print("=== RESULT ===")
    print("Completed:", len(rows))
    print("04_ARBEIT:", counts["04_ARBEIT"])
    print("03_PRIVAT:", counts["03_PRIVAT"])
    print("01_INBOX:", counts["01_INBOX"])
    print("Confidence < 0.70:", low)
    print("Model disagreements:", disagreements)
    print("Report:", report)


if __name__ == "__main__":
    main()