from __future__ import annotations

import argparse
import sys
from pathlib import Path

import cv2

REPO_ROOT = Path(__file__).resolve().parents[1]
SRC_ROOT = REPO_ROOT / "src"
for path in (REPO_ROOT, SRC_ROOT):
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))

from orisys import FFCFrameProcessor
from tools.ffc.roi_polygon_tool import load_roi_mask, select_polygon_roi


def parse_source(value):
    value = str(value).strip()
    for prefix in ("camera:", "cam:", "webcam:"):
        if value.lower().startswith(prefix):
            value = value[len(prefix):].strip()
            break
    try:
        return int(value)
    except ValueError:
        return value


def read_image_if_possible(source):
    if isinstance(source, int):
        return None
    path = Path(source)
    if not path.exists() or not path.is_file():
        return None
    image = cv2.imread(str(path), cv2.IMREAD_COLOR)
    return image


def draw_roi_outline(image, roi_mask):
    view = image.copy()
    contours, _ = cv2.findContours(roi_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(view, contours, -1, (0, 255, 255), 1)
    return view


def parse_args():
    parser = argparse.ArgumentParser(description="Test ROI drawing and flat-field correction.")
    parser.add_argument("--source", default="0", help="Camera id, image path, or video path.")
    parser.add_argument("--roi-config", default=None, help="Existing ROI JSON config.")
    parser.add_argument("--roi-output", default="roi_polygon.json", help="Path to save ROI JSON.")
    parser.add_argument("--weights", default=None, help="Existing FFC weights .npz to load.")
    parser.add_argument("--save-weights", default="ffc_weights.npz", help="Path to save calibrated weights.")
    parser.add_argument("--blur-kernel", type=int, default=151)
    parser.add_argument("--blur-sigma", type=float, default=0.0)
    parser.add_argument("--epsilon", type=float, default=1e-6)
    parser.add_argument("--keep-outside", action="store_true", help="Keep pixels outside ROI.")
    return parser.parse_args()


def prepare_corrector(frame, args):
    ffc = FFCFrameProcessor(
        blur_kernel=args.blur_kernel,
        blur_sigma=args.blur_sigma,
        epsilon=args.epsilon,
        masked_outside=not args.keep_outside,
    )

    if args.weights:
        ffc.load(args.weights)
        return ffc, ffc.roi_mask

    if args.roi_config:
        roi_mask = load_roi_mask(args.roi_config, frame.shape)
    else:
        roi = select_polygon_roi(frame, output=args.roi_output)
        if roi is None:
            raise RuntimeError("ROI selection cancelled.")
        roi_mask = roi["mask"]

    ffc.calibrate(frame, roi_mask=roi_mask)
    if args.save_weights:
        ffc.save(args.save_weights)
    return ffc, roi_mask


def run_single_image(image, args):
    ffc, roi_mask = prepare_corrector(image, args)
    corrected = ffc.correct(image)
    cv2.imshow("Input ROI", draw_roi_outline(image, roi_mask))
    cv2.imshow("Corrected ROI", draw_roi_outline(corrected, roi_mask))
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def run_video(source, args):
    cap = cv2.VideoCapture(source)
    if not cap.isOpened():
        raise RuntimeError(f"Cannot open source: {source}")

    try:
        ok, frame = cap.read()
        if not ok:
            raise RuntimeError("Failed to read first frame.")

        ffc, roi_mask = prepare_corrector(frame, args)

        while True:
            corrected = ffc.correct(frame)
            cv2.imshow("Input ROI", draw_roi_outline(frame, roi_mask))
            cv2.imshow("Corrected ROI", draw_roi_outline(corrected, roi_mask))

            key = cv2.waitKey(1) & 0xFF
            if key in (27, ord("q")):
                break
            if key == ord("r"):
                ffc.calibrate(frame, roi_mask=roi_mask)
                if args.save_weights:
                    ffc.save(args.save_weights)
                print("Recalibrated from current frame.")

            ok, frame = cap.read()
            if not ok:
                break
    finally:
        cap.release()
        cv2.destroyAllWindows()


def main():
    args = parse_args()
    source = parse_source(args.source)

    image = read_image_if_possible(source)
    if image is not None:
        run_single_image(image, args)
    else:
        run_video(source, args)


if __name__ == "__main__":
    main()