#!/usr/bin/env python3
"""Convert PaddlePaddle {lang}_PP-OCRv5_mobile_rec_onnx to NCNN for PaddleOCR-ncnn-CPP."""
from __future__ import annotations
import argparse, re, subprocess, sys
from pathlib import Path
import yaml

def fix_param(text: str) -> str:
    lines = []
    for line in text.splitlines():
        if line.startswith("Input"):
            line = line.replace(
                "Input                    in0                      0 1 in0",
                "Input                    input                    0 1 input",
            )
            lines.append(line)
            continue
        parts = line.split()
        if len(parts) >= 4:
            for i in range(len(parts) - 1, 3, -1):
                if parts[i] == "in0": parts[i] = "input"
                elif parts[i] == "out0": parts[i] = "output"
            line = " ".join(parts)
        lines.append(line)
    return "\n".join(lines) + "\n"

def num_classes(param_text: str) -> int:
    num = 0
    for line in param_text.splitlines():
        if line.startswith("Gemm"):
            m = re.search(r" 8=(\d+)", line)
            if m: num = int(m.group(1))
    return num

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("lang", help="e.g. arabic, ta, en")
    ap.add_argument("--onnx-dir", type=Path, default=Path("onnx"))
    ap.add_argument("--out-dir", type=Path, default=Path("."))
    ap.add_argument("--pnnx", default="pnnx")
    args = ap.parse_args()
    lang = args.lang
    onnx = args.onnx_dir / "inference.onnx"
    yml = args.onnx_dir / "inference.yml"
    if not onnx.exists():
        print("missing", onnx, file=sys.stderr); sys.exit(1)
    subprocess.run([args.pnnx, str(onnx), "inputshape=[1,3,48,320]", "inputshape2=[1,3,48,640]", "fp16=0"], check=True, cwd=args.onnx_dir)
    param_text = fix_param((args.onnx_dir / "inference.ncnn.param").read_text())
    ncls = num_classes(param_text)
    chars = yaml.safe_load(yml.read_text(encoding="utf-8"))["PostProcess"]["character_dict"]
    keys = [""] + [str(c) for c in chars]
    while len(keys) < ncls: keys.append(" ")
    keys = keys[:ncls]
    args.out_dir.mkdir(parents=True, exist_ok=True)
    (args.out_dir / f"{lang}_PP-OCRv5_mobile_rec.param").write_text(param_text)
    (args.out_dir / f"{lang}_PP-OCRv5_mobile_rec.bin").write_bytes((args.onnx_dir / "inference.ncnn.bin").read_bytes())
    (args.out_dir / f"ppocr_keys_{lang}.txt").write_text("\n".join(keys) + "\n", encoding="utf-8")
    print(f"wrote {lang}: classes={ncls}")

if __name__ == "__main__":
    main()