import json
from pathlib import Path
import numpy as np
from Tbd.helper import  CameraObject

def _np_to_list(a: np.ndarray):
    return None if a is None else a.tolist()

def save_calibration_json(camera_list, file_path: str, logger=None):
    payload = {
        "version": 1,
        "cameras": []
    }

    for cam in camera_list:
        payload["cameras"].append({
            "index": int(cam.index),
            "pxWidth": int(cam.pxWidth),
            "pxHeight": int(cam.pxHeight),

            "camera_matrix": _np_to_list(cam.camera_matrix),
            "distortion_coefficients": _np_to_list(cam.distortion_coefficients),

            "rotation_matrix_world_to_camera": _np_to_list(cam.rotation_matrix_world_to_camera),
            "translation_vector_world_to_camera": _np_to_list(cam.translation_vector_world_to_camera),

            "camera_projection_matrix": _np_to_list(cam.camera_projection_matrix),
        })

    Path(file_path).parent.mkdir(parents=True, exist_ok=True)
    with open(file_path, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2)

    if logger:
        logger.log(f"[Calibration] Saved calibration to {file_path}")

def _list_to_np(x, shape=None):
    if x is None:
        return None
    a = np.array(x, dtype=np.float64)
    if shape is not None:
        a = a.reshape(shape)
    return a

def load_calibration_json(file_path: str, logger=None):
    file_path = str(file_path)
    if not Path(file_path).exists():
        raise FileNotFoundError(file_path)

    with open(file_path, "r", encoding="utf-8") as f:
        payload = json.load(f)

    cams = []
    for c in payload.get("cameras", []):
        cam = CameraObject(
            capture=None,
            index=int(c["index"]),
            pxWidth=int(c.get("pxWidth", 1536)),
            pxHeight=int(c.get("pxHeight", 2048)),
        )

        cam.camera_matrix = _list_to_np(c.get("camera_matrix"), shape=(3,3))
        cam.distortion_coefficients = _list_to_np(c.get("distortion_coefficients"))  # keep native shape

        cam.rotation_matrix_world_to_camera = _list_to_np(c.get("rotation_matrix_world_to_camera"), shape=(3,3))
        tv = _list_to_np(c.get("translation_vector_world_to_camera"))
        if tv is not None:
            cam.translation_vector_world_to_camera = tv.reshape(3,1)

        cam.camera_projection_matrix = _list_to_np(c.get("camera_projection_matrix"), shape=(3,4))
        cam.load_video_capture()
        cams.append(cam)

    if logger:
        logger.log(f"[Calibration] Loaded calibration from {file_path} ({len(cams)} cameras)")

    return cams