#!/usr/bin/env python3
"""
Extract structured data from the CSE internship D2/D3 forms.

Input mix (see README):
  .docx                  -> word/document.xml (text + tables)
  .doc                   -> LibreOffice headless conversion, then as .docx
  .pdf with real text    -> pdftotext -layout
  .pdf scan / OCR-garbage-> pdftoppm -r 200 -> vision-model transcription
  .jpeg                  -> EXIF-rotate -> vision-model transcription

The transcript is then mapped to a fixed JSON schema by a text model.
Everything is cached per file, so the run is resumable and re-runnable.

Credentials: OPENCODE_GO_API_KEY in the environment or in ./.env

Usage:
  python3 extract_forms.py --limit 5                  # smoke test
  python3 extract_forms.py --stage acquire            # text only, no LLM structuring
  python3 extract_forms.py                            # full run (acquire + structure)
  python3 extract_forms.py --aggregate                # rebuild outputs from cache
  python3 extract_forms.py --report                   # coverage / issues summary
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import requests
from xml.etree import ElementTree as ET

BASE_URL = "https://opencode.ai/zen/go/v1"
USER_AGENT = "internship-cse-extractor/1.0"
SESSION_ID = "cse-forms-extract-v1"

# OCR runs on ai-box qwen3.8-flash (images are rejected by, and paid on, some Go models).
VISION_MODEL = "qwen3.8-flash"
VISION_BASE_URL = "https://api.ai-box.vn/v1"
VISION_KEY_ENV = "AIBOX_API_KEY"
# (no paid vision model is used: the project forbids deepseek-v4-flash-vision-exp)
# Free (limited-time) models first, paid fallback per project decision.
TEXT_MODEL_CANDIDATES = ["step-5-preview-free", "longcat-2.5-preview-free", "deepseek-v4.1-flash"]

FORMS_DIR = Path("./forms")
D2_DIR = FORMS_DIR / "d2"
D3_DIR = FORMS_DIR / "d3"
INFO_FILE = FORMS_DIR / "company_info.json"
OUT_DIR = FORMS_DIR / "extracted"
CACHE_DIR = OUT_DIR / "cache"

RENDER_DPI = 200
REAL_TEXT_MIN_SCORE = 5  # occurrences of common Vietnamese form words
REAL_TEXT_MIN_LEN = 200
MAX_TEXT_CHARS = 150_000  # guard for the structuring prompt; the longest D2 is 40 pages (~90k chars)

# Words that survive pdftotext intact only when the PDF has a real text layer;
# scanned-then-OCR'd files come back with broken diacritics ("cONc rY TNHU").
REAL_TEXT_MARKERS = ["thực tập", "sinh viên", "chương trình", "Địa chỉ", "Công ty", "nghiệp"]

W = "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}"

VISION_PROMPT = (
    "Bạn là công cụ OCR cho văn bản hành chính tiếng Việt. "
    "Chép lại TOÀN BỘ chữ trong ảnh, giữ nguyên tiếng Việt có dấu, đúng chính tả. "
    "Giữ cấu trúc bảng: mỗi hàng một dòng, các ô ngăn bằng dấu |. "
    "Đánh dấu đoạn không đọc được là [?]. "
    "Không bình luận, không tóm tắt, không dịch, không thêm gì ngoài nội dung đọc được."
)

SCHEMA_D2 = {
    "company_name": "Tên đầy đủ của doanh nghiệp (string|null)",
    "company_address": "Địa chỉ trụ sở/địa chỉ văn phòng chính (string|null)",
    "work_address": "Địa chỉ làm việc thực tế nếu văn bản ghi riêng, khác địa chỉ trụ sở (string|null)",
    "tax_code": "Mã số thuế (string|null, giữ nguyên chữ số)",
    "phone": "string|null",
    "fax": "string|null",
    "email": "string|null",
    "website": "string|null",
    "semester": "Học kỳ, ví dụ 'HK253' hoặc 'HK252' (string|null)",
    "period": "Thời gian thực tập thực tế, ví dụ '15/06-14/08/2026' (string|null)",
    "intro": "Toàn bộ mục 'Giới thiệu sơ bộ về DN' (string|null)",
    "slots": "Số lượng suất thực tập dự kiến (integer|null, chỉ lấy số đầu tiên)",
    "slots_note": "Phần ghi chú kèm theo số suất (string|null)",
    "requirements": "Danh sách yêu cầu/đối tượng SV, mỗi mục là một string",
    "accepted_students": [
        {
            "no": "số thứ tự (integer|null)",
            "student_id": "Mã số sinh viên, giữ nguyên chữ số (string|null)",
            "ho_lot": "string|null",
            "ten": "string|null",
            "note": "Ghi chú (string|null)",
        }
    ],
    "duration": "Tổng thời gian làm việc của đợt thực tập (string|null)",
    "daily_hours": "Thời gian làm việc trong ngày (string|null)",
    "location": "Địa điểm SV làm việc thực tế theo văn bản (string|null)",
    "benefits": "Ưu đãi (string|null)",
    "extra_info": "Thông tin thêm (string|null)",
    "projects": [
        {
            "no": "số thứ tự đề tài (integer|null)",
            "title": "Tên đề tài (string|null)",
            "languages": "Danh sách ngôn ngữ lập trình (list of string)",
            "frameworks": "Danh sách framework/công nghệ (list of string)",
            "timeline": [{"week": "tuần (string)", "task": "nội dung công việc (string)"}],
        }
    ],
    "dossier": "Danh sách hồ sơ SV phải nộp (list of string)",
    "submission_channel": "Địa điểm/kênh tiếp nhận hồ sơ (string|null)",
    "deadline": "Hạn nộp hồ sơ (string|null)",
    "interview": "Thông tin phỏng vấn/xét tuyển (string|null)",
    "contact": {
        "name": "string|null",
        "gender": "string|null",
        "phone": "string|null",
        "email": "string|null",
    },
    "signature_place_date": "Địa danh và ngày ký ở cuối văn bản, nguyên văn (string|null)",
    "signer_name": "Họ tên người ký ở cuối văn bản (string|null)",
    "signer_title": "Chức vụ người ký (string|null)",
    "signature_date": "NGÀY KÝ của văn bản, chuẩn hoá ISO 'YYYY-MM-DD' (string|null)",
    "signature_date_raw": "Ngày ký nguyên văn trong văn bản, ví dụ 'ngày 04 tháng 05 năm 2026' (string|null)",
    "signed_and_stamped": "CHỈ true khi tài liệu là bản scan và nhìn thấy chữ ký tay hoặc dấu đỏ; file .docx (bản nháp) luôn là false",
    "is_official_template": "true nếu văn bản dùng đúng mẫu chương trình của Khoa (có mục I. NỘI DUNG, bảng tiêu chí, cam kết với Khoa)",
    "quality_notes": "Ghi chú về chất lượng bản scan/nội dung bất thường (string|null)",
}

SCHEMA_D3 = {
    "company_name": "Tên doanh nghiệp trong văn bản (string|null)",
    "date_text": "Địa danh, ngày tháng năm ở đầu công văn, nguyên văn (string|null)",
    "signature_date": "NGÀY KÝ của công văn, chuẩn hoá ISO 'YYYY-MM-DD'; nếu phần chữ ký có ngày viết tay thì ưu tiên ngày đó, nếu không lấy ngày ở đầu văn bản (string|null)",
    "signature_date_raw": "Ngày ký nguyên văn, ví dụ 'ngày 14 tháng 06 năm 2026' (string|null)",
    "semester": "Học kỳ, ví dụ 'HK253' (string|null)",
    "recipient": "Nơi nhận / Kính gửi (string|null)",
    "students": [
        {
            "no": "số thứ tự (integer|null)",
            "student_id": "Mã số sinh viên, giữ nguyên chữ số (string|null)",
            "ho_lot": "string|null",
            "ten": "string|null",
            "note": "Ghi chú (string|null)",
        }
    ],
    "signature_present": "true nếu có chữ ký/đóng dấu trong văn bản scan",
    "signer_name": "Họ tên người ký ở cuối công văn (string|null)",
    "signer_title": "Chức vụ người ký (string|null)",
    "quality_notes": "string|null",
}


# --------------------------------------------------------------------------- utils


def log(msg: str) -> None:
    print(msg, file=sys.stderr, flush=True)


def read_api_key(name: str = "OPENCODE_GO_API_KEY") -> str:
    key = os.environ.get(name, "").strip()
    if not key:
        env = Path(".env")
        if env.exists():
            for line in env.read_text(encoding="utf-8").splitlines():
                if line.strip().startswith(name):
                    key = line.split("=", 1)[1].strip().strip('"').strip("'")
                    break
    if not key:
        sys.exit(f"{name} not set (env or ./.env)")
    return key


def safe_name(path: Path) -> str:
    """Must match download_forms.py's sanitisation (non-word chars dropped)."""
    return re.sub(r"[^\w\s-]", "", path.name).strip().replace(" ", "_") or "unnamed"


