#!/usr/bin/env python3
"""
Apply calibration poses to a ZED Fusion configuration file.

This script takes the output from `calibrate_extrinsics.py` and updates the 
`FusionConfiguration.pose` entries in a fusion configuration JSON file 
(e.g., `inside_network.json`).

Input Formats:
1. Calibration JSON (--calibration-json):
   {
     "12345678": {
       "pose": "1.0 0.0 0.0 0.0 0.0 1.0 0.0 0.0 0.0 0.0 1.0 0.0 0.0 0.0 0.0 1.0",
       ...
     }
   }
   The pose is a space-separated string of 16 floats representing a 4x4 matrix 
   in row-major order (T_world_from_cam).

2. Fusion Config JSON (--fusion-config-json):
   {
     "12345678": {
       "FusionConfiguration": {
         "pose": "...",
         ...
       },
       ...
     }
   }

Usage Example:
    uv run apply_calibration_to_fusion_config.py \\
        --calibration-json output/calibration.json \\
        --fusion-config-json zed_settings/inside_network.json \\
        --output-json output/updated_fusion_config.json
"""

import json
import click
import sys


def validate_pose_string(pose_str: str) -> bool:
    """Ensures the pose string contains exactly 16 floats."""
    try:
        parts = pose_str.split()
        if len(parts) != 16:
            return False
        [float(p) for p in parts]
        return True
    except (ValueError, AttributeError):
        return False


@click.command()
@click.option(
    "--calibration-json",
    required=True,
    type=click.Path(exists=True),
    help="Path to calibration output JSON.",
)
@click.option(
    "--fusion-config-json",
    required=True,
    type=click.Path(exists=True),
    help="Path to fusion configuration JSON.",
)
@click.option(
    "--output-json",
    required=True,
    type=click.Path(),
    help="Path to save the updated fusion configuration.",
)
@click.option(
    "--strict/--no-strict",
    default=False,
    help="Fail if a calibration serial is missing in fusion config.",
)
def main(
    calibration_json: str,
    fusion_config_json: str,
    output_json: str,
    strict: bool,
) -> None:
    """
    Apply calibration poses to a ZED Fusion configuration file.
    """
    with open(calibration_json, "r") as f:
        calib_data: dict[str, dict[str, str]] = json.load(f)

    with open(fusion_config_json, "r") as f:
        fusion_data: dict[str, dict[str, dict[str, str]]] = json.load(f)

    updated_count = 0
    missing_serials: list[str] = []
    untouched_fusion_serials = set(fusion_data.keys())

    # Sort serials for deterministic reporting
    calib_serials = sorted(calib_data.keys())

    for serial in calib_serials:
        pose_str = calib_data[serial].get("pose")
        if not pose_str:
            click.echo(
                f"Warning: Serial {serial} in calibration file has no 'pose' entry. Skipping.",
                err=True,
            )
            continue

        if not validate_pose_string(pose_str):
            click.echo(
                f"Error: Invalid pose string for serial {serial}. Must be 16 floats.",
                err=True,
            )
            sys.exit(1)

        if serial in fusion_data:
            if "FusionConfiguration" in fusion_data[serial]:
                fusion_data[serial]["FusionConfiguration"]["pose"] = pose_str
                updated_count += 1
                untouched_fusion_serials.discard(serial)
            else:
                click.echo(
                    f"Warning: Serial {serial} found in fusion config but lacks 'FusionConfiguration' key.",
                    err=True,
                )
                if strict:
                    sys.exit(1)
        else:
            missing_serials.append(serial)
            if strict:
                click.echo(
                    f"Error: Serial {serial} from calibration not found in fusion config (strict mode).",
                    err=True,
                )
                sys.exit(1)

    with open(output_json, "w") as f:
        json.dump(fusion_data, f, indent=4)

    click.echo("\n--- Summary ---")
    click.echo(f"Updated: {updated_count}")

    if missing_serials:
        click.echo(
            f"Serials in calibration but missing in fusion config: {', '.join(sorted(missing_serials))}"
        )
    else:
        click.echo("All calibration serials were found in fusion config.")

    if untouched_fusion_serials:
        click.echo(
            f"Fusion entries untouched: {', '.join(sorted(untouched_fusion_serials))}"
        )
    else:
        click.echo("All fusion entries were updated.")

    click.echo(f"Output written to: {output_json}")


if __name__ == "__main__":
    main()