from __future__ import annotations

import argparse
import json
import os
from datetime import datetime
from pathlib import Path

import cv2
import numpy as np


WINDOW_NAME = "ROI Editor"
POLYGON_WINDOW_NAME = "ROI Polygon Editor"


def to_uint8_image(image):
    array = np.asarray(image)
    if array.dtype == np.uint8:
        return array
    array = np.nan_to_num(array.astype(np.float32), copy=False)
    if array.size > 0 and float(np.max(array)) <= 1.0:
        array = array * 255.0
    return np.clip(array, 0, 255).astype(np.uint8)


def normalize_editor_frame(frame):
    if not isinstance(frame, np.ndarray):
        raise TypeError("frame must be a numpy ndarray with shape HxW or HxWxC")

    frame = to_uint8_image(frame)
    if frame.ndim == 2:
        return cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
    if frame.ndim != 3:
        raise ValueError("frame must have shape HxW or HxWxC")

    channels = frame.shape[2]
    if channels == 1:
        return cv2.cvtColor(frame[:, :, 0], cv2.COLOR_GRAY2BGR)
    if channels == 3:
        return np.ascontiguousarray(frame)
    if channels == 4:
        return cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR)
    raise ValueError("frame channel count must be 1, 3, or 4")


def polygon_to_mask(points, image_size):
    width, height = int(image_size[0]), int(image_size[1])
    mask = np.zeros((height, width), dtype=np.uint8)
    polygon = np.asarray(points, dtype=np.int32).reshape(-1, 1, 2)
    cv2.fillPoly(mask, [polygon], 255)
    return mask


def scale_points(points, source_size, target_size):
    source_w, source_h = int(source_size[0]), int(source_size[1])
    target_w, target_h = int(target_size[0]), int(target_size[1])
    pts = np.asarray(points, dtype=np.float32).reshape(-1, 2)
    if source_w > 0 and source_h > 0 and (source_w, source_h) != (target_w, target_h):
        pts[:, 0] *= target_w / float(source_w)
        pts[:, 1] *= target_h / float(source_h)
    pts[:, 0] = np.clip(pts[:, 0], 0, max(target_w - 1, 0))
    pts[:, 1] = np.clip(pts[:, 1], 0, max(target_h - 1, 0))
    return np.rint(pts).astype(np.int32)


def normalize_mask(mask, shape):
    if mask is None:
        return np.ones(shape, dtype=np.uint8) * 255
    mask = np.asarray(mask)
    if mask.shape[:2] != shape:
        mask = cv2.resize(mask, (shape[1], shape[0]), interpolation=cv2.INTER_NEAREST)
    return np.where(mask > 0, 255, 0).astype(np.uint8)


def build_roi_payload(image_size, points, mask_path=None, metadata=None):
    """Build the shared ``roi`` object written into sensor / FFC JSON configs."""
    points = np.asarray(points, dtype=np.int32).reshape(-1, 2)
    if len(points) < 3:
        raise ValueError("polygon ROI requires at least 3 points")
    payload = {
        "type": "polygon",
        "image_size": [int(image_size[0]), int(image_size[1])],
        "points": [[int(x), int(y)] for x, y in points],
        "created_at": datetime.now().isoformat(timespec="seconds"),
    }

    if mask_path is not None:
        payload["mask_path"] = str(mask_path)
    if metadata:
        payload["metadata"] = metadata
    return payload


def extract_roi_section(data):
    """Return the ROI dict from either a nested ``roi`` field or a legacy top-level ROI file."""
    if not isinstance(data, dict):
        raise TypeError("ROI config must be a JSON object")
    if isinstance(data.get("roi"), dict):
        return data["roi"]
    return data


def load_roi_config(path):
    path = Path(path)
    with path.open("r", encoding="utf-8") as f:
        data = json.load(f)
    data["_config_path"] = str(path)
    return data