def parse_form_date(name: str) -> str | None:
    dates = re.findall(r"(20\d{2})-(\d{2})-(\d{2})", name)
    if not dates:
        return None
    y, m, d = dates[-1]
    return f"{y}-{m}-{d}"


# ------------------------------------------------------------------------ LLM call


class Client:
    def __init__(self, key: str, workers: int, base_url: str = BASE_URL):
        self.key = key
        self.base_url = base_url.rstrip("/")
        self.session = requests.Session()
        self.session.headers.update(
            {
                "Authorization": f"Bearer {key}",
                "Content-Type": "application/json",
                "User-Agent": USER_AGENT,
                "x-opencode-session": SESSION_ID,
            }
        )
        adapter = requests.adapters.HTTPAdapter(pool_connections=workers, pool_maxsize=workers)
        self.session.mount("https://", adapter)
        self.lock = threading.Lock()
        self.usage = {"calls": 0, "prompt_tokens": 0, "completion_tokens": 0}
        self.vision_model = VISION_MODEL

    def chat(self, model: str, messages: list, max_tokens: int = 8000, temperature: float = 0.0) -> tuple[str, dict]:
        body = {"model": model, "messages": messages, "max_tokens": max_tokens, "temperature": temperature}
        delay = 2.0
        last_err = None
        for attempt in range(6):
            try:
                resp = self.session.post(f"{self.base_url}/chat/completions", json=body, timeout=600)
                if resp.status_code in (429, 500, 502, 503, 504, 529):
                    last_err = f"HTTP {resp.status_code}: {resp.text[:200]}"
                    time.sleep(delay)
                    delay = min(delay * 2, 60)
                    continue
                resp.raise_for_status()
                data = resp.json()
                usage = data.get("usage") or {}
                with self.lock:
                    self.usage["calls"] += 1
                    self.usage["prompt_tokens"] += usage.get("prompt_tokens", 0)
                    self.usage["completion_tokens"] += usage.get("completion_tokens", 0)
                return data["choices"][0]["message"].get("content") or "", usage
            except requests.HTTPError as e:
                code = e.response.status_code if e.response is not None else 0
                if code in (401, 403):
                    sys.exit(f"auth/permission failure ({code}): {e.response.text[:300]}")
                last_err = f"HTTP {code}: {e.response.text[:200] if e.response is not None else e}"
                time.sleep(delay)
                delay = min(delay * 2, 60)
            except Exception as e:  # network / JSON errors
                last_err = f"{type(e).__name__}: {e}"
                time.sleep(delay)
                delay = min(delay * 2, 60)
        raise RuntimeError(f"chat failed after retries: {last_err}")

    def vision(self, model: str, prompt: str, image: Path, max_tokens: int = 4000) -> str:
        return vision_call(self, model, prompt, image, max_tokens)

    def text_model(self) -> str:
        """First candidate that answers; remembers the choice per process."""
        if getattr(self, "_text_model", None) is None:
            self._text_model = TEXT_MODEL_CANDIDATES[0]
        return self._text_model

    def demote_text_model(self, failing: str, reason: str) -> None:
        with self.lock:
            if failing in TEXT_MODEL_CANDIDATES:
                idx = TEXT_MODEL_CANDIDATES.index(failing)
                nxt = TEXT_MODEL_CANDIDATES[idx + 1] if idx + 1 < len(TEXT_MODEL_CANDIDATES) else failing
                if nxt != self._text_model:
                    log(f"! text model {failing} unusable ({reason}); switching to {nxt}")
                self._text_model = nxt


