#!/usr/bin/env python3

import argparse
import csv
import gzip
import hashlib
import re
import shutil
import sys
import time
import xml.etree.ElementTree as ET
from pathlib import Path

import duckdb
import pyarrow as pa
import pyarrow.parquet as pq

LANGS = ["cs", "hr", "pl", "sk", "sl"]
TEXT_EXTS = ("en", *LANGS)

FILENAME_RE = re.compile(
    r"^(?P<corpus>.+)\.(?P<langpair>[a-z]{2}-[a-z]{2})\.(?P<ext>en|cs|hr|pl|sk|sl)$"
)

# md5 of the normalised English text. Normalisation is surface-level only:
# trim, collapse internal whitespace runs, casefold. No punctuation or
# semantic folding -- those would merge genuinely distinct strings.
NORMALISE = r"lower(regexp_replace(trim(en_text), '\s+', ' ', 'g'))"
MERGE_KEY = f"md5({NORMALISE})"


# ---------------------------------------------------------------------------
# stage 1: extract
# ---------------------------------------------------------------------------

def read_lines(path: Path) -> list[str]:
    with open(path, encoding="utf-8") as fh:
        return [line.rstrip("\n") for line in fh]


def drop_separator_pairs(en: list[str], tgt: list[str]) -> tuple[list[str], list[str]]:
    """
    Remove line pairs blank on BOTH sides.

    Some exports insert a blank line into both files at once as a document
    or domain separator -- CCAligned.cs-en has 255 of them, exactly matching
    its 255 <linkGrp> elements. They are not links and must go before line N
    can equal link N.

    Only *simultaneously* blank pairs are dropped. A real link with a missing
    translation is blank on one side only, so genuine data cannot be caught
    by this.
    """
    keep_en, keep_tgt = [], []
    for e, t in zip(en, tgt):
        if e == "" and t == "":
            continue
        keep_en.append(e)
        keep_tgt.append(t)
    return keep_en, keep_tgt


def _side_is_english(doc_path: str) -> bool:
    """OPUS doc paths are language-prefixed: 'en/foo.xml.gz', 'cs/bar/baz.xml.gz'."""
    return doc_path.split("/", 1)[0] == "en"


def parse_ids_file(path: Path) -> list[tuple[int, int]]:
    """
    Tab-separated: doc_a, doc_b, ids_a, ids_b -- ids space-separated within
    a column. Returns (n_english_ids, n_target_ids) per link.
    """
    out = []
    with open(path, encoding="utf-8") as fh:
        for line in fh:
            parts = line.rstrip("\n").split("\t")
            if len(parts) < 4:
                parts = re.split(r"\s{2,}", line.rstrip("\n"))
            if len(parts) < 4:
                out.append((0, 0))  # unparseable -> never counts as clean
                continue
            doc_a, _doc_b, ids_a, ids_b = parts[0], parts[1], parts[2], parts[3]
            n_a, n_b = len(ids_a.split()), len(ids_b.split())
            out.append((n_a, n_b) if _side_is_english(doc_a) else (n_b, n_a))
    return out


def parse_xml_file(path: Path) -> list[tuple[int, int]]:
    """
    XCES alignment. Streamed with iterparse and cleared as we go -- some of
    these files are 500MB+ and a full DOM is not affordable.

    Multiple <linkGrp> elements are concatenated in document order, matching
    how the flat text file numbers its lines globally.
    """
    out = []
    english_left = True
    for event, elem in ET.iterparse(str(path), events=("start", "end")):
        if event == "start" and elem.tag == "linkGrp":
            english_left = _side_is_english(elem.get("fromDoc", ""))
        elif event == "end":
            if elem.tag == "link":
                sides = elem.get("xtargets", "").split(";")
                if len(sides) != 2:
                    out.append((0, 0))
                else:
                    a, b = len(sides[0].split()), len(sides[1].split())
                    out.append((a, b) if english_left else (b, a))
                elem.clear()
            elif elem.tag == "linkGrp":
                elem.clear()
    return out