def save_polygon_roi_config(path, points, image_size, mask_path=None, metadata=None):
    """Backward-compatible standalone polygon ROI JSON writer."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    data = build_roi_payload(
        image_size,
        points,
        mask_path=mask_path,
        metadata=metadata,
    )
    data["version"] = "1.0"
    with path.open("w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)
    return data


def upsert_roi_in_config(path, roi_payload, *, create_if_missing=True):
    """Read-modify-write ``roi`` into an existing sensor config, preserving other fields."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.exists():
        with path.open("r", encoding="utf-8") as f:
            data = json.load(f)
        if not isinstance(data, dict):
            raise ValueError(f"Config must be a JSON object: {path}")
    elif create_if_missing:
        data = {"version": "1.0"}
    else:
        raise FileNotFoundError(f"Config not found: {path}")

    data["roi"] = roi_payload
    with path.open("w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)
        f.write("\n")
    return data


def load_roi_mask(config_path, frame_shape):
    config_path = Path(config_path)
    data = load_roi_config(config_path)
    roi = extract_roi_section(data)

    height, width = frame_shape[:2]
    target_size = (width, height)
    source_size = roi.get("image_size") or data.get("image_size") or [width, height]
    points = roi.get("points")
    points = roi.get("points") or roi.get("polygon") or roi.get("roi_polygon")
    if points:
        polygon = scale_points(points, source_size, target_size).reshape(-1, 1, 2)
        output = np.zeros((height, width), dtype=np.uint8)
        cv2.fillPoly(output, [polygon], 255)
        return output

    mask_path = roi.get("mask_path") or data.get("mask_path") or data.get("roi_mask")
    if mask_path:
        mask_path = Path(mask_path)
        if not mask_path.is_absolute():
            mask_path = config_path.parent / mask_path
        loaded = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
        if loaded is None:
            raise FileNotFoundError(f"Failed to read ROI mask: {mask_path}")
        return normalize_mask(loaded, (height, width))

    raise ValueError(f"Unsupported ROI config: {config_path}")


def resolve_video_source(video):
    """Accept camera index (int/str digit) or filesystem path."""
    if isinstance(video, int):
        return video
    text = str(video).strip()
    if text.isdigit():
        return int(text)
    return text


def open_video_capture(video):
    source = resolve_video_source(video)
    # Avoid MSMF HW transform stalls on Windows camera capture.
    os.environ.setdefault("OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS", "0")
    if isinstance(source, int):
        cap = cv2.VideoCapture(source, cv2.CAP_DSHOW)
        if not cap.isOpened():
            cap.release()
            cap = cv2.VideoCapture(source)
    else:
        path = Path(source)
        if not path.exists():
            raise FileNotFoundError(f"Video file not found: {path}")
        cap = cv2.VideoCapture(str(path))
    if not cap.isOpened():
        raise RuntimeError(f"Failed to open video source: {video}")
    return cap


def capture_frame(cap, *, max_attempts=60):
    """Grab the next readable frame from an opened capture."""
    last_error = None
    for _ in range(max_attempts):
        ok, frame = cap.read()
        if ok and frame is not None and getattr(frame, "size", 0) > 0:
            return normalize_editor_frame(frame)
        last_error = "empty frame" if ok else "read failed"
    raise RuntimeError(f"Failed to capture frame from video source ({last_error})")


def capture_frame_from_source(video):
    cap = open_video_capture(video)
    try:
        return capture_frame(cap)
    finally:
        cap.release()


class RoiPolygonEditor:
    def __init__(self, frame):
        self.frame = normalize_editor_frame(frame)
        self.points = []

    def set_frame(self, frame):
        self.frame = normalize_editor_frame(frame)
        self.points = []

    def on_mouse(self, event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONDOWN:
            self.points.append((int(x), int(y)))
        elif event == cv2.EVENT_RBUTTONDOWN and self.points:
            self.points.pop()

    def clear(self):
        self.points = []

    def undo(self):
        if self.points:
            self.points.pop()

    def render(self):
        view = self.frame.copy()
        if len(self.points) >= 3:
            overlay = view.copy()
            polygon = np.asarray(self.points, dtype=np.int32).reshape(-1, 1, 2)
            cv2.fillPoly(overlay, [polygon], (0, 128, 255))
            view = cv2.addWeighted(overlay, 0.25, view, 0.75, 0.0)
            cv2.polylines(view, [polygon], True, (0, 255, 255), 2, cv2.LINE_AA)
        elif len(self.points) >= 2:
            cv2.polylines(
                view,
                [np.asarray(self.points, dtype=np.int32).reshape(-1, 1, 2)],
                False,
                (0, 255, 255),
                2,
                cv2.LINE_AA,
            )

        for idx, (x, y) in enumerate(self.points):
            cv2.circle(view, (x, y), 4, (0, 255, 0), -1, cv2.LINE_AA)
            cv2.putText(
                view,
                str(idx + 1),
                (x + 6, y - 6),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.45,
                (0, 255, 0),
                1,
                cv2.LINE_AA,
            )

        lines = [
            "LClick add  RClick/U undo  C clear  R recapture  S/Enter save  Q/Esc quit",
            f"points={len(self.points)} resolution={view.shape[1]}x{view.shape[0]}",
        ]
        for i, text in enumerate(lines):
            y = 24 + i * 24
            cv2.putText(view, text, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 3, cv2.LINE_AA)
            cv2.putText(view, text, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1, cv2.LINE_AA)
        return view


def build_roi_result(
    points,
    image_size,
    mask,
    config=None,
    config_path=None,
    mask_path=None,
):
    result = {
        "type": "polygon",
        "points": [[int(x), int(y)] for x, y in points] if points is not None else None,
        "image_size": [int(image_size[0]), int(image_size[1])],
        "mask": mask,
        "config": config,
        "config_path": str(config_path) if config_path is not None else None,
        "mask_path": str(mask_path) if mask_path is not None else None,
    }
    return result


def _save_roi_outputs(
    *,
    image_size,
    output,
    points,
    metadata=None,
    save_mask=True,
    nest_in_config=True,
):
    mask = polygon_to_mask(points, image_size)
    config = None
    config_path = None
    mask_path = None

    if output is not None:
        config_path = Path(output)
        config_path.parent.mkdir(parents=True, exist_ok=True)
        saved_mask = None
        if save_mask:
            mask_path = config_path.with_name(config_path.stem + "_mask.png")
            cv2.imwrite(str(mask_path), mask)
            saved_mask = mask_path.name

        roi_payload = build_roi_payload(
            image_size,
            points,
            mask_path=saved_mask,
            metadata=metadata,
        )
        if nest_in_config:
            config = upsert_roi_in_config(config_path, roi_payload)
        else:
            # Standalone ROI file (legacy FFC polygon style).
            standalone = dict(roi_payload)
            standalone["version"] = "1.0"
            with config_path.open("w", encoding="utf-8") as f:
                json.dump(standalone, f, indent=2, ensure_ascii=False)
                f.write("\n")
            config = standalone

    return build_roi_result(
        points,
        image_size,
        mask,
        config=config,
        config_path=config_path,
        mask_path=mask_path,
    )


def select_polygon_roi(
    frame,
    output=None,
    window_name=POLYGON_WINDOW_NAME,
    metadata=None,
    save_mask=True,
    destroy_window=True,
    recapture_callback=None,
    nest_in_config=False,
):
    """Open an OpenCV window, draw a polygon ROI, and return points/mask/config info."""
    editor = RoiPolygonEditor(frame)
    cv2.namedWindow(window_name, cv2.WINDOW_AUTOSIZE)
    cv2.setMouseCallback(window_name, editor.on_mouse)

    try:
        while True:
            cv2.imshow(window_name, editor.render())
            key = cv2.waitKey(16) & 0xFF
            if key in (27, ord("q")):
                return None
            if key in (8, ord("u")):
                editor.undo()
            elif key == ord("c"):
                editor.clear()
            elif key == ord("r") and recapture_callback is not None:
                editor.set_frame(recapture_callback())
            elif key in (13, ord("s")):
                if len(editor.points) < 3:
                    print("Need at least 3 points to save ROI.")
                    continue

                image_size = (editor.frame.shape[1], editor.frame.shape[0])
                return _save_roi_outputs(
                    image_size=image_size,
                    output=output,
                    points=editor.points,
                    metadata=metadata,
                    save_mask=save_mask,
                    nest_in_config=nest_in_config,
                )
    finally:
        if destroy_window:
            cv2.destroyWindow(window_name)


select_roi = select_polygon_roi


def parse_args(argv=None):
    parser = argparse.ArgumentParser(
        description="Capture a frame from a camera/video and save a polygon ROI to JSON."
    )
    parser.add_argument(
        "--video",
        "-v",
        required=True,
        help="Camera index (e.g. 0) or path to a video file",
    )
    parser.add_argument(
        "--config",
        "-c",
        required=True,
        help="Output JSON config path (existing sensor configs are updated in place)",
    )
    parser.add_argument(
        "--frame-output",
        default=None,
        help="Optional path to save the captured reference frame (PNG/JPG)",
    )
    parser.add_argument(
        "--no-mask",
        action="store_true",
        help="Do not write a companion *_mask.png next to the config",
    )
    return parser.parse_args(argv)


def main(argv=None):
    args = parse_args(argv)
    video = resolve_video_source(args.video)
    cap = open_video_capture(video)

    def recapture():
        # For file sources, rewind so R always gets a usable frame.
        if not isinstance(video, int):
            cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
        return capture_frame(cap)

    try:
        frame = capture_frame(cap)
        if args.frame_output:
            frame_path = Path(args.frame_output)
            frame_path.parent.mkdir(parents=True, exist_ok=True)
            cv2.imwrite(str(frame_path), frame)
            print(f"Saved reference frame: {frame_path}")

        result = select_polygon_roi(
            frame,
            output=args.config,
            metadata={"source": str(args.video)},
            save_mask=not args.no_mask,
            recapture_callback=recapture,
            nest_in_config=True,
        )
    finally:
        cap.release()
        cv2.destroyAllWindows()

    if result is None:
        print("ROI capture cancelled.")
        return 1

    print(f"Saved ROI type={result['type']} image_size={result['image_size']}")
    if result.get("points") is not None:
        print(f"points={result['points']}")
    if result.get("config_path"):
        print(f"config={result['config_path']}")
    if result.get("mask_path"):
        print(f"mask={result['mask_path']}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())