def vision_call(client: "Client", model: str, prompt: str, image: Path, max_tokens: int = 4000) -> str:
    """Single-image OCR call against the given client (which may be a different provider)."""
    b64 = base64.b64encode(image.read_bytes()).decode()
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": prompt},
                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}},
            ],
        }
    ]
    return client.chat(model, messages, max_tokens=max_tokens)[0]


# ---------------------------------------------------------------- text acquisition


def docx_text(path: Path) -> str:
    root = ET.fromstring(_read_docx_document(path))
    body = root.find(W + "body")
    lines: list[str] = []

    def walk(node):
        for ch in node:
            if ch.tag == W + "p":
                lines.append("".join(t.text or "" for t in ch.iter(W + "t")))
            elif ch.tag == W + "tbl":
                for tr in ch.findall(W + "tr"):
                    cells = []
                    for tc in tr.findall(W + "tc"):
                        txt = " ".join("".join(t.text or "" for t in p.iter(W + "t")) for p in tc.iter(W + "p"))
                        cells.append(re.sub(r"\s+", " ", txt).strip())
                    lines.append(" | ".join(cells))
            else:
                walk(ch)

    walk(body)
    return "\n".join(lines)


def _read_docx_document(path: Path) -> bytes:
    import zipfile

    with zipfile.ZipFile(path) as z:
        return z.read("word/document.xml")


def legacy_doc_text(path: Path, tmp: Path) -> str:
    out = subprocess.run(
        ["soffice", "--headless", "--convert-to", "docx", "--outdir", str(tmp), str(path)],
        capture_output=True,
        text=True,
        timeout=300,
    )
    produced = tmp / (path.stem + ".docx")
    if not produced.exists():
        raise RuntimeError(f"soffice failed: {out.stdout[-200:]} {out.stderr[-200:]}")
    return docx_text(produced)


def pdf_pages(path: Path) -> int:
    out = subprocess.run(["pdfinfo", str(path)], capture_output=True, text=True, timeout=60).stdout
    m = re.search(r"^Pages:\s+(\d+)", out, re.M)
    return int(m.group(1)) if m else 0


def pdf_text(path: Path) -> str:
    return subprocess.run(
        ["pdftotext", "-layout", str(path), "-"], capture_output=True, text=True, timeout=300
    ).stdout


def has_real_text(text: str) -> bool:
    if len(text.strip()) < REAL_TEXT_MIN_LEN:
        return False
    return sum(text.count(w) for w in REAL_TEXT_MARKERS) >= REAL_TEXT_MIN_SCORE


MAX_IMAGE_LONG_SIDE = 2400  # px; oversized PDFs (24x35 in pages) blow past the API's image limit
MAX_IMAGE_BYTES = 4_000_000


def fit_image(path: Path) -> Path:
    """Downscale / re-encode a rendered page so the vision API accepts it.

    Oversized pages (e.g. 1759x2519 pt) rendered at 200 dpi reach ~4900x7000 px
    and are rejected upstream; scale the long side down and fall back to JPEG if
    the PNG is still too heavy."""
    from PIL import Image

    im = Image.open(path)
    long_side = max(im.size)
    if long_side > MAX_IMAGE_LONG_SIDE:
        scale = MAX_IMAGE_LONG_SIDE / long_side
        im = im.resize((max(1, round(im.width * scale)), max(1, round(im.height * scale))), Image.LANCZOS)
    if im.mode not in ("RGB", "L"):
        im = im.convert("RGB")
    im.save(path, "PNG", optimize=True)
    if path.stat().st_size <= MAX_IMAGE_BYTES:
        return path
    jpg = path.with_suffix(".jpg")
    im.save(jpg, "JPEG", quality=88, optimize=True)
    path.unlink()
    return jpg


def render_pages(path: Path, out_dir: Path) -> list[Path]:
    out_dir.mkdir(parents=True, exist_ok=True)
    subprocess.run(
        ["pdftoppm", "-r", str(RENDER_DPI), "-png", str(path), str(out_dir / "page")],
        capture_output=True,
        timeout=1800,
        check=True,
    )
    return [fit_image(p) for p in sorted(out_dir.glob("page-*.png"))]