def extract_one(root: Path, corpus: str, langpair: str) -> pa.Table:
    a, b = langpair.split("-")
    lang = b if a == "en" else a

    en_path = root / f"{corpus}.{langpair}.en"
    tgt_path = root / f"{corpus}.{langpair}.{lang}"
    xml_path = root / f"{corpus}.{langpair}.xml"
    ids_path = root / f"{corpus}.{langpair}.ids"

    en_raw, tgt_raw = read_lines(en_path), read_lines(tgt_path)
    if len(en_raw) != len(tgt_raw):
        raise ValueError(
            f"text files disagree: en={len(en_raw)} {lang}={len(tgt_raw)}"
        )

    en_lines, tgt_lines = drop_separator_pairs(en_raw, tgt_raw)
    n = len(en_lines)

    if xml_path.exists():
        counts, align_source = parse_xml_file(xml_path), "xml"
    elif ids_path.exists():
        counts, align_source = parse_ids_file(ids_path), "ids"
    else:
        counts, align_source = [(1, 1)] * n, "none"

    if len(counts) != n:
        raise ValueError(
            f"{align_source} has {len(counts)} links but {n} usable text lines "
            "-- refusing to pair them positionally"
        )

    # A link whose English side is empty has no reachable key: drop it.
    # A link with English present but the translation empty is kept -- the
    # English key survives and that language is simply NULL.
    keep = [i for i in range(n) if en_lines[i].strip() != ""]

    return pa.table({
        "corpus": pa.array([corpus] * len(keep)),
        "langpair": pa.array([langpair] * len(keep)),
        "lang": pa.array([lang] * len(keep)),
        "line_no": pa.array([i + 1 for i in keep], type=pa.int64()),
        "en_text": pa.array([en_lines[i] for i in keep]),
        "lang_text": pa.array(
            [tgt_lines[i] if tgt_lines[i].strip() != "" else None for i in keep]
        ),
        "en_n_ids": pa.array([counts[i][0] for i in keep], type=pa.int32()),
        "lang_n_ids": pa.array([counts[i][1] for i in keep], type=pa.int32()),
        "clean": pa.array([counts[i] == (1, 1) for i in keep]),
        "align_source": pa.array([align_source] * len(keep)),
    })


def discover(root: Path, exclude: set[str]) -> list[tuple[str, str]]:
    found: dict[tuple[str, str], set[str]] = {}
    for path in root.iterdir():
        if not path.is_file():
            continue
        m = FILENAME_RE.match(path.name)
        if not m:
            continue
        corpus, langpair, ext = m["corpus"], m["langpair"], m["ext"]
        if corpus in exclude:
            continue
        if "en" not in langpair.split("-"):
            continue
        found.setdefault((corpus, langpair), set()).add(ext)

    combos = []
    for (corpus, langpair), exts in found.items():
        a, b = langpair.split("-")
        target = b if a == "en" else a
        if "en" in exts and target in exts:
            combos.append((corpus, langpair))
    return sorted(combos)


def stage_extract(root: Path, out_dir: Path, exclude: set[str], report_path: Path) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    combos = discover(root, exclude)
    print(f"[extract] {len(combos)} corpus x langpair combos "
          f"({len(exclude)} corpora excluded)")

    rows = []
    for i, (corpus, langpair) in enumerate(combos, 1):
        dest = out_dir / f"{corpus}.{langpair}.parquet"
        if dest.exists():
            print(f"[extract] {i}/{len(combos)} {corpus}.{langpair} -- already done")
            continue
        t0 = time.time()
        tmp = dest.with_suffix(".parquet.partial")
        try:
            table = extract_one(root, corpus, langpair)
            pq.write_table(table, tmp)
            tmp.replace(dest)  # atomic: a partial file never looks complete
            n = table.num_rows
            n_clean = sum(table.column("clean").to_pylist())
            src = table.column("align_source")[0].as_py() if n else ""
            pct = round(100 * n_clean / n, 2) if n else 0.0
            print(f"[extract] {i}/{len(combos)} {corpus}.{langpair} -- "
                  f"{n:,} rows, {pct}% clean, {src}, {time.time()-t0:.0f}s")
            rows.append(dict(corpus=corpus, langpair=langpair, status="ok",
                             rows=n, clean_pct=pct, align_source=src, error=""))
        except Exception as exc:
            tmp.unlink(missing_ok=True)
            print(f"[extract] {i}/{len(combos)} {corpus}.{langpair} -- FAILED: {exc}")
            rows.append(dict(corpus=corpus, langpair=langpair, status="failed",
                             rows=0, clean_pct=0.0, align_source="", error=str(exc)))

    if rows:
        write_header = not report_path.exists()
        with open(report_path, "a", newline="", encoding="utf-8") as fh:
            w = csv.DictWriter(fh, fieldnames=["corpus", "langpair", "status", "rows",
                                               "clean_pct", "align_source", "error"])
            if write_header:
                w.writeheader()
            w.writerows(rows)
        n_fail = sum(1 for r in rows if r["status"] == "failed")
        print(f"[extract] {len(rows)-n_fail} ok, {n_fail} failed -> {report_path}")


