#!/usr/bin/env python3

import argparse
import glob
import json
import os
import time
from collections import defaultdict


# ------------------------------------------------------------

def load_detection_json(path):

    with open(path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    return data


# ------------------------------------------------------------

def find_detection_files(directory, timestamp):

    pattern = os.path.join(
        directory,
        f'*_{timestamp}_aruco_detection.json'
    )

    files = sorted(glob.glob(pattern))

    return files


# ------------------------------------------------------------

def build_scene(detection_files, timestamp):

    scene = {

        "schema_version": "1.0",

        "scene_timestamp": str(timestamp),

        "created_utc": time.strftime(
            "%Y-%m-%dT%H:%M:%SZ",
            time.gmtime()
        ),

        "cameras": {},

        "markers": {}
    }

    marker_map = defaultdict(list)

    # --------------------------------------------------------
    # Read all camera detection jsons
    # --------------------------------------------------------

    for json_file in detection_files:

        data = load_detection_json(json_file)

        cam = data["camera"]
        img = data["image"]

        basename = os.path.basename(json_file)
        # snapshot_video1_1778819665744_aruco_detection.json
        parts = basename.split('_')
        video_part = parts[1]          # "video1"
        camera_id = video_part.replace('video', '')

        # ----------------------------------------------------
        # Camera information
        # ----------------------------------------------------

        scene["cameras"][camera_id] = {

            "camera_id": camera_id,

            "intrinsics_file": cam["intrinsics_file"],

            "camera_matrix": cam["camera_matrix"],

            "distortion_coefficients":
                cam["distortion_coefficients"],

            "image": {

                "image_file": img["image_file"],

                "image_sha256": img["image_sha256"],

                "width_px": img["width_px"],

                "height_px": img["height_px"]
            }
        }

        # ----------------------------------------------------
        # Marker observations
        # ----------------------------------------------------

        for obs in data["observations"]:

            marker_id = str(obs["marker_id"])

            observation = {

                "camera_id": camera_id,

                "center_px": obs["center_px"],

                "corners_px": obs["corners_px"],

                "bbox_px": obs["bbox_px"],

                "quality": obs["quality"]
            }

            marker_map[marker_id].append(observation)

    # --------------------------------------------------------
    # Convert marker map
    # --------------------------------------------------------

    for marker_id, observations in marker_map.items():

        scene["markers"][marker_id] = {

            "marker_id": int(marker_id),

            "num_observations": len(observations),

            "observations": observations
        }

    return scene


# ------------------------------------------------------------

def main():

    parser = argparse.ArgumentParser()

    parser.add_argument(
        '-timestamp',
        required=True,
        help='Scene timestamp'
    )

    parser.add_argument(
        '-dir',
        default='.',
        help='Directory containing detection json files'
    )

    parser.add_argument(
        '-o',
        '--output',
        default=None
    )

    args = parser.parse_args()

    detection_files = find_detection_files(
        args.dir,
        args.timestamp
    )

    if len(detection_files) == 0:

        raise RuntimeError(
            f'No detection json files found for timestamp '
            f'{args.timestamp}'
        )

    print('Found detection files:')

    for f in detection_files:
        print('  ', f)

    scene = build_scene(
        detection_files,
        args.timestamp
    )

    # --------------------------------------------------------
    # Output filename
    # --------------------------------------------------------

    if args.output is None:

        out_json = os.path.join(
            args.dir,
            f'scene_{args.timestamp}.json'
        )

    else:

        out_json = args.output

    with open(out_json, 'w', encoding='utf-8') as f:

        json.dump(scene, f, indent=2)

    print()
    print(f'Saved scene: {out_json}')


# ------------------------------------------------------------

if __name__ == '__main__':
    main()