#!/usr/bin/env python3
"""Audit saved VJudge MHTML files without copying problem statements.

The script reports the metadata needed by the course and verifies that each
snapshot contains an embedded statement, rather than only the VJudge shell.
It intentionally stores no statement text.
"""

from __future__ import annotations

import argparse
import html
import json
import re
import sys
from email import policy
from email.parser import BytesParser
from html.parser import HTMLParser
from pathlib import Path


class TextExtractor(HTMLParser):
    def __init__(self) -> None:
        super().__init__()
        self.parts: list[str] = []

    def handle_data(self, data: str) -> None:
        if data.strip():
            self.parts.append(data.strip())

    def text(self) -> str:
        return " ".join(self.parts)


def plain_text(fragment: str) -> str:
    parser = TextExtractor()
    parser.feed(fragment)
    return html.unescape(parser.text())


def decode_part(part) -> str:
    payload = part.get_payload(decode=True) or b""
    charset = part.get_content_charset() or "utf-8"
    return payload.decode(charset, errors="replace")


def parse_snapshot(path: Path) -> dict[str, object]:
    message = BytesParser(policy=policy.default).parsebytes(path.read_bytes())
    snapshot_url = str(message.get("Snapshot-Content-Location", ""))
    subject = str(message.get("Subject", path.stem))

    html_parts = [
        part
        for part in message.walk()
        if not part.is_multipart() and part.get_content_type() == "text/html"
    ]
    main_part = next(
        (
            part
            for part in html_parts
            if str(part.get("Content-Location", "")) == snapshot_url
        ),
        None,
    )
    if main_part is None:
        raise ValueError("missing VJudge main HTML part")

    main_html = decode_part(main_part)
    match = re.search(
        r'<textarea[^>]*name=["\']dataJson["\'][^>]*>(.*?)</textarea>',
        main_html,
        flags=re.DOTALL,
    )
    if match is None:
        raise ValueError("missing dataJson metadata")
    metadata = json.loads(html.unescape(match.group(1)))

    official_keys = {
        str(item["key"])
        for item in metadata.get("descBriefs", [])
        if item.get("mainOfficial") or (item.get("official") and item.get("type") == "main")
    }
    statement_candidates: list[tuple[bool, str]] = []
    for part in html_parts:
        location = str(part.get("Content-Location", ""))
        if "/problem/description/" not in location:
            continue
        body = decode_part(part)
        key_match = re.search(r"/problem/description/(\d+)", location)
        is_official = bool(key_match and key_match.group(1) in official_keys)
        statement_candidates.append((is_official, body))

    if not statement_candidates:
        raise ValueError("missing embedded statement HTML")
    statement_candidates.sort(key=lambda item: (item[0], len(item[1])), reverse=True)
    statement_html = statement_candidates[0][1]
    statement_text = plain_text(statement_html)

    properties = {
        str(item.get("title", "")): plain_text(str(item.get("content", "")))
        for item in metadata.get("properties", [])
    }
    title = re.sub(r"\s+-\s+Virtual Judge$", "", subject)
    title = re.sub(r"\s+-\s+[^-]+\s+[^-]+$", "", title)
    has_input = bool(re.search(r"\bInput\b|输入", statement_text, re.IGNORECASE))
    has_output = bool(re.search(r"\bOutput\b|输出", statement_text, re.IGNORECASE))
    sample_count = len(re.findall(r'class=["\']vjudge_sample["\']', statement_html))

    return {
        "file": path.as_posix(),
        "query": f"{metadata.get('oj', '?')}-{metadata.get('prob', '?')}",
        "title": title,
        "source": properties.get("source", ""),
        "difficulty": properties.get("difficulty", ""),
        "time_limit": properties.get("time_limit", ""),
        "memory_limit": properties.get("mem_limit", ""),
        "statement_chars": len(statement_text),
        "samples": sample_count,
        "has_input": has_input,
        "has_output": has_output,
        "official_statement": statement_candidates[0][0],
        "url": snapshot_url,
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("root", nargs="?", default="problems")
    parser.add_argument("--json", action="store_true", dest="as_json")
    args = parser.parse_args()

    paths = sorted(Path(args.root).rglob("*.mhtml"))
    if not paths:
        print(f"no .mhtml files found under {args.root}", file=sys.stderr)
        return 2

    rows: list[dict[str, object]] = []
    errors: list[str] = []
    for path in paths:
        try:
            rows.append(parse_snapshot(path))
        except Exception as exc:  # keep auditing the remaining snapshots
            errors.append(f"{path}: {exc}")

    if args.as_json:
        print(json.dumps(rows, ensure_ascii=False, indent=2))
    else:
        print("query\ttitle\tdifficulty\tsource\tstatement\tsamples\tofficial")
        for row in rows:
            statement = "ok" if row["has_input"] and row["has_output"] else "incomplete"
            official = "yes" if row["official_statement"] else "fallback"
            print(
                f"{row['query']}\t{row['title']}\t{row['difficulty'] or '-'}\t"
                f"{row['source'] or '-'}\t{statement}\t{row['samples']}\t{official}"
            )

    for error in errors:
        print(error, file=sys.stderr)
    incomplete = [
        row for row in rows if not row["has_input"] or not row["has_output"] or row["statement_chars"] < 200
    ]
    return 1 if errors or incomplete else 0


if __name__ == "__main__":
    raise SystemExit(main())