def normalize_image(path: Path, out_dir: Path) -> Path:
    """EXIF-rotate phone photos, re-encode to a vision-API-friendly PNG."""
    from PIL import Image, ImageOps

    out_dir.mkdir(parents=True, exist_ok=True)
    dest = out_dir / "page-1.png"
    with Image.open(path) as im:
        im = ImageOps.exif_transpose(im)
        if im.mode not in ("RGB", "L"):
            im = im.convert("RGB")
        im.save(dest, "PNG")
    return fit_image(dest)


def acquire_text(client: Client, path: Path, tmp_root: Path, use_ocr: bool = True) -> tuple[str, str]:
    """Return (text, source). source in {docx, doc-legacy, pdf-text, pdf-ocr, jpeg-ocr}."""
    name = path.name.lower()
    if name.endswith("docx"):
        return docx_text(path), "docx"
    if name.endswith("doc"):
        return legacy_doc_text(path, tmp_root), "doc-legacy"
    if name.endswith("jpeg") or name.endswith("jpg") or name.endswith("png"):
        if not use_ocr:
            return "", "jpeg-ocr"
        img = normalize_image(path, tmp_root / "img")
        return vision_call(client.ocr, client.vision_model, VISION_PROMPT, img), "jpeg-ocr"
    if name.endswith("pdf"):
        text = pdf_text(path)
        if has_real_text(text):
            return text, "pdf-text"
        if not use_ocr:
            return text, "pdf-scan"
        pages = render_pages(path, tmp_root / "pages")
        chunks = []
        for i, page in enumerate(pages, 1):
            chunks.append(f"=== TRANG {i} ===\n" + vision_call(client.ocr, client.vision_model, VISION_PROMPT, page))
        return "\n\n".join(chunks), "pdf-ocr"
    raise RuntimeError(f"unsupported file type: {path.name}")


# ------------------------------------------------------------------ struct parsing


def extract_json(raw: str) -> dict | None:
    if not raw:
        return None
    raw = raw.strip()
    if raw.startswith("```"):
        raw = re.sub(r"^```[a-zA-Z]*\n?", "", raw)
        raw = re.sub(r"\n?```\s*$", "", raw)
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        pass
    start = raw.find("{")
    while start != -1:
        depth = 0
        for i in range(start, len(raw)):
            if raw[i] == "{":
                depth += 1
            elif raw[i] == "}":
                depth -= 1
                if depth == 0:
                    try:
                        return json.loads(raw[start : i + 1])
                    except json.JSONDecodeError:
                        break
        start = raw.find("{", start + 1)
    return None


def structure_prompt(kind: str, text: str) -> list[dict]:
    schema = SCHEMA_D2 if kind == "d2" else SCHEMA_D3
    doc = "chương trình thực tập (form D2)" if kind == "d2" else "công văn xác nhận kết quả xét tuyển (form D3)"
    system = (
        "Bạn trích xuất dữ liệu có cấu trúc từ văn bản tiếng Việt của Khoa KH&KT Máy tính (ĐH Bách Khoa). "
        "Chỉ trả về MỘT đối tượng JSON hợp lệ, không giải thích, không markdown."
    )
    user = f"""Văn bản dưới đây là {doc} do doanh nghiệp gửi Khoa.

Trích xuất theo đúng JSON schema sau (giữ nguyên tên khoá, đủ tất cả khoá):
{json.dumps(schema, ensure_ascii=False, indent=2)}

Quy tắc:
- Không bịa. Thông tin không có trong văn bản thì để null (hoặc [] cho mảng).
- Giữ nguyên tiếng Việt có dấu; sửa lỗi OCR hiển nhiên (ví dụ "HQC" -> "HỌC") khi rõ nghĩa.
- Mã số sinh viên giữ nguyên chuỗi chữ số (7-8 chữ số).
- Ngày ký: luôn điền 'signature_date' dạng YYYY-MM-DD. Ngày Việt Nam kiểu "ngày 04 tháng 05 năm 2026" là 2026-05-04. Nếu bản scan có ngày viết tay cạnh chữ ký/đóng dấu thì dùng ngày đó; nếu không có thì dùng ngày ghi ở đầu/cuối văn bản. Không suy đoán nếu văn bản không ghi ngày nào.
- Nếu văn bản là scan lỗi/nội dung không liên quan, vẫn trả JSON với các trường null và mô tả trong quality_notes.
- 'slots' chỉ lấy số đầu tiên; 'projects.timeline' giữ đúng số tuần có trong bảng.

VĂN BẢN:
---
{text[:MAX_TEXT_CHARS]}
---
"""
    return [{"role": "system", "content": system}, {"role": "user", "content": user}]


def structure(client: Client, kind: str, text: str) -> tuple[dict | None, str]:
    """Returns (fields, model_used).

    Free models are tried first; a model is demoted only when it answers with
    unusable JSON or fails twice in this call (transient network errors are
    retried on the same model).
    """
    msgs = structure_prompt(kind, text)
    errors_here = 0
    used = client.text_model()
    for attempt in range(4):
        model = client.text_model()
        used = model
        try:
            raw, _u = client.chat(model, msgs, max_tokens=32000)
        except Exception as e:
            errors_here += 1
            if errors_here >= 2:
                client.demote_text_model(model, f"repeated error {e}")
                errors_here = 0
            continue
        fields = extract_json(raw)
        if fields is not None:
            return fields, model
        client.demote_text_model(model, "invalid JSON")
        errors_here = 0
    return None, used


# ------------------------------------------------------------------- validation


def norm_mssv(value) -> str | None:
    if value is None:
        return None
    digits = re.sub(r"\D", "", str(value))
    return digits or None


