import os
import sys
import glob
import argparse
import numpy as np
import cv2

ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.append(ROOT)

from model_data_processing.Crop_And_Fill import (
    process_complete,
    smart_crop,
    remove_noise,
    masked_hist_equalize_color_bgr,
)
from Load_process.file_processing import Process_File


def list_images(path):
    exts = ("*.png", "*.jpg", "*.jpeg", "*.bmp", "*.tif", "*.tiff")
    files = []
    for ext in exts:
        files.extend(glob.glob(os.path.join(path, ext)))
    return sorted(files)


def make_synthetic_images(save_dir, n=3):
    pf = Process_File()
    pf.JudgeRoot_MakeDir(save_dir)
    rng = np.random.default_rng(123)
    paths = []
    for i in range(n):
        h, w = 300, 300
        img = np.zeros((h, w, 3), dtype=np.uint8)
        border = int(rng.integers(15, 40))
        content = rng.integers(40, 180, size=(h - 2*border, w - 2*border, 3), dtype=np.uint8)
        img[border:h-border, border:w-border] = content
        path = os.path.join(save_dir, f"synthetic_{i+1}.png")
        cv2.imwrite(path, img)
        paths.append(path)
    return paths


def run_single(image_path, output_dir, args):
    pf = Process_File()
    pf.JudgeRoot_MakeDir(output_dir)
    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"無法讀取: {image_path}")

    denoised = img if args.no_denoise else remove_noise(img, method=args.denoise_method, strength=args.denoise_strength)
    eq = masked_hist_equalize_color_bgr(
        denoised,
        ignore_below=(args.ignore_below if args.ignore_below is not None else args.threshold),
        space=args.equalize_space,
    )
    cropped, borders, success = smart_crop(eq, args.threshold, args.min_content_ratio, keep_margin=0, show_info=False)

    result = process_complete(
        img,
        threshold=args.threshold,
        min_content_ratio=args.min_content_ratio,
        fill_method=args.fill_method,
        denoise=(not args.no_denoise),
        denoise_method=args.denoise_method,
        denoise_strength=args.denoise_strength,
        equalize=(not args.no_equalize),
        equalize_ignore_below=(args.ignore_below if args.ignore_below is not None else None),
        equalize_space=args.equalize_space,
    )

    base = os.path.basename(image_path)
    name, _ = os.path.splitext(base)
    p_cropped = os.path.join(output_dir, f"{name}_cropped.png")
    p_equalized = os.path.join(output_dir, f"{name}_equalized.png")
    p_processed = os.path.join(output_dir, f"{name}_processed.png")
    cv2.imwrite(p_cropped, cropped)
    cv2.imwrite(p_equalized, eq)
    cv2.imwrite(p_processed, result)
    print(p_cropped)
    print(p_equalized)
    print(p_processed)


def run_batch(input_dir, output_dir, args):
    images = list_images(input_dir)
    if len(images) == 0:
        raise FileNotFoundError(f"資料夾中沒有影像: {input_dir}")
    for p in images:
        run_single(p, output_dir, args)
    print(output_dir)


def main():
    parser = argparse.ArgumentParser(description="測試 Crop_And_Fill.py 流程並輸出裁切、等化、處理後影像")
    parser.add_argument("--input", type=str, default=None)
    parser.add_argument("--output", type=str, default=os.path.join(ROOT, "outputs", "crop_fill_test"))
    parser.add_argument("--threshold", type=int, default=30)
    parser.add_argument("--min_content_ratio", type=float, default=0.3)
    parser.add_argument("--fill_method", type=str, default="inpaint_ns")
    parser.add_argument("--no_denoise", action="store_true")
    parser.add_argument("--denoise_method", type=str, default="median")
    parser.add_argument("--denoise_strength", type=str, default="medium")
    parser.add_argument("--no_equalize", action="store_true")
    parser.add_argument("--ignore_below", type=int, default=None)
    parser.add_argument("--equalize_space", type=str, default="LAB")
    args = parser.parse_args()

    pf = Process_File()
    pf.JudgeRoot_MakeDir(args.output)

    if args.input is None:
        synth_dir = os.path.join(ROOT, "tests", "sample_images")
        paths = make_synthetic_images(synth_dir, n=3)
        for p in paths:
            run_single(p, args.output, args)
    else:
        if os.path.isdir(args.input):
            run_batch(args.input, args.output, args)
        else:
            run_single(args.input, args.output, args)


if __name__ == "__main__":
    main()