#!/usr/bin/env python3
"""Interactive zone calibration tool.

Click 4 points for truck area, 4 for counting area, 2 for count line.
Keys: s=save, c=clear/restart, q/ESC=quit.
"""

import argparse
import json
import sys
import cv2
import numpy as np

ZONE_KEYS = ["truck_area", "counting_area"]
ZONE_SIZES = [4, 4]
LINE_POINTS = 2


def load_zones(path: str) -> dict:
    try:
        with open(path) as f:
            return json.load(f)
    except FileNotFoundError:
        return {"presets": {}, "active_preset": None}


def save_zones(config: dict, path: str):
    with open(path, "w") as f:
        json.dump(config, f, indent=2)


def draw_state(frame, points, current_zone, zone_idx, zone_sizes, zone_names):
    vis = frame.copy()
    colors = [(0, 255, 0), (255, 0, 0), (0, 0, 255)]
    total_zones = len(zone_names)

    all_zones = []
    offset = 0
    for i, sz in enumerate(zone_sizes):
        all_zones.append(points[offset : offset + sz])
        offset += sz
    if current_zone < total_zones:
        all_zones.append(points[offset:])

    for i, zone_pts in enumerate(all_zones):
        if not zone_pts:
            continue
        color = colors[i % len(colors)]
        if i < total_zones:
            name = zone_names[i]
            need = zone_sizes[i]
        else:
            name = "count_line"
            need = LINE_POINTS

        if len(zone_pts) > 1:
            pts = np.array(zone_pts, dtype=np.int32)
            if len(zone_pts) >= 3 and i < total_zones:
                cv2.polylines(vis, [pts], True, color, 2)
            else:
                cv2.polylines(vis, [pts], False, color, 2)

        for j, pt in enumerate(zone_pts):
            cv2.circle(vis, pt, 6, color, -1)
            cv2.putText(vis, str(j + 1), (pt[0] + 8, pt[1] - 8),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

        if i == current_zone:
            status = f"[{name}] {len(zone_pts)}/{need} pts"
        else:
            status = f"{name}: {len(zone_pts)}/{need}"
        y_pos = 30 + i * 25
        cv2.putText(vis, status, (10, y_pos),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

    if current_zone < total_zones:
        cur_name = zone_names[current_zone]
    else:
        cur_name = "count_line"
    hint = f"Click {cur_name} point. s=save c=clear q=quit"
    cv2.putText(vis, hint, (10, vis.shape[0] - 20),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
    return vis


def run(video_path: str, preset_name: str, zones_path: str):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Error: cannot open {video_path}")
        sys.exit(1)

    ret, frame = cap.read()
    cap.release()
    if not ret:
        print("Error: cannot read frame from video")
        sys.exit(1)

    h, w = frame.shape[:2]
    print(f"Video: {w}x{h}")

    points = []
    all_sizes = list(ZONE_SIZES) + [LINE_POINTS]
    all_names = list(ZONE_KEYS) + ["count_line"]
    zone_idx = 0

    def mouse_cb(event, x, y, flags, param):
        nonlocal points, zone_idx
        if event == cv2.EVENT_LBUTTONDOWN:
            total = sum(all_sizes)
            if len(points) < total:
                points.append((x, y))
                cur_total = sum(all_sizes[:zone_idx])
                if zone_idx < len(all_sizes):
                    cur_total += len(points) - sum(all_sizes[:zone_idx])
                    if cur_total >= all_sizes[zone_idx] and zone_idx < len(all_sizes) - 1:
                        zone_idx += 1

    win = "Calibrate Zone"
    cv2.namedWindow(win, cv2.WINDOW_NORMAL)
    cv2.resizeWindow(win, min(w, 1280), min(h, 720))
    cv2.setMouseCallback(win, mouse_cb)

    while True:
        vis = draw_state(frame, points, zone_idx, all_sizes, all_names)
        cv2.imshow(win, vis)
        key = cv2.waitKey(30) & 0xFF

        if key in (ord("q"), 27):
            cv2.destroyAllWindows()
            print("Quit.")
            return

        if key == ord("c"):
            points.clear()
            zone_idx = 0
            print("Cleared.")

        if key == ord("s"):
            total_needed = sum(all_sizes)
            if len(points) < total_needed:
                print(f"Need {total_needed} points, have {len(points)}. Continue clicking.")
                continue

            offset = 0
            preset = {}
            for i, sz in enumerate(ZONE_SIZES):
                preset[ZONE_KEYS[i]] = [list(p) for p in points[offset : offset + sz]]
                offset += sz

            x_start = points[offset][0]
            x_end = points[offset + 1][0]
            line_y = min(points[offset][1], points[offset + 1][1])
            preset["count_line"] = {"x_start": x_start, "x_end": x_end, "y": line_y}

            config = load_zones(zones_path)
            config.setdefault("presets", {})[preset_name] = preset
            save_zones(config, zones_path)
            print(f"Saved preset '{preset_name}' to {zones_path}")
            cv2.destroyAllWindows()
            return

    cv2.destroyAllWindows()


def main():
    parser = argparse.ArgumentParser(description="Interactive zone calibration tool")
    parser.add_argument("video_path", help="Path to video file")
    parser.add_argument("--preset", default="calibrated", help="Preset name (default: calibrated)")
    parser.add_argument("--zones", default="zones.json", help="Path to zones.json")
    args = parser.parse_args()
    run(args.video_path, args.preset, args.zones)


if __name__ == "__main__":
    main()