def validate(kind: str, f: dict) -> list[str]:
    issues = []
    sd = f.get("signature_date")
    if sd and not re.fullmatch(r"\d{4}-\d{2}-\d{2}", str(sd).strip()):
        issues.append(f"signature_date not ISO: {sd!r}")
    if kind == "d2":
        students = f.get("accepted_students") or []
        for s in students:
            sid = norm_mssv(s.get("student_id"))
            if sid and not (6 <= len(sid) <= 9):
                issues.append(f"student_id suspicious: {s.get('student_id')!r}")
        tax = f.get("tax_code")
        if tax and not re.fullmatch(r"\d{10}(-?\d{3})?", str(tax).strip()):
            digits = re.sub(r"\D", "", str(tax))
            issues.append(f"tax_code not 10/13 digits in source: {tax!r} ({len(digits)} digits)")
        if f.get("slots") is not None and not isinstance(f["slots"], int):
            issues.append(f"slots not int: {f['slots']!r}")
        for p in f.get("projects") or []:
            if not p.get("title"):
                issues.append("project without title")
    else:
        if not (f.get("students") or []):
            issues.append("no students extracted")
        for s in f.get("students") or []:
            sid = norm_mssv(s.get("student_id"))
            if sid and not (6 <= len(sid) <= 9):
                issues.append(f"student_id suspicious: {s.get('student_id')!r}")
    return issues


# ---------------------------------------------------------------------- pipeline


def cache_path(kind: str, path: Path) -> Path:
    return CACHE_DIR / f"{kind}_{safe_name(path)}.json"


