#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
预生成合成车牌池 (供 C++ 管线使用)
===================================
为每种特殊车牌类型批量生成纯车牌图，C++ 运行时直接读内存池，
避免每张 CCPD 原图都重新渲染车牌。

用法:
  python gen_plate_pool.py --num 5000 --output ./output/plate_pool

输出目录结构:
  plate_pool/
    yellow/        教练车牌 (440x140)
    green_truck/   新能源车牌 (480x140)
    white/         警车车牌
    white_army/    军车车牌
    black/         港澳/领馆车牌
    black_shi/     使馆车牌
"""

import os
import sys
import argparse
from pathlib import Path

import cv2
import numpy as np
from tqdm import tqdm

# 路径
_SCRIPT_DIR = Path(__file__).resolve().parent
_PROJECT_DIR = _SCRIPT_DIR.parent.parent
_GENERATOR_DIR = _PROJECT_DIR / 'chinese_license_plate_generator'
sys.path.insert(0, str(_GENERATOR_DIR))

from generate_multi_plate import MultiPlateGenerator
from plate_number import (
    generate_plate_number_white,
    generate_plate_number_army,
    generate_plate_number_black_gangao,
    generate_plate_number_black_shi_old,
    generate_plate_number_black_shi_new,
    generate_plate_number_black_ling_old,
    generate_plate_number_black_ling_new,
    generate_plate_number_yellow_xue,
    generate_plate_number_blue,
)

# 类型配置
PLATE_TYPES = [
    ('blue',        'blue',        generate_plate_number_blue),
    ('green',       'green_car',   lambda: generate_plate_number_blue(length=8)),
    ('police',      'white',       generate_plate_number_white),
    ('army',        'white_army',  generate_plate_number_army),
    ('gangao',      'black',       generate_plate_number_black_gangao),
    ('shiguan_old', 'black_shi',   generate_plate_number_black_shi_old),
    ('shiguan_new', 'black_shi',   generate_plate_number_black_shi_new),
    ('lingguan_old','black',       generate_plate_number_black_ling_old),
    ('lingguan_new','black',       generate_plate_number_black_ling_new),
    ('coach',       'yellow',      generate_plate_number_yellow_xue),
    ('new_energy',  'green_truck', lambda: generate_plate_number_blue(length=8)),
]


def main():
    parser = argparse.ArgumentParser(description='预生成合成车牌池')
    parser.add_argument('--num', type=int, default=5000,
                        help='每种类型生成数量 (默认: 5000)')
    parser.add_argument('--output', type=str, default='./output/plate_pool',
                        help='输出根目录')
    args = parser.parse_args()

    # 初始化生成器
    pm = str(_GENERATOR_DIR / 'plate_model')
    fm = str(_GENERATOR_DIR / 'font_model')
    gen = MultiPlateGenerator(pm, fm)

    for plate_type, bg_color, num_func in PLATE_TYPES:
        out_dir = Path(args.output) / bg_color
        out_dir.mkdir(parents=True, exist_ok=True)

        existing = len(list(out_dir.glob('*.jpg')))
        if existing >= args.num:
            print(f"[SKIP] {plate_type:12s} ({bg_color:12s}): 已有 {existing} 张")
            continue

        need = args.num - existing
        print(f"[GEN]  {plate_type:12s} ({bg_color:12s}): {need} 张...", end=' ', flush=True)

        success = 0
        pbar = tqdm(total=need, desc=f'{plate_type}', unit='张', leave=False)
        attempts = 0
        while success < need and attempts < need * 3:
            attempts += 1
            try:
                number = num_func()
                plate = gen.generate_plate_special(
                    number, bg_color, is_double=False, enhance=False
                )
                # 用 imencode 支持中文文件名
                fname = f"{number}.jpg"
                ok, buf = cv2.imencode('.jpg', plate, [cv2.IMWRITE_JPEG_QUALITY, 95])
                if ok:
                    with open(str(out_dir / fname), 'wb') as f:
                        f.write(buf.tobytes())
                    success += 1
                    pbar.update(1)
            except Exception:
                pass
        pbar.close()
        print(f"OK ({success} 张)")

    print(f"\n完成! 输出: {args.output}")
    # 统计
    for bg_color in ['yellow','green_truck','white','white_army','black','black_shi']:
        d = Path(args.output) / bg_color
        count = len(list(d.glob('*.jpg'))) if d.is_dir() else 0
        print(f"  {bg_color:15s}: {count} 张")


if __name__ == '__main__':
    import os
    os.chdir(str(_PROJECT_DIR))
    main()