import argparse
import csv
import sys
from pathlib import Path

import cv2

from colorchecker import (
    detect_colorchecker,
    PATCH_NAMES,
    ResultCode,
)


# ============================================================
# Main
# ============================================================

def main():

    parser = argparse.ArgumentParser(
        description="画像からColorCheckerを検出し、RGB値を計測します"
    )
    parser.add_argument(
        "filename",
        help="入力画像ファイルのパス"
    )
    args = parser.parse_args()
    code, rgb, images = detect_colorchecker(args.filename)
    if code != ResultCode.SUCCESS:
        print(
            f"ColorChecker処理失敗: {code.name} ({int(code)}): {args.filename}",
            file=sys.stderr,
        )
        return int(code)

    results = [
        {
            "index": index + 1,
            "name": name,
            "R": float(rgb[0, index]),
            "G": float(rgb[1, index]),
            "B": float(rgb[2, index]),
        }
        for index, name in enumerate(PATCH_NAMES)
    ]

    print()
    print( f"{'No.':>3} " f"{'Name':<18} " f"{'R':>8} " f"{'G':>8} " f"{'B':>8}" )
    for r in results:
        print(
            f"{r['index']:3d} "
            f"{r['name']:<18} "
            f"{r['R']:8.2f} "
            f"{r['G']:8.2f} "
            f"{r['B']:8.2f}"
        )

    # --------------------------------------------------------
    # CSV保存
    # --------------------------------------------------------
    input_path = Path(args.filename)
    output_path = input_path.with_name(
        f"{input_path.stem}_rgb.csv"
    )
    with open(
        output_path,
        "w",
        newline="",
        encoding="utf-8"
    ) as f:

        writer = csv.writer(f)

        writer.writerow( [ "No", "Name", "R", "G", "B" ] )
        for r in results:
            writer.writerow([
                r["index"],
                r["name"],
                r["R"],
                r["G"],
                r["B"]
            ])
    print()
    print( f"{output_path} saved" )

    # --------------------------------------------------------
    # デバッグ画像保存
    # --------------------------------------------------------
    for suffix, key in (("warped", "warped"), ("patches", "debug"), ("mask", "mask")):
        output_path = input_path.with_name(
            f"{input_path.stem}_{suffix}{input_path.suffix}"
        )
        cv2.imwrite(str(output_path), images[key])
        print(f"{output_path} saved")

    return int(ResultCode.SUCCESS)


if __name__ == "__main__":
    sys.exit(main())