#!/usr/bin/env python3
"""
Scope triage v1 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.
- Daily model first; Architect only for uncertain/fragile cases.
- If Daily and Architect disagree, route conservatively to 01_INBOX.
- 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-scope classifier.
Your only job is to decide the PRIMARY life/work area of ONE file.

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

Allowed scope values:
- 04_ARBEIT = primarily employer/professional work material
- 03_PRIVAT = primarily the user's personal/private material
- 01_INBOX = genuinely unclear, mixed, or not safe to decide

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-related.
- Choose 04_ARBEIT for professional correspondence, employer projects, work events,
  work-team records, work-area administration, professional prevention work,
  work reimbursements/expenses, and material primarily owned by the employer/work role.
- Choose 03_PRIVAT for the user's own personal records, personal banking/customer documents,
  private travel/event tickets, personal planning, personal study/reference material,
  and the user's own employment/contract records kept as personal documentation.
- 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.
- Mentions of an employer inside a personal document do not automatically make it work.
- Mentions of a person/private name inside a work document do not automatically make it private.
- confidence must be 0.0 to 1.0. Use >=0.85 only when genuinely clear.
- reason: maximum 12 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]:
    # Scope output is tiny; timeouts still mirror the model families used by V9.
    name = model.lower()
    if "gpt-oss" in name or "architect" in name:
        return 450, 180
    if "35b" in name or "debug" in name:
        return 500, 240
    return 300, 90


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 needs_fallback(result: Optional[dict]) -> bool:
    if not isinstance(result, dict):
        return True
    if result.get("scope") not in {"04_ARBEIT", "03_PRIVAT", "01_INBOX"}:
        return True
    try:
        if float(result.get("confidence", 0)) < 0.85:
            return True
    except (TypeError, ValueError):
        return True

    status = str(result.get("content_status", ""))
    if "no-text" in status or status.startswith("error:") or status.startswith("missing:"):
        return True
    return False


def classify(item: dict, endpoint: str, model: str, fallback_model: str) -> dict:
    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 = bool(fallback_model) and (
        primary is None or needs_fallback(primary)
    )

    if not use_fallback:
        result = dict(primary)
        result["model_path"] = "primary"
        return result

    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:
        return {
            "scope": "01_INBOX",
            "confidence": 0.0,
            "reason": f"both models failed: {primary_error}; {fallback_error}"[:160],
            "content_status": "model-error",
            "model_path": "failure",
        }

    if primary is None:
        result = dict(fallback)
        result["model_path"] = "fallback-only"
        return result

    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]
        result["model_path"] = "primary-fallback-failed"
        return result

    if primary["scope"] == fallback["scope"]:
        result = dict(fallback)
        result["model_path"] = "models-agree"
        return result

    return {
        "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", ""),
        "model_path": "models-disagree",
    }


FIELDS = [
    "key", "origin", "relative_path", "source_context",
    "filename", "extension", "date_guess", "date_source",
    "scope", "confidence", "reason", "content_status", "model_path",
]


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_v1.tsv"
            if args.pilot_mixed
            else "scope_triage_readable_v1.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"],
        }

        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()