def load_cache(kind: str, path: Path) -> dict:
    cp = cache_path(kind, path)
    if cp.exists():
        try:
            return json.loads(cp.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            return {}
    return {}


def save_cache(kind: str, path: Path, rec: dict) -> None:
    cp = cache_path(kind, path)
    cp.parent.mkdir(parents=True, exist_ok=True)
    tmp = cp.with_suffix(".tmp")
    tmp.write_text(json.dumps(rec, ensure_ascii=False), encoding="utf-8")
    tmp.replace(cp)


def process(client: Client, kind: str, path: Path, stage: str, force_text: bool = False, force_fields: bool = False) -> str:
    rec = load_cache(kind, path)
    rec.update({"file": path.name, "kind": kind, "form_date": parse_form_date(path.name)})
    rec.pop("error", None)  # a previous failure must not stick after a successful re-run
    need_text = stage in ("all", "acquire") and (force_text or not rec.get("text"))
    need_fields = stage in ("all", "structure") and (force_fields or not rec.get("fields"))

    if not need_text and not need_fields:
        return "cached"

    tmp_root = Path(tempfile.mkdtemp(prefix="forms_ocr_"))
    try:
        if need_text:
            text, source = acquire_text(client, path, tmp_root, use_ocr=True)
            rec["text"] = text
            rec["text_source"] = source
            rec["text_chars"] = len(text)
            if source and source.endswith("ocr"):
                rec["vision_model"] = client.vision_model
            save_cache(kind, path, rec)
            if stage == "acquire":
                return f"text:{source}"
        if need_fields:
            text = rec.get("text", "")
            # An image-only PDF without OCR text cannot be structured.
            if len(text.strip()) < 80:
                rec["fields"] = None
                rec["issues"] = ["no text"]
                save_cache(kind, path, rec)
                return "no-text"
            fields, model = structure(client, kind, text)
            rec["fields"] = fields
            rec["struct_model"] = model
            rec["issues"] = validate(kind, fields) if fields else ["structure failed"]
            rec["ts"] = time.time()
            save_cache(kind, path, rec)
            return "structured" if fields else "structure-failed"
        return "text-only"
    except Exception as e:
        rec["error"] = f"{type(e).__name__}: {e}"
        save_cache(kind, path, rec)
        return f"error: {e}"
    finally:
        shutil.rmtree(tmp_root, ignore_errors=True)


def collect(kind: str, only: str | None = None) -> list[Path]:
    directory = D2_DIR if kind == "d2" else D3_DIR
    if not directory.exists():
        return []
    files = [p for p in sorted(directory.iterdir()) if p.is_file()]
    if only:
        files = [p for p in files if only.lower() in p.name.lower()]
    return files


def is_ocr_candidate(path: Path) -> bool:
    """True when acquire_text would need the vision model for this file."""
    name = path.name.lower()
    if name.endswith("docx") or name.endswith("doc"):
        return False
    if name.endswith(("jpeg", "jpg", "png")):
        return True
    if name.endswith("pdf"):
        return not has_real_text(pdf_text(path))
    return False


def needs_retry(kind: str, path: Path) -> bool:
    rec = load_cache(kind, path)
    if rec.get("error") or not rec.get("text"):
        return True
    if len(rec.get("text", "").strip()) < 80:
        return False  # image-only file legitimately has no text
    return not rec.get("fields")


def run(args) -> None:
    key = args.api_key or read_api_key("OPENCODE_GO_API_KEY")
    client = Client(key, args.workers, base_url=args.base_url or BASE_URL)
    if args.text_model:
        client._text_model = args.text_model
    # OCR runs on its own provider (ai-box qwen3.8-flash by default).
    client.ocr = Client(
        args.vision_api_key or read_api_key(VISION_KEY_ENV),
        args.workers,
        base_url=args.vision_base_url or VISION_BASE_URL,
    )
    client.vision_model = args.vision_model or VISION_MODEL
    kinds = ["d2", "d3"] if args.kind == "both" else [args.kind]

    for kind in kinds:
        files = collect(kind, args.only)
        if args.limit:
            files = files[: args.limit]
        if args.retry_failed:
            files = [f for f in files if needs_retry(kind, f)]
        if args.only_ocr:
            files = [f for f in files if is_ocr_candidate(f)]
        if not files:
            log(f"[{kind}] nothing to do")
            continue
        log(f"[{kind}] {len(files)} files, stage={args.stage}, workers={args.workers}")
        done = {"cached": 0, "ok": 0, "fail": 0}
        t0 = time.time()
        with ThreadPoolExecutor(max_workers=args.workers) as pool:
            futures = {
                pool.submit(process, client, kind, f, args.stage, args.force_text, args.force_fields): f
                for f in files
            }
            for i, fut in enumerate(as_completed(futures), 1):
                status = fut.result()
                if status == "cached":
                    done["cached"] += 1
                elif status.startswith("error"):
                    done["fail"] += 1
                    log(f"  [{kind} {i}/{len(files)}] {futures[fut].name}: {status}")
                else:
                    done["ok"] += 1
                if i % 25 == 0 or i == len(files):
                    rate = i / max(time.time() - t0, 1e-9)
                    log(
                        f"  [{kind} {i}/{len(files)}] ok={done['ok']} cached={done['cached']} "
                        f"fail={done['fail']} {rate:.1f} files/s"
                    )
        log(
            f"[{kind}] done: ok={done['ok']} cached={done['cached']} fail={done['fail']} "
            f"in {time.time() - t0:.0f}s; text tokens in/out="
            f"{client.usage['prompt_tokens']}/{client.usage['completion_tokens']}; "
            f"ocr calls={client.ocr.usage['calls']} tokens in/out="
            f"{client.ocr.usage['prompt_tokens']}/{client.ocr.usage['completion_tokens']}"
        )


# --------------------------------------------------------------------- aggregates


def portal_index() -> dict[str, dict]:
    """local filename -> portal company/form metadata (mirrors download_forms.py naming)."""
    index: dict[str, dict] = {}
    if not INFO_FILE.exists():
        return index
    companies = json.loads(INFO_FILE.read_text(encoding="utf-8"))
    for c in companies:
        safe_company = re.sub(r"[^\w\s-]", "", c.get("shortname", "")).strip().replace(" ", "_")
        for kind, key in (("d2", "d2_forms"), ("d3", "d3_forms")):
            for name in c.get(key) or []:
                safe_form = re.sub(r"[^\w\s-]", "", name).strip().replace(" ", "_")
                index[f"{kind}_{safe_company}_{safe_form}.json"] = {
                    "company_id": c.get("id"),
                    "company_shortname": c.get("shortname"),
                    "company_fullname": c.get("fullname"),
                    "portal_location": c.get("location"),
                    "portal_max_slot": c.get("max_slot"),
                    "portal_student_accepted": c.get("student_accepted"),
                    "portal_form_name": name,
                }
    return index


def iter_cache(kind: str):
    for cp in sorted(CACHE_DIR.glob(f"{kind}_*.json")):
        try:
            rec = json.loads(cp.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            continue
        rec["cache_file"] = cp.name
        yield rec


VN_DATE_RE = re.compile(r"ng[àa]y\s*(\d{1,2})\s*th[áa]ng\s*(\d{1,2})\s*n[ăa]m\s*(\d{4})", re.IGNORECASE)


def iso_from_vn_date(text: str | None) -> str | None:
    """Parse 'ngày 04 tháng 05 năm 2026' -> '2026-05-04' (fallback when the model left it null)."""
    if not text:
        return None
    m = VN_DATE_RE.search(str(text))
    if not m:
        return None
    d, mo, y = (int(x) for x in m.groups())
    if 1 <= d <= 31 and 1 <= mo <= 12:
        return f"{y:04d}-{mo:02d}-{d:02d}"
    return None


def clean_signer_name(name: str | None) -> str | None:
    """Drop boilerplate that OCR appends to the signer block ('Ban Lãnh đạo', '(Ký & ghi rõ họ tên…)')."""
    if not name:
        return None
    s = re.sub(r"\(?K[ýy][^)]*\)?", " ", name)
    s = re.sub(r"Ban L[ãa]nh [Đđ]ạo.*$", " ", s, flags=re.IGNORECASE)
    s = re.sub(r"\s{2,}", " ", s).strip(" ,;:-.")
    return s or None


BLANK_DATE_RE = re.compile(r"ng[àa]y\s*[.…\s]{2,}th[áa]ng", re.IGNORECASE)


def sign_date_from_text(text: str | None) -> str | None:
    """Fallback: last date in the signature block, only if the form doesn't leave it blank."""
    if not text:
        return None
    tail = text[-800:]
    if BLANK_DATE_RE.search(tail):  # 'ngày …… tháng …… năm ……' — company did not date the form
        return None
    hits = list(VN_DATE_RE.finditer(tail))
    return iso_from_vn_date(hits[-1].group(0)) if hits else None


def sign_date_of(f: dict, text: str | None = None) -> tuple[str | None, str | None]:
    """(signature_date, source) — model value, else parsed from the raw string, else from the text."""
    sd = f.get("signature_date")
    if sd:
        return sd, "model"
    parsed = iso_from_vn_date(f.get("signature_date_raw"))
    if parsed:
        return parsed, "raw_parsed"
    parsed = sign_date_from_text(text)
    if parsed:
        return parsed, "text_parsed"
    return None, None


def aggregate() -> None:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    index = portal_index()
    import csv

    for kind in ("d2", "d3"):
        records = []
        for rec in iter_cache(kind):
            meta = index.get(rec["cache_file"], {})
            f = rec.get("fields") or {}
            sd, sd_source = sign_date_of(f, rec.get("text"))
            record = {
                "file": rec.get("file"),
                "form_date": rec.get("form_date"),
                "signature_date": sd,
                "signature_date_source": sd_source,
                "signature_date_raw": f.get("signature_date_raw"),
                "signer_name": clean_signer_name(f.get("signer_name")),
                "signer_title": f.get("signer_title"),
                "tax_code": re.sub(r"[^\d-]", "", str(f.get("tax_code"))) if f.get("tax_code") else None,
                "text_source": rec.get("text_source"),
                "vision_model": rec.get("vision_model"),
                "text_chars": rec.get("text_chars"),
                "struct_model": rec.get("struct_model"),
                "issues": rec.get("issues") or [],
                "error": rec.get("error"),
                **meta,
                "fields": f,
            }
            records.append(record)
        # deterministic order: company shortname, then file
        records.sort(key=lambda r: (str(r.get("company_shortname") or ""), str(r.get("file") or "")))
        out = OUT_DIR / f"{kind}.json"
        out.write_text(json.dumps(records, ensure_ascii=False, indent=2), encoding="utf-8")
        log(f"wrote {out} ({len(records)} records)")

    # one row per form file: sign date vs. filename date vs. extracted semester
    index_rows = []
    for kind in ("d2", "d3"):
        for rec in iter_cache(kind):
            meta = index.get(rec["cache_file"], {})
            f = rec.get("fields") or {}
            sd, sd_source = sign_date_of(f, rec.get("text"))
            index_rows.append(
                {
                    "company_shortname": meta.get("company_shortname"),
                    "company_id": meta.get("company_id"),
                    "form_type": kind.upper(),
                    "file": rec.get("file"),
                    "filename_date": rec.get("form_date"),
                    "signature_date": sd,
                    "signature_date_source": sd_source,
                    "signature_date_raw": f.get("signature_date_raw"),
                    "signer_name": f.get("signer_name"),
                    "signer_title": f.get("signer_title"),
                    "semester": f.get("semester"),
                    "text_source": rec.get("text_source"),
                    "portal_form_name": meta.get("portal_form_name"),
                    "error": rec.get("error"),
                }
            )
    index_rows.sort(key=lambda r: (str(r.get("company_shortname") or ""), r["form_type"], r["file"] or ""))
    with (OUT_DIR / "forms_index.csv").open("w", newline="", encoding="utf-8") as fh:
        wr = csv.DictWriter(fh, fieldnames=list(index_rows[0].keys()) if index_rows else ["file"])
        wr.writeheader()
        wr.writerows(index_rows)
    log(f"wrote {OUT_DIR / 'forms_index.csv'} ({len(index_rows)} rows)")

    # accepted students (D3 primary, D2 accepted lists as secondary)
    rows = []
    for kind in ("d3", "d2"):
        for rec in iter_cache(kind):
            f = rec.get("fields") or {}
            students = f.get("students") if kind == "d3" else f.get("accepted_students")
            for s in students or []:
                rows.append(
                    {
                        "company_id": index.get(rec["cache_file"], {}).get("company_id"),
                        "company_shortname": index.get(rec["cache_file"], {}).get("company_shortname"),
                        "form_type": kind.upper(),
                        "form_date": rec.get("form_date"),
                        "source_file": rec.get("file"),
                        "no": s.get("no"),
                        "student_id": norm_mssv(s.get("student_id")) or s.get("student_id"),
                        "ho_lot": (s.get("ho_lot") or "").strip(),
                        "ten": (s.get("ten") or "").strip(),
                        "note": (s.get("note") or "").strip(),
                    }
                )
    order = {"D3": 0, "D2": 1}
    rows.sort(key=lambda r: (str(r.get("company_shortname") or ""), order[r["form_type"]], r["form_date"] or ""))
    with (OUT_DIR / "accepted_students.csv").open("w", newline="", encoding="utf-8") as fh:
        wr = csv.DictWriter(fh, fieldnames=list(rows[0].keys()) if rows else ["company_id"])
        wr.writeheader()
        wr.writerows(rows)
    log(f"wrote {OUT_DIR / 'accepted_students.csv'} ({len(rows)} rows)")

    # projects
    prows = []
    for rec in iter_cache("d2"):
        f = rec.get("fields") or {}
        for p in f.get("projects") or []:
            prows.append(
                {
                    "company_shortname": index.get(rec["cache_file"], {}).get("company_shortname"),
                    "source_file": rec.get("file"),
                    "form_date": rec.get("form_date"),
                    "no": p.get("no"),
                    "title": p.get("title"),
                    "languages": "; ".join(p.get("languages") or []),
                    "frameworks": "; ".join(p.get("frameworks") or []),
                    "weeks": len(p.get("timeline") or []),
                }
            )
    prows.sort(key=lambda r: (str(r.get("company_shortname") or ""), r.get("no") or 0))
    with (OUT_DIR / "projects.csv").open("w", newline="", encoding="utf-8") as fh:
        wr = csv.DictWriter(fh, fieldnames=list(prows[0].keys()) if prows else ["company_shortname"])
        wr.writeheader()
        wr.writerows(prows)
    log(f"wrote {OUT_DIR / 'projects.csv'} ({len(prows)} rows)")


def report() -> None:
    for kind in ("d2", "d3"):
        total = len(collect(kind))
        cached = structured = failed = no_text = 0
        sources: dict[str, int] = {}
        issues: dict[str, int] = {}
        for rec in iter_cache(kind):
            cached += 1
            if rec.get("fields"):
                structured += 1
            if rec.get("error"):
                failed += 1
            if "no text" in (rec.get("issues") or []):
                no_text += 1
            sources[rec.get("text_source", "?")] = sources.get(rec.get("text_source", "?"), 0) + 1
            for i in rec.get("issues") or []:
                key = re.sub(r":.*", "", i)
                issues[key] = issues.get(key, 0) + 1
        print(f"[{kind}] files={total} cached={cached} structured={structured} failed={failed} no_text={no_text}")
        print(f"       sources={sources}")
        if issues:
            print(f"       issues={dict(sorted(issues.items(), key=lambda kv: -kv[1]))}")


def verify() -> None:
    """Cross-checks across the extracted records; prints and writes qa_report.json."""
    index = portal_index()
    problems: dict[str, list] = {
        "missing_signature_date": [],
        "signature_date_out_of_range": [],
        "signature_date_differs_from_filename": [],
        "semester_mismatch": [],
        "students_only_in_d3": [],
        "students_only_in_d2": [],
        "duplicate_student": [],
        "structuring_failed": [],
        "low_confidence_scan": [],
    }
    d2_by_company: dict[str, set] = {}
    d3_by_company: dict[str, set] = {}
    date_range = ("2025-11-01", "2026-12-31")

    for kind in ("d2", "d3"):
        for rec in iter_cache(kind):
            meta = index.get(rec["cache_file"], {})
            company = meta.get("company_shortname") or rec.get("file")
            f = rec.get("fields") or {}
            sd, _sd_source = sign_date_of(f, rec.get("text"))
            fd = rec.get("form_date")
            tag = f"{kind}:{rec.get('file')}"
            if not f:
                problems["structuring_failed"].append(tag)
                continue
            if not sd:
                problems["missing_signature_date"].append(tag)
            else:
                if not (date_range[0] <= sd <= date_range[1]):
                    problems["signature_date_out_of_range"].append(f"{tag} -> {sd}")
                if fd and sd != fd:
                    problems["signature_date_differs_from_filename"].append(f"{tag} sign={sd} filename={fd}")
            sem = f.get("semester")
            if sem and "253" not in str(sem) and "HK252" in str(sem):
                problems["semester_mismatch"].append(f"{tag} -> {sem}")
            if f.get("quality_notes"):
                problems["low_confidence_scan"].append(f"{tag} -> {str(f['quality_notes'])[:140]}")
            students = f.get("students") if kind == "d3" else f.get("accepted_students")
            ids = set()
            for s in students or []:
                sid = norm_mssv(s.get("student_id"))
                if not sid:
                    continue
                if sid in ids:
                    problems["duplicate_student"].append(f"{tag} -> {sid}")
                ids.add(sid)
            if ids:
                (d3_by_company if kind == "d3" else d2_by_company).setdefault(company, set()).update(ids)

    for company in sorted(set(d2_by_company) | set(d3_by_company)):
        d2s = d2_by_company.get(company, set())
        d3s = d3_by_company.get(company, set())
        if d2s and d3s:
            for sid in sorted(d3s - d2s):
                problems["students_only_in_d3"].append(f"{company} -> {sid}")
            for sid in sorted(d2s - d3s):
                problems["students_only_in_d2"].append(f"{company} -> {sid}")

    summary = {
        "counts": {k: len(v) for k, v in problems.items()},
        "signature_dates": {
            "d2_total": len(list(iter_cache("d2"))),
            "d2_with_date": sum(1 for r in iter_cache("d2") if sign_date_of(r.get("fields") or {}, r.get("text"))[0]),
            "d3_total": len(list(iter_cache("d3"))),
            "d3_with_date": sum(1 for r in iter_cache("d3") if sign_date_of(r.get("fields") or {}, r.get("text"))[0]),
        },
        "problems": problems,
    }
    out = OUT_DIR / "qa_report.json"
    out.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(summary["counts"], ensure_ascii=False, indent=2))
    print("signature dates:", json.dumps(summary["signature_dates"], ensure_ascii=False))
    print(f"wrote {out}")


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--kind", choices=["d2", "d3", "both"], default="both")
    ap.add_argument("--stage", choices=["all", "acquire", "structure"], default="all")
    ap.add_argument("--workers", type=int, default=16)
    ap.add_argument("--limit", type=int, default=0, help="only first N files per kind (smoke test)")
    ap.add_argument("--only", type=str, default=None, help="substring filter on file name")
    ap.add_argument("--aggregate", action="store_true", help="rebuild outputs from cache and exit")
    ap.add_argument("--report", action="store_true", help="print coverage summary from cache and exit")
    ap.add_argument("--verify", action="store_true", help="cross-check extracted records, write qa_report.json")
    ap.add_argument("--text-model", type=str, default=None, help="force a text model for structuring")
    ap.add_argument("--vision-model", type=str, default=None, help="override the OCR/vision model")
    ap.add_argument("--vision-base-url", type=str, default=None, help="OCR provider endpoint (default: ai-box)")
    ap.add_argument("--vision-api-key", type=str, default=None, help="OCR provider key (default: AIBOX_API_KEY)")
    ap.add_argument("--base-url", type=str, default=None, help="OpenAI-compatible endpoint (default: OpenCode Go)")
    ap.add_argument("--api-key", type=str, default=None, help="override OPENCODE_GO_API_KEY (e.g. ai-box key)")
    ap.add_argument("--force-text", action="store_true", help="redo text acquisition (ignores cache)")
    ap.add_argument("--force-fields", action="store_true", help="redo structuring (ignores cache)")
    ap.add_argument("--retry-failed", action="store_true", help="only files missing text or fields / with errors")
    ap.add_argument("--only-ocr", action="store_true", help="only files that need the vision model (scans/photos)")
    args = ap.parse_args()

    if args.aggregate:
        aggregate()
        return
    if args.report:
        report()
        return
    if args.verify:
        verify()
        return
    run(args)
    aggregate()


if __name__ == "__main__":
    main()