"""Camera registration + intrinsics I/O for ar_view.

Cameras are referenced by ROLE ("scene" / "wrist"), not by raw index:
  cameras.json            -> {role: {"index": N, "width": W, "height": H}}
  intrinsics/<role>.json  -> {"width","height","K","dist","rms","num_frames"}

`open_named` resolves a role to a live capture. If the registered index no longer has the
registered native resolution (indices shuffled), it scans for the index whose resolution
matches and uses that — so identical-resolution cameras still need a stable port, but
differently-sized cameras (e.g. a webcam vs the B0332) self-identify regardless of index.
"""
import json
import os
import time

import numpy as np
import cv2

HERE = os.path.dirname(os.path.abspath(__file__))
CAMS_FILE = os.path.join(HERE, "cameras.json")
INTR_DIR = os.path.join(HERE, "intrinsics")
DEFAULT_INDEX = {"scene": 0, "wrist": 1}


CAP_WH = (1920, 1080)   # requested capture resolution (cameras like the iPhone default low otherwise)


def open_index(idx, warm=40, cap_wh=CAP_WH):
    """Open + warm up a camera index at a high resolution. Returns (cap, first_frame) or (None, None)."""
    for backend in (getattr(cv2, "CAP_AVFOUNDATION", cv2.CAP_ANY), cv2.CAP_ANY):
        cap = cv2.VideoCapture(idx, backend)
        if not cap.isOpened():
            continue
        if cap_wh:
            cap.set(cv2.CAP_PROP_FRAME_WIDTH, cap_wh[0])
            cap.set(cv2.CAP_PROP_FRAME_HEIGHT, cap_wh[1])
        for _ in range(warm):
            ok, fr = cap.read()
            if ok and fr is not None:
                return cap, fr
            time.sleep(0.05)
        return cap, None
    return None, None


def load_cameras():
    try:
        return json.load(open(CAMS_FILE))
    except FileNotFoundError:
        return {}


def register(role, index, width, height):
    cams = load_cameras()
    cams[role] = {"index": int(index), "width": int(width), "height": int(height)}
    json.dump(cams, open(CAMS_FILE, "w"), indent=2)


def load_intrinsics(role):
    p = os.path.join(INTR_DIR, f"{role}.json")
    if not os.path.exists(p):
        return None
    d = json.load(open(p))
    d["K"] = np.array(d["K"], float)
    d["dist"] = np.array(d["dist"], float)
    return d


def save_intrinsics(role, K, dist, width, height, rms, num_frames):
    os.makedirs(INTR_DIR, exist_ok=True)
    json.dump({"width": int(width), "height": int(height),
               "K": np.asarray(K).tolist(), "dist": np.asarray(dist).ravel().tolist(),
               "rms": float(rms), "num_frames": int(num_frames)},
              open(os.path.join(INTR_DIR, f"{role}.json"), "w"), indent=2)


def _res(fr):
    return (fr.shape[1], fr.shape[0]) if fr is not None else None


def open_named(role, override_index=None, scan=5):
    """-> (cap, index, first_frame). An explicit override_index is used as-is; otherwise
    resolve via the registered index, auto-healing by the registered resolution."""
    if override_index is not None:
        cap, fr = open_index(override_index)
        return cap, override_index, fr

    entry = load_cameras().get(role, {})
    want = (entry["width"], entry["height"]) if entry.get("width") else None
    idx = entry.get("index", DEFAULT_INDEX.get(role, 0))

    cap, fr = open_index(idx)
    if cap is not None and (want is None or _res(fr) == want):
        return cap, idx, fr
    if cap is not None:
        cap.release()
    if want is not None:                       # try to find the registered resolution elsewhere
        for j in range(scan):
            if j == idx:
                continue
            c, f = open_index(j)
            if c is not None and _res(f) == want:
                print(f"  [{role}] registered index {idx} moved; using index {j} (matched {want[0]}x{want[1]})")
                return c, j, f
            if c is not None:
                c.release()
    cap, fr = open_index(idx)                  # give up: use the registered index as-is
    return cap, idx, fr