"""Preprocess assigned source files; final deduplication and splits happen at merge."""

import argparse
import json
import sys
import unicodedata
from collections import Counter
from contextlib import ExitStack, closing
from itertools import islice
from pathlib import Path

ROOT = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(ROOT / "src"))

from vimeml.data.corpus_parts import (AUDITS, PART_SCHEMA_VERSION, load_plan, parse_part,
                                    create_database, process_content, aggregate_metrics)
from vimeml.data.build import sha, file_sha, display_path, dump_json
from vimeml.data.readers import iter_fineweb, iter_tatoeba


def run(config_path, part, output=None, list_only=False):
    index, total = parse_part(part)
    config, inventory, policies, _, policy_hashes, code_hashes, signature = load_plan(config_path)
    selected = [item for number, item in enumerate(inventory) if number % total == index]
    if not selected:
        raise ValueError("This part has no inputs; use fewer parts.")
    if list_only:
        print(json.dumps({"part": part, "files": [p for p, _, _ in selected]}, ensure_ascii=False, indent=2))
        return
    paths = [ROOT / p for p, _, _ in selected]
    for path in paths:
        if not path.is_file():
            raise FileNotFoundError(f"Missing assigned input: {path}")
    snapshots = {p: (p.stat().st_size, p.stat().st_mtime_ns) for p in paths}
    settings = config["build"]
    output = (ROOT / (output or f"{settings['output_dir']}-part-{index:03d}-of-{total:03d}")).resolve()
    selected_documents = Counter({p: 0 for p, _, _ in selected})
    raw_characters = 0
    with ExitStack() as stack:
        db, annotation_count = create_database(output, policies.get("annotations"))
        stack.callback(db.close)
        audits = {name: stack.enter_context((output / f"{name}.jsonl").open("w", encoding="utf-8", newline="\n")) for name in AUDITS}
        for relative, source, limit in selected:
            path = ROOT / relative
            reader = iter_fineweb(path, settings["batch_size"]) if source == "fineweb" else iter_tatoeba(path)
            with closing(reader):
                for document in islice(reader, limit) if limit else reader:
                    selected_documents[relative] += 1
                    raw_characters += len(document.text)
                    normalized = unicodedata.normalize("NFC", document.text.replace("\r\n", "\n").replace("\r", "\n")).lstrip("\ufeff")
                    doc_hash = sha(normalized)
                    first = db.execute("INSERT OR IGNORE INTO contents(doc_hash) VALUES (?)", (doc_hash,)).rowcount == 1
                    identity = {"source": source, "doc_id": document.doc_id, "source_id": document.source_id,
                                "source_file": document.source_file, "row_index": document.row_index,
                                "doc_hash": doc_hash, "metadata": document.metadata}
                    db.execute("INSERT INTO origins VALUES (?,?,?,?,?,?)", (doc_hash, document.doc_id, source,
                               document.source_file, document.row_index, json.dumps(identity, ensure_ascii=False, allow_nan=False)))
                    if first:
                        metrics = process_content(db, document, normalized, doc_hash, settings, audits)
                        db.execute("INSERT INTO document_metrics VALUES (?,?)", (doc_hash, json.dumps(metrics, ensure_ascii=False)))
                    count = sum(selected_documents.values())
                    if count % settings.get("commit_every", 5000) == 0:
                        db.commit()
                    if count % settings.get("progress_every", 5000) == 0:
                        print(f"Part {part}: read {count} document origins", flush=True)
        db.commit()
        unmatched_here = [row[0] for row in db.execute("SELECT a.id FROM annotations a JOIN contents c ON c.doc_hash=a.doc_hash WHERE a.matched=0")]
        if config.get("review", {}).get("require_all_annotations", False) and unmatched_here:
            raise ValueError(f"Annotations in assigned documents did not match: {unmatched_here}")
        metrics = aggregate_metrics(db)
        stats = {**metrics, "document_origins": sum(selected_documents.values()), "raw_characters": raw_characters,
                 "annotation_records": annotation_count,
                 "matched_annotations": db.execute("SELECT COUNT(*) FROM annotations WHERE matched=1").fetchone()[0]}
        dump_json(output / "part-stats.json", stats)
    for path, before in snapshots.items():
        if (path.stat().st_size, path.stat().st_mtime_ns) != before:
            raise ValueError(f"Input changed during processing: {path}")
    if load_plan(config_path)[-1] != signature:
        raise ValueError("Code or policy changed during processing; part is incomplete.")
    print(f"Part {part}: hashing completed files", flush=True)
    artifact_hashes = {name: file_sha(output / name) for name in ("index.sqlite", "part-stats.json", *(f"{name}.jsonl" for name in AUDITS))}
    dump_json(output / "manifest.json", {
        "status": "complete", "stage": "preprocessed_part", "part_schema_version": PART_SCHEMA_VERSION,
        "partition": {"index": index, "count": total}, "pipeline_signature": signature,
        "config": config, "config_sha256": file_sha(config_path), "python_version": sys.version,
        "planned_inputs": [p for p, _, _ in inventory], "selected_documents": dict(selected_documents),
        "input_sha256": {display_path(p): file_sha(p) for p in paths},
        "policy_sha256": policy_hashes, "code_sha256": code_hashes, "artifact_sha256": artifact_hashes,
        "fingerprint_normalization": "Code and policy: UTF-8 text with LF line endings, no BOM; input/artifact hashes: exact bytes.",
        "ready_for_lm_training": False,
    })
    print(json.dumps({"part": part, "document_origins": stats["document_origins"], "output": str(output)}, ensure_ascii=False))


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, default=ROOT / "configs/corpus-sharded.toml")
    parser.add_argument("--part", required=True, help="Zero-based INDEX/COUNT; 0/2 and 1/2 for two machines.")
    parser.add_argument("--output")
    parser.add_argument("--list-only", action="store_true", help="Show assigned files; no raw data needed.")
    args = parser.parse_args()
    run(args.config.resolve(), args.part, args.output, args.list_only)


if __name__ == "__main__":
    main()