#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Convert PPOCRLabel output into PaddleOCR-style training folders.

Input PPOCRLabel project example:
  label_project/
    images/
    Label.txt
    rec_gt.txt
    crop_img/

Output:
  train_data/japanese_ocr/
    det/images/train, det/images/val
    det/train_label.txt, det/val_label.txt
    rec/images/train, rec/images/val
    rec/rec_gt_train.txt, rec/rec_gt_val.txt
    dict/ppocr_japanese_product_dict.txt
"""

from __future__ import annotations

import argparse
import json
import random
import shutil
from pathlib import Path
from typing import Iterable, List, Tuple


IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}


def read_lines(path: Path) -> List[str]:
    if not path.exists():
        return []
    return [line.rstrip("\n") for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]


def safe_copy(src: Path, dst: Path) -> None:
    dst.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(src, dst)


def split_items(items: List[str], val_ratio: float, seed: int) -> Tuple[List[str], List[str]]:
    items = list(items)
    rnd = random.Random(seed)
    rnd.shuffle(items)
    val_count = max(1, int(round(len(items) * val_ratio))) if len(items) > 1 else 0
    val = items[:val_count]
    train = items[val_count:]
    return train, val


def resolve_image(project_dir: Path, rel_path: str) -> Path:
    p = project_dir / rel_path
    if p.exists():
        return p

    # PPOCRLabel may store only the file name.
    p2 = project_dir / "images" / rel_path
    if p2.exists():
        return p2

    candidates = list(project_dir.rglob(Path(rel_path).name))
    if candidates:
        return candidates[0]

    raise FileNotFoundError(f"image not found: {rel_path} under {project_dir}")


def parse_det_line(line: str) -> Tuple[str, str]:
    if "\t" not in line:
        raise ValueError(f"invalid detection label line, tab is required: {line}")
    img_rel, label_json = line.split("\t", 1)

    # Validate JSON early.
    json.loads(label_json)
    return img_rel, label_json


def parse_rec_line(line: str) -> Tuple[str, str]:
    if "\t" not in line:
        raise ValueError(f"invalid recognition label line, tab is required: {line}")
    img_rel, text = line.split("\t", 1)
    return img_rel, text


def build_dict(rec_lines: Iterable[str]) -> List[str]:
    chars = []
    seen = set()

    for line in rec_lines:
        _, text = parse_rec_line(line)
        for ch in text:
            if ch not in seen and ch not in {"\n", "\r", "\t"}:
                seen.add(ch)
                chars.append(ch)

    # Stable order: common symbols first if present, then Unicode order for reproducibility.
    return sorted(chars)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--ppocrlabel_dir", nargs="+", required=True)
    parser.add_argument("--out_dir", required=True)
    parser.add_argument("--val_ratio", type=float, default=0.1)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    project_dirs = [Path(p).resolve() for p in args.ppocrlabel_dir]
    out_dir = Path(args.out_dir).resolve()

    det_items = []
    rec_items = []

    for project_dir in project_dirs:
        label_txt = project_dir / "Label.txt"
        rec_gt = project_dir / "rec_gt.txt"

        for line in read_lines(label_txt):
            img_rel, label_json = parse_det_line(line)
            det_items.append((project_dir, img_rel, label_json))

        for line in read_lines(rec_gt):
            img_rel, text = parse_rec_line(line)
            rec_items.append((project_dir, img_rel, text))

    if not det_items:
        raise RuntimeError("No detection labels found. Please check Label.txt.")
    if not rec_items:
        raise RuntimeError("No recognition labels found. Please check rec_gt.txt.")

    det_indices = [str(i) for i in range(len(det_items))]
    rec_indices = [str(i) for i in range(len(rec_items))]

    det_train_ids, det_val_ids = split_items(det_indices, args.val_ratio, args.seed)
    rec_train_ids, rec_val_ids = split_items(rec_indices, args.val_ratio, args.seed + 1)

    det_train_ids = set(map(int, det_train_ids))
    det_val_ids = set(map(int, det_val_ids))
    rec_train_ids = set(map(int, rec_train_ids))
    rec_val_ids = set(map(int, rec_val_ids))

    det_train_lines = []
    det_val_lines = []

    for i, (project_dir, img_rel, label_json) in enumerate(det_items):
        src = resolve_image(project_dir, img_rel)
        subset = "val" if i in det_val_ids else "train"
        dst_rel = f"images/{subset}/{src.name}"
        dst = out_dir / "det" / dst_rel
        safe_copy(src, dst)

        line = f"{dst_rel}\t{label_json}"
        if subset == "val":
            det_val_lines.append(line)
        else:
            det_train_lines.append(line)

    rec_train_lines = []
    rec_val_lines = []

    for i, (project_dir, img_rel, text) in enumerate(rec_items):
        src = resolve_image(project_dir, img_rel)
        subset = "val" if i in rec_val_ids else "train"
        dst_rel = f"images/{subset}/{src.name}"
        dst = out_dir / "rec" / dst_rel
        safe_copy(src, dst)

        line = f"{dst_rel}\t{text}"
        if subset == "val":
            rec_val_lines.append(line)
        else:
            rec_train_lines.append(line)

    (out_dir / "det").mkdir(parents=True, exist_ok=True)
    (out_dir / "rec").mkdir(parents=True, exist_ok=True)
    (out_dir / "dict").mkdir(parents=True, exist_ok=True)

    (out_dir / "det" / "train_label.txt").write_text("\n".join(det_train_lines) + "\n", encoding="utf-8")
    (out_dir / "det" / "val_label.txt").write_text("\n".join(det_val_lines) + "\n", encoding="utf-8")
    (out_dir / "rec" / "rec_gt_train.txt").write_text("\n".join(rec_train_lines) + "\n", encoding="utf-8")
    (out_dir / "rec" / "rec_gt_val.txt").write_text("\n".join(rec_val_lines) + "\n", encoding="utf-8")

    chars = build_dict([f"{img}\t{text}" for _, img, text in rec_items])
    (out_dir / "dict" / "ppocr_japanese_product_dict.txt").write_text(
        "\n".join(chars) + "\n",
        encoding="utf-8",
    )

    print("done:", out_dir)
    print("det train:", len(det_train_lines))
    print("det val:", len(det_val_lines))
    print("rec train:", len(rec_train_lines))
    print("rec val:", len(rec_val_lines))
    print("dict chars:", len(chars))


if __name__ == "__main__":
    main()
