verify_qisi_corpus.py

# -*- coding: utf-8 -*-
"""Verify a qisi corpus output directory."""

from __future__ import annotations

import argparse
import json
import sqlite3
from pathlib import Path
from typing import Any


def load_json(path: Path) -> Any:
    return json.loads(path.read_text(encoding="utf-8"))


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    rows = []
    with path.open("r", encoding="utf-8") as f:
        for line_no, line in enumerate(f, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError as exc:
                raise ValueError(f"{path}:{line_no}: invalid JSONL: {exc}") from exc
    return rows


def sqlite_count(path: Path) -> int:
    conn = sqlite3.connect(path)
    try:
        return int(conn.execute("SELECT COUNT(*) FROM statuses").fetchone()[0])
    finally:
        conn.close()


def verify_gemini(dir_path: Path) -> tuple[list[str], list[str]]:
    errors = []
    warnings = []
    gemini_dir = dir_path / "gemini"
    manifest_path = gemini_dir / "manifest.json"
    if not gemini_dir.exists():
        warnings.append("missing gemini/ directory; run export_for_gemini.py")
        return errors, warnings
    if not manifest_path.exists():
        errors.append("missing gemini/manifest.json")
        return errors, warnings
    manifest = load_json(manifest_path)
    for item in manifest:
        chunk = gemini_dir / item["file"]
        if not chunk.exists():
            errors.append(f"missing Gemini chunk: {chunk}")
        elif chunk.stat().st_size == 0:
            errors.append(f"empty Gemini chunk: {chunk}")
    if not manifest:
        warnings.append("Gemini manifest is empty")
    return errors, warnings


def build_report(dir_path: Path, require_complete: bool) -> dict[str, Any]:
    errors = []
    warnings = []

    summary_path = dir_path / "summary.json"
    jsonl_path = dir_path / "timeline_clean.jsonl"
    markdown_path = dir_path / "timeline_clean.md"
    sqlite_path = dir_path / "timeline.sqlite"
    raw_pages_path = dir_path / "timeline_raw_pages.json"

    for path in [summary_path, jsonl_path, markdown_path, sqlite_path, raw_pages_path]:
        if not path.exists():
            errors.append(f"missing required file: {path.name}")

    summary = load_json(summary_path) if summary_path.exists() else {}
    rows = load_jsonl(jsonl_path) if jsonl_path.exists() else []
    raw_pages = load_json(raw_pages_path) if raw_pages_path.exists() else []
    db_count = sqlite_count(sqlite_path) if sqlite_path.exists() else None

    ids = [row.get("id") for row in rows]
    unique_ids = {item for item in ids if item is not None}
    if len(unique_ids) != len(ids):
        errors.append("timeline_clean.jsonl contains duplicate or missing ids")

    if summary:
        if summary.get("uid") != "9199129225":
            errors.append(f"unexpected uid in summary: {summary.get('uid')}")
        if summary.get("status_count") != len(rows):
            errors.append(
                f"summary status_count={summary.get('status_count')} but JSONL rows={len(rows)}"
            )
        if summary.get("pages_saved") != len(raw_pages):
            errors.append(
                f"summary pages_saved={summary.get('pages_saved')} but raw pages={len(raw_pages)}"
            )
        advertised_total = summary.get("advertised_total")
        if advertised_total and advertised_total != len(rows):
            warnings.append(
                f"advertised_total={advertised_total} differs from accessible JSONL rows={len(rows)}"
            )
        if require_complete and not summary.get("complete"):
            errors.append(
                "corpus is not complete according to summary.json "
                f"(stop_reason={summary.get('stop_reason')}, failed_page={summary.get('failed_page')})"
            )
        if not summary.get("complete"):
            warnings.append(
                "summary complete=false; this is expected for anonymous/test runs but not for final corpus"
            )

    if db_count is not None and db_count != len(rows):
        errors.append(f"SQLite statuses count={db_count} but JSONL rows={len(rows)}")

    column_count = sum(1 for row in rows if row.get("is_column"))
    column_fulltext_count = sum(1 for row in rows if row.get("is_column") and row.get("full_raw"))
    if column_count and column_fulltext_count < column_count:
        warnings.append(
            f"only {column_fulltext_count}/{column_count} column articles have full_raw text"
        )

    gemini_errors, gemini_warnings = verify_gemini(dir_path)
    errors.extend(gemini_errors)
    warnings.extend(gemini_warnings)

    report = {
        "output_dir": str(dir_path.resolve()),
        "ok": not errors,
        "complete": bool(summary.get("complete")),
        "status_count": len(rows),
        "sqlite_count": db_count,
        "raw_pages_count": len(raw_pages),
        "column_count": column_count,
        "column_fulltext_count": column_fulltext_count,
        "advertised_total": summary.get("advertised_total"),
        "advertised_max_page": summary.get("advertised_max_page"),
        "requested_max_page": summary.get("requested_max_page"),
        "stop_reason": summary.get("stop_reason"),
        "failed_page": summary.get("failed_page"),
        "errors": errors,
        "warnings": warnings,
    }
    return report


def write_markdown_report(path: Path, report: dict[str, Any]) -> None:
    lines = [
        "# Qisi Corpus Verification",
        "",
        f"- ok: {report['ok']}",
        f"- complete: {report['complete']}",
        f"- status_count: {report['status_count']}",
        f"- sqlite_count: {report['sqlite_count']}",
        f"- raw_pages_count: {report['raw_pages_count']}",
        f"- column_fulltext: {report['column_fulltext_count']}/{report['column_count']}",
        f"- advertised_total: {report['advertised_total']}",
        f"- advertised_max_page: {report['advertised_max_page']}",
        f"- requested_max_page: {report['requested_max_page']}",
        f"- stop_reason: {report['stop_reason']}",
        f"- failed_page: {report['failed_page']}",
        "",
        "## Errors",
        "",
    ]
    lines.extend([f"- {item}" for item in report["errors"]] or ["- none"])
    lines.extend(["", "## Warnings", ""])
    lines.extend([f"- {item}" for item in report["warnings"]] or ["- none"])
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Verify qisi corpus output.")
    parser.add_argument("output_dir")
    parser.add_argument("--require-complete", action="store_true")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    output_dir = Path(args.output_dir)
    report = build_report(output_dir, args.require_complete)
    (output_dir / "verification_report.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    write_markdown_report(output_dir / "verification_report.md", report)
    print(json.dumps(report, ensure_ascii=False, indent=2))
    return 0 if report["ok"] else 2


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