# ---------------------------------------------------------------------------
# stage 2: bucket
# ---------------------------------------------------------------------------

def stage_bucket(con, out_dir: Path, keyed_dir: Path, n_buckets: int) -> None:
    """
    One streaming pass. Every row gets its merge_key and a bucket number.
    Rows sharing a merge_key always land in the same bucket, so each bucket
    can later be merged in isolation with no cross-bucket reconciliation.

    This is a projection plus a write -- no aggregation, so it does not build
    the large hash table that made the single-query merge run out of memory.
    """
    print(f"[bucket] partitioning into {n_buckets} buckets -> {keyed_dir}/")
    t0 = time.time()
    staging = keyed_dir.with_name(keyed_dir.name + ".partial")
    if staging.exists():
        shutil.rmtree(staging)
    con.execute(f"""
        COPY (
            SELECT
                {MERGE_KEY} AS merge_key,
                en_text, lang, lang_text, clean,
                en_n_ids, lang_n_ids, corpus, langpair,
                abs(hash({MERGE_KEY})) % {n_buckets} AS bucket
            FROM read_parquet('{out_dir}/*.parquet')
        ) TO '{staging}' (FORMAT PARQUET, PARTITION_BY (bucket), OVERWRITE_OR_IGNORE true)
    """)
    staging.replace(keyed_dir)
    print(f"[bucket] done in {time.time()-t0:.0f}s")


# ---------------------------------------------------------------------------
# stage 3: merge
# ---------------------------------------------------------------------------

def per_language_columns() -> str:
    """
    For each language: pick ONE source row and take its text, flag, grouping
    counts and corpus from that same row.

    This deliberately avoids separate any_value() calls per field. Those are
    independent aggregates -- given several candidate rows for one key and
    language (exactly what happens with common boilerplate), each call may
    pick a different row, so the clean flag could describe a translation that
    was not the one kept. Packing the fields into a struct and unpacking after
    aggregation keeps them from a single row.

    max() on a struct compares field-by-field in declaration order, so putting
    is_clean first makes a strict 1:1 link win over a grouped one, with the
    text as a deterministic tiebreak. That is a policy choice, not a neutral
    one: among equally valid candidates it prefers the cleanly-aligned link.

    Three counters accompany each language:

      <lang>_n_rows      how many source rows supplied a translation
      <lang>_n_distinct  how many DIFFERENT translations among them
      <lang>_n_corpora    how many distinct corpora contributed

    n_distinct is the one that matters for quality: = 1 means every source
    agreed, so the pick was arbitrary only between identical strings. The gap
    between n_rows and n_distinct measures pure duplication -- on a sample it
    ran ~80%, largely because several corpora in the download appear to be the
    same data under different names. n_corpora separates "one corpus repeated"
    from "many corpora agreeing".
    """
    parts = []
    for lang in LANGS:
        pick = (
            f"max(CASE WHEN lang = '{lang}' THEN struct_pack("
            f"is_clean := clean, txt := lang_text, "
            f"n_en := en_n_ids, n_tgt := lang_n_ids, src := corpus) END)"
        )
        parts.append(f"{pick}.txt       AS {lang}_text")
        parts.append(f"{pick}.is_clean  AS {lang}_clean")
        parts.append(f"{pick}.n_en      AS {lang}_en_n_ids")
        parts.append(f"{pick}.n_tgt     AS {lang}_n_ids")
        parts.append(f"{pick}.src       AS {lang}_corpus")
        parts.append(
            f"count(*) FILTER (WHERE lang = '{lang}' AND lang_text IS NOT NULL) "
            f"AS {lang}_n_rows"
        )
        parts.append(
            f"count(DISTINCT CASE WHEN lang = '{lang}' THEN lang_text END) "
            f"AS {lang}_n_distinct"
        )
        parts.append(
            f"count(DISTINCT CASE WHEN lang = '{lang}' AND lang_text IS NOT NULL "
            f"THEN corpus END) AS {lang}_n_corpora"
        )
    return ",\n            ".join(parts)


