#!/usr/bin/env python3
from __future__ import annotations

import argparse
import os
from pathlib import Path

import ujson


def main() -> None:
    parser = argparse.ArgumentParser(description="Rewrite CoVT JSON image paths to portable basenames.")
    parser.add_argument("source_json", type=Path)
    parser.add_argument("image_dir", type=Path)
    parser.add_argument("output_json", type=Path)
    args = parser.parse_args()

    args.output_json.parent.mkdir(parents=True, exist_ok=True)
    count = 0
    missing = 0
    first = True
    with args.source_json.open("r", encoding="utf-8") as src, args.output_json.open(
        "w", encoding="utf-8"
    ) as dst:
        dst.write("[\n")
        for raw in src:
            line = raw.strip()
            if not line or line in ("[", "]"):
                continue
            if line.endswith(","):
                line = line[:-1]
            item = ujson.loads(line)
            images = item.get("image")
            if isinstance(images, str):
                basename = os.path.basename(images)
                item["image"] = basename
                missing += int(not (args.image_dir / basename).is_file())
            elif isinstance(images, list):
                basenames = [os.path.basename(value) for value in images]
                item["image"] = basenames
                missing += sum(not (args.image_dir / value).is_file() for value in basenames)
            else:
                raise TypeError(f"row {count} has invalid image field: {type(images).__name__}")
            if not first:
                dst.write(",\n")
            dst.write("  " + ujson.dumps(item, ensure_ascii=False))
            first = False
            count += 1
        dst.write("\n]\n")
    if missing:
        raise FileNotFoundError(f"{missing} image references were not found")
    print(f"portable_rows={count} missing_images={missing} output={args.output_json}")


if __name__ == "__main__":
    main()