#!/usr/bin/env python3
"""Prepare a small Persian PP-OCRv6 recognition pilot on a remote GPU."""

from __future__ import annotations

import argparse
import io
import json
import unicodedata
from collections import Counter
from pathlib import Path

import pyarrow.parquet as pq
from huggingface_hub import HfApi, hf_hub_download
from PIL import Image


def clean_text(value: str) -> str:
    value = unicodedata.normalize("NFC", value).replace("\t", " ")
    return " ".join(value.split())


def image_bytes(value: object) -> bytes:
    if isinstance(value, dict) and value.get("bytes") is not None:
        return value["bytes"]
    raise ValueError("Expected an embedded Hugging Face image payload")


def read_gt(path: Path) -> list[tuple[Path, str]]:
    rows = []
    for raw in path.read_text(encoding="utf-8").splitlines():
        key, text = raw.split(maxsplit=1)
        rows.append((path.parent / f"{key}.png", clean_text(text)))
    return rows


def first_parquet(repo: str) -> str:
    files = HfApi().list_repo_files(repo, repo_type="dataset")
    candidates = sorted(x for x in files if x.endswith(".parquet") and x.startswith("data/"))
    if not candidates:
        raise RuntimeError(f"No data parquet found in {repo}")
    return candidates[0]


def extract_printed(repo: str, output: Path, train_count: int, eval_count: int) -> tuple[list, list]:
    source = hf_hub_download(
        repo, first_parquet(repo), repo_type="dataset", local_dir=output / "hf"
    )
    table = pq.read_table(source, columns=["image", "label"]).slice(0, train_count + eval_count)
    rows = table.to_pylist()
    if len(rows) < train_count + eval_count:
        raise RuntimeError("Printed pilot shard does not contain enough rows")

    destination = output / "images" / "printed"
    destination.mkdir(parents=True, exist_ok=True)
    converted = []
    for index, row in enumerate(rows):
        path = destination / f"{index:07d}.jpg"
        with Image.open(io.BytesIO(image_bytes(row["image"]))) as image:
            image.convert("RGB").save(path, "JPEG", quality=95)
        converted.append((path, clean_text(row["label"])))
    return converted[:train_count], converted[train_count:]


def load_hardword_text(repo: str, output: Path) -> list[str]:
    source = hf_hub_download(
        repo, first_parquet(repo), repo_type="dataset", local_dir=output / "hf"
    )
    return [clean_text(x) for x in pq.read_table(source, columns=["sentence"])["sentence"].to_pylist()]


def write_list(path: Path, rows: list[tuple[Path, str]], root: Path) -> None:
    with path.open("w", encoding="utf-8") as handle:
        for image, text in rows:
            handle.write(f"{image.relative_to(root)}\t{text}\n")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--phtd", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--printed-train", type=int, default=4000)
    parser.add_argument("--printed-eval", type=int, default=500)
    args = parser.parse_args()

    args.output.mkdir(parents=True, exist_ok=True)
    handwriting_train = read_gt(args.phtd / "train" / "gt.txt")
    handwriting_eval = read_gt(args.phtd / "val" / "gt.txt")
    printed_train, printed_eval = extract_printed(
        "Reza2kn/persian-printed-ocr-3.5m",
        args.output,
        args.printed_train,
        args.printed_eval,
    )
    hardwords = load_hardword_text("Reza2kn/visualears-hardword-sentences", args.output)

    train = handwriting_train + printed_train
    evaluation = handwriting_eval + printed_eval
    data_root = args.output.parent
    write_list(args.output / "train_list.txt", train, data_root)
    write_list(args.output / "val_list.txt", evaluation, data_root)

    frequencies = Counter("".join(text for _, text in train + evaluation) + "".join(hardwords))
    characters = sorted(
        char for char in frequencies
        if char != " " and char.isprintable() and char not in "\r\n\t"
    )
    (args.output / "persian_dict.txt").write_text(
        "".join(f"{char}\n" for char in characters), encoding="utf-8"
    )
    summary = {
        "train_rows": len(train),
        "eval_rows": len(evaluation),
        "handwriting_train": len(handwriting_train),
        "handwriting_eval": len(handwriting_eval),
        "printed_train": len(printed_train),
        "printed_eval": len(printed_eval),
        "hardword_sentences_for_charset": len(hardwords),
        "characters": len(characters),
        "max_train_characters": max(map(lambda row: len(row[1]), train)),
        "max_eval_characters": max(map(lambda row: len(row[1]), evaluation)),
    }
    (args.output / "summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    print(json.dumps(summary, ensure_ascii=False))


if __name__ == "__main__":
    main()