def stage_merge(con, keyed_dir: Path, merged_dir: Path, n_buckets: int) -> None:
    merged_dir.mkdir(parents=True, exist_ok=True)
    cols = per_language_columns()
    total = 0
    for i in range(n_buckets):
        src_dir = keyed_dir / f"bucket={i}"
        dest = merged_dir / f"bucket_{i:04d}.parquet"
        if dest.exists():
            total += con.execute(f"SELECT count(*) FROM '{dest}'").fetchone()[0]
            continue
        if not src_dir.exists() or not any(src_dir.glob("*.parquet")):
            print(f"[merge] bucket {i} empty, skipping")
            continue
        t0 = time.time()
        tmp = dest.with_suffix(".parquet.partial")
        con.execute(f"""
            COPY (
                SELECT
                    merge_key,
                    any_value(en_text) AS en_text,
                    {cols}
                FROM read_parquet('{src_dir}/*.parquet')
                GROUP BY merge_key
            ) TO '{tmp}' (FORMAT PARQUET)
        """)
        tmp.replace(dest)
        n = con.execute(f"SELECT count(*) FROM '{dest}'").fetchone()[0]
        total += n
        print(f"[merge] bucket {i}/{n_buckets-1}: {n:,} keys, {time.time()-t0:.0f}s")
    print(f"[merge] {total:,} unique English keys across {n_buckets} buckets")


# ---------------------------------------------------------------------------
# summary
# ---------------------------------------------------------------------------

def summarise(con, merged_dir: Path) -> None:
    glob = f"{merged_dir}/*.parquet"
    total = con.execute(f"SELECT count(*) FROM read_parquet('{glob}')").fetchone()[0]
    print(f"\n{'='*64}\nunique English keys: {total:,}\n{'='*64}")
    print(f"{'lang':<6}{'covered':>16}{'%':>8}{'clean':>16}{'ambiguous':>14}")
    for lang in LANGS:
        cov, clean, ambiguous = con.execute(f"""
            SELECT count(*) FILTER (WHERE {lang}_text IS NOT NULL),
                   count(*) FILTER (WHERE {lang}_clean),
                   count(*) FILTER (WHERE {lang}_n_distinct > 1)
            FROM read_parquet('{glob}')
        """).fetchone()
        pct = 100 * cov / total if total else 0
        print(f"{lang:<6}{cov:>16,}{pct:>7.1f}%{clean:>16,}{ambiguous:>14,}")
    print("\n'ambiguous' = keys where sources supplied genuinely DIFFERENT "
          "translations.\nFilter on <lang>_n_distinct = 1 for the unambiguous "
          "subset, <lang>_clean for strict 1:1 links.\nNote: <lang>_clean is "
          "inflated by design -- the tiebreak prefers clean candidates.")


# ---------------------------------------------------------------------------

def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--root", required=True, type=Path,
                    help="directory holding the OPUS plain-text download")
    ap.add_argument("--work", default=Path("build"), type=Path,
                    help="working directory for all intermediate and final output")
    ap.add_argument("--buckets", type=int, default=128,
                    help="raise if a bucket runs out of memory during merge")
    ap.add_argument("--memory-limit", default="100GB",
                    help="keep comfortably under the container's ceiling")
    ap.add_argument("--exclude", default="",
                    help="comma-separated corpus names to skip entirely, e.g. "
                         "MultiHPLT,MultiCCAligned,MultiParaCrawl,MultiMaCoCu")
    ap.add_argument("--force-stage", default="", choices=["", "bucket", "merge"],
                    help="redo a stage that is already on disk")
    args = ap.parse_args()

    work = args.work
    out_dir = work / "extracted"
    keyed_dir = work / "keyed"
    merged_dir = work / "merged"
    tmp_dir = work / "duckdb_tmp"
    for d in (work, tmp_dir):
        d.mkdir(parents=True, exist_ok=True)

    exclude = {c.strip() for c in args.exclude.split(",") if c.strip()}

    con = duckdb.connect()
    con.execute("PRAGMA enable_progress_bar=true")
    con.execute(f"SET memory_limit='{args.memory_limit}'")
    con.execute(f"SET temp_directory='{tmp_dir}'")
    con.execute("SET preserve_insertion_order=false")

    started = time.time()

    stage_extract(args.root, out_dir, exclude, work / "extract_report.csv")

    if args.force_stage == "bucket" and keyed_dir.exists():
        shutil.rmtree(keyed_dir)
    if keyed_dir.exists():
        print(f"[bucket] {keyed_dir}/ present -- skipping "
              f"(--force-stage bucket to redo)")
    else:
        stage_bucket(con, out_dir, keyed_dir, args.buckets)

    if args.force_stage == "merge" and merged_dir.exists():
        shutil.rmtree(merged_dir)
    stage_merge(con, keyed_dir, merged_dir, args.buckets)

    summarise(con, merged_dir)
    print(f"\ntotal wall time {(time.time()-started)/60:.1f} min")
    print(f"query with: SELECT * FROM read_parquet('{merged_dir}/*.parquet')")


if __name__ == "__main__":
    main()