#!/usr/bin/env python3
"""Materialize LAB E19 from immutable E14 detections, masks, and LiDAR.

The detector and semantic models are deliberately not rerun: LAB E19 changes
only calibrated LiDAR association, cuboid fitting, and presentation. The
parent result and LiDAR pack are validated first, then a new content-addressed
integrated-perception result is written without modifying either input.
"""

from __future__ import annotations

import argparse
import copy
import hashlib
import json
import os
import platform
import shutil
import tempfile
from collections import Counter
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

import numpy as np
from e10_fusion_runtime import (
    CuboidCompletionTracker,
    LidarReplayPack,
    WorldStateProjector,
    _completion_profile,
    _object_support_ground_filter_profile,
    canonical_json,
    clearance,
    distance_history,
    fuse_tracks,
    fusion_document,
    project_points,
    sha256,
)

from k1link.compute.integrated_perception import (
    FUSION_SCHEMA,
    IDENTITY_SCHEMA,
    REPORT_SCHEMA,
    RESULT_SCHEMA,
    SEMANTIC_SCHEMA,
    WORLD_SCHEMA,
    validate_integrated_perception_result,
)

PIPELINE_ID = "fixed-e14-yolox-eomt-e19-ground-aware-refusion/v1"


def arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--job", type=Path, required=True)
    parser.add_argument("--parent-result", type=Path, required=True)
    parser.add_argument("--lidar-packs-root", type=Path, required=True)
    parser.add_argument("--profile", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    return parser.parse_args()


def read_object(path: Path) -> dict[str, Any]:
    value = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise RuntimeError(f"JSON root is not an object: {path}")
    return value


def read_rows(path: Path) -> list[dict[str, Any]]:
    rows = []
    with path.open(encoding="utf-8") as stream:
        for line in stream:
            value = json.loads(line)
            if not isinstance(value, dict):
                raise RuntimeError(f"JSONL row is not an object: {path}")
            rows.append(value)
    return rows


def write_json(path: Path, value: object) -> None:
    path.write_bytes(canonical_json(value) + b"\n")
    os.chmod(path, 0o600)


def write_row(stream: Any, value: object) -> None:
    stream.write(canonical_json(value).decode() + "\n")


def artifact(
    path: Path,
    kind: str,
    media_type: str,
    schema_version: str | None = None,
) -> dict[str, Any]:
    value = {
        "kind": kind,
        "path": path.name,
        "media_type": media_type,
        "byte_length": path.stat().st_size,
        "sha256": sha256(path),
    }
    if schema_version is not None:
        value["schema_version"] = schema_version
    return value


def read_profile(path: Path) -> tuple[dict[str, Any], str]:
    resolved = path.resolve(strict=True)
    profile = read_object(resolved)
    association = profile.get("association")
    if (
        profile.get("schema_version") != "missioncore.e10-integrated-perception-profile/v1"
        or profile.get("profile_id") != "lab-e19-ground-aware-cuboids-v1"
        or profile.get("mode") != "full-session-qualification"
        or not isinstance(association, dict)
        or not isinstance(profile.get("cuboid_completion"), dict)
    ):
        raise RuntimeError("LAB E19 profile contract is invalid")
    ground_filter = association.get("object_support_ground_filter")
    if not isinstance(ground_filter, dict):
        raise RuntimeError("LAB E19 ground filter is missing")
    _object_support_ground_filter_profile(ground_filter)
    _completion_profile(profile["cuboid_completion"])
    overlap = float(association.get("support_duplicate_overlap_threshold", 0.0))
    if not 0.0 < overlap <= 1.0:
        raise RuntimeError("LAB E19 support overlap threshold is invalid")
    return profile, sha256(resolved)


def color_for(group: str, track_id: int) -> tuple[int, int, int]:
    seed = hashlib.sha256(f"e10:{group}:{track_id}".encode()).digest()
    return tuple(64 + value % 176 for value in seed[:3])


def materialize(args: argparse.Namespace) -> Path:
    parent = validate_integrated_perception_result(
        args.job,
        args.parent_result,
        args.lidar_packs_root,
    )
    if not parent.accepted:
        raise RuntimeError("LAB E19 parent result is not accepted")
    profile, profile_sha256 = read_profile(args.profile)
    parent_result = read_object(parent.result_root / "result.json")
    parent_report = read_object(parent.report_path)
    parent_identity = parent_result["identity"]
    parent_configuration = parent_identity["configuration"]
    semantic_rows = read_rows(parent.semantic_path)
    parent_fusion_rows = read_rows(parent.fusion_path)
    parent_world_rows = read_rows(parent.world_path)
    if not (
        len(parent_fusion_rows) == len(parent_world_rows) == parent.frame_count
        and profile["selection"]["required_frame_count"] == parent.frame_count
    ):
        raise RuntimeError("LAB E19 parent selection changed")

    identity = {
        "schema_version": IDENTITY_SCHEMA,
        "job_id": parent_identity["job_id"],
        "input_sha256": parent_identity["input_sha256"],
        "session_id": parent_identity["session_id"],
        "source_id": parent_identity["source_id"],
        "lidar_pack_id": parent_identity["lidar_pack_id"],
        "selection": copy.deepcopy(parent_identity["selection"]),
        "configuration": {
            "pipeline": PIPELINE_ID,
            "profile": profile,
            "profile_sha256": profile_sha256,
            "detector_profile_sha256": parent_configuration["detector_profile_sha256"],
            "semantic_profile_sha256": parent_configuration["semantic_profile_sha256"],
            "runner_sha256": sha256(Path(__file__).resolve(strict=True)),
            "fusion_runtime_sha256": sha256(
                Path(__file__).with_name("e10_fusion_runtime.py").resolve(strict=True)
            ),
            "orchestrator_sha256": sha256(Path(__file__).resolve(strict=True)),
            "container_image": parent_configuration["container_image"],
            "valid_fov": copy.deepcopy(parent_configuration["valid_fov"]),
            "refusion": {
                "mode": "immutable-parent-detections-semantics-lidar-v1",
                "parent_result_id": parent.result_id,
                "parent_result_json_sha256": sha256(parent.result_root / "result.json"),
                "parent_fusion_sha256": sha256(parent.fusion_path),
                "parent_arrays_sha256": sha256(parent.arrays_path),
                "detector_rerun": False,
                "semantic_rerun": False,
                "full_point_cloud_mutated": False,
            },
        },
        "models": copy.deepcopy(parent_identity["models"]),
    }
    identity_sha256 = hashlib.sha256(canonical_json(identity)).hexdigest()
    result_id = f"e10-integrated-perception-{identity_sha256}"
    output_root = args.output_root.expanduser().absolute()
    output_root.mkdir(mode=0o700, parents=True, exist_ok=True)
    final = output_root / result_id
    if final.exists():
        validate_integrated_perception_result(args.job, final, args.lidar_packs_root)
        return final

    staging = Path(tempfile.mkdtemp(prefix=".e19-refusion-", dir=output_root))
    os.chmod(staging, 0o700)
    lidar = LidarReplayPack(parent.pack_root, expected_job_id=parent.job.job_id)
    final_created = False
    try:
        semantic_path = staging / "semantic-frames.jsonl"
        shutil.copyfile(parent.semantic_path, semantic_path)
        os.chmod(semantic_path, 0o600)
        fusion_path = staging / "fusion-frames.jsonl"
        world_path = staging / "world-state.jsonl"
        gpu_path = staging / "gpu-telemetry.jsonl"

        with np.load(parent.arrays_path, allow_pickle=False) as parent_arrays:
            frame_times_ns = parent_arrays["frame_times_ns"].copy()
            semantic_frame_indices = parent_arrays["semantic_frame_indices"].copy()
            semantic_masks = parent_arrays["semantic_masks"].copy()
        if not np.array_equal(
            semantic_frame_indices,
            np.asarray([row["frame_index"] for row in semantic_rows], dtype=np.int64),
        ):
            raise RuntimeError("LAB E19 semantic mask binding changed")
        semantic_by_source = {
            int(row["source_frame_index"]): mask
            for row, mask in zip(semantic_rows, semantic_masks, strict=True)
        }

        history = distance_history(int(profile["association"]["distance_history_frames"]))
        completion = CuboidCompletionTracker(profile["cuboid_completion"])
        projector = WorldStateProjector(float(profile["world_state"]["velocity_history_limit_s"]))
        support_offsets = [0]
        support_points: list[np.ndarray] = []
        support_colors: list[np.ndarray] = []
        box_offsets = [0]
        box_centers: list[tuple[float, float, float]] = []
        box_half_sizes: list[tuple[float, float, float]] = []
        box_quaternions: list[tuple[float, float, float, float]] = []
        box_colors: list[tuple[int, int, int, int]] = []
        status_counts: Counter[str] = Counter()
        fusion_state_counts: Counter[str] = Counter()
        accepted_cuboids = 0
        old_accepted_cuboids = 0
        ground_rejected_points = 0
        tracks_with_ground_rejections = 0
        duplicate_support_rejections = 0
        fused_frames = 0

        with (
            fusion_path.open("x", encoding="utf-8", newline="\n") as fusion_stream,
            world_path.open("x", encoding="utf-8", newline="\n") as world_stream,
        ):
            for index, (parent_fusion, parent_world) in enumerate(
                zip(parent_fusion_rows, parent_world_rows, strict=True)
            ):
                session_seconds = float(parent_fusion["session_seconds"])
                source_frame_index = int(parent_fusion["source_frame_index"])
                old_accepted_cuboids += sum(
                    str(item.get("cuboid_status", "")).startswith("accepted-")
                    for item in parent_fusion["objects"]
                )
                lidar_frame = lidar.frame(index)
                points_lidar = np.empty((0, 3), dtype=np.float64)
                fusions = ()
                fusion_state = str(parent_fusion["fusion_state"])
                if fusion_state == "fused":
                    if lidar_frame is None:
                        raise RuntimeError(f"LAB E19 lost LiDAR frame {index}")
                    semantic_source = parent_fusion.get("semantic_source_frame_index")
                    semantic_map = semantic_by_source.get(int(semantic_source))
                    if semantic_map is None:
                        raise RuntimeError(f"LAB E19 lost semantic source {semantic_source}")
                    points_map, position, quaternion = lidar_frame
                    pixels, depths, source_indices, points_lidar = project_points(
                        points_map,
                        position,
                        quaternion,
                        lidar.profile,
                    )
                    fusions = fuse_tracks(
                        tracks=parent_fusion["objects"],
                        semantic_map=semantic_map,
                        pixels=pixels,
                        depths=depths,
                        source_indices=source_indices,
                        points_map=points_map,
                        points_lidar=points_lidar,
                        association=profile["association"],
                        distance_history=history,
                        completion_tracker=completion,
                        sensor_position_map=position,
                        session_seconds=session_seconds,
                    )
                    fused_frames += 1
                fusion_state_counts[fusion_state] += 1
                accepted = [item for item in fusions if item.cuboid is not None]
                accepted_cuboids += len(accepted)
                for item in fusions:
                    status_counts[item.status] += 1
                    ground_rejected_points += item.ground_rejected_points
                    tracks_with_ground_rejections += item.ground_rejected_points > 0
                    duplicate_support_rejections += (
                        item.status == "rejected-duplicate-lidar-support"
                    )

                frame_support = []
                frame_support_colors = []
                for item in accepted:
                    if lidar_frame is None:
                        raise RuntimeError("LAB E19 accepted a cuboid without LiDAR")
                    color = color_for(item.association_group, item.track_id)
                    values = lidar_frame[0][item.source_indices].astype(np.float32)
                    frame_support.append(values)
                    frame_support_colors.append(
                        np.tile(np.asarray([color], dtype=np.uint8), (values.shape[0], 1))
                    )
                    box_centers.append(item.cuboid.center_map)
                    box_half_sizes.append(item.cuboid.half_size)
                    box_quaternions.append(item.cuboid.quaternion_xyzw)
                    box_colors.append((*color, 88))
                if frame_support:
                    support = np.concatenate(frame_support)
                    colors = np.concatenate(frame_support_colors)
                    support_points.append(support)
                    support_colors.append(colors)
                    support_offsets.append(support_offsets[-1] + support.shape[0])
                else:
                    support_offsets.append(support_offsets[-1])
                box_offsets.append(box_offsets[-1] + len(accepted))

                clearance_state = clearance(
                    points_lidar,
                    profile["world_state"]["clearance"],
                )
                world = projector.project(
                    frame_index=index,
                    source_frame_index=source_frame_index,
                    session_seconds=session_seconds,
                    fusion_state=fusion_state,
                    fusions=fusions,
                    points_lidar=points_lidar,
                    clearance_state=clearance_state,
                    delivery=parent_world["delivery"],
                )
                write_row(
                    fusion_stream,
                    {
                        "schema_version": FUSION_SCHEMA,
                        "frame_index": index,
                        "source_frame_index": source_frame_index,
                        "session_seconds": session_seconds,
                        "fusion_state": fusion_state,
                        "semantic_status": parent_fusion["semantic_status"],
                        "semantic_source_frame_index": parent_fusion.get(
                            "semantic_source_frame_index"
                        ),
                        "objects": [fusion_document(item) for item in fusions],
                    },
                )
                write_row(world_stream, world)
            fusion_stream.flush()
            world_stream.flush()
            os.fsync(fusion_stream.fileno())
            os.fsync(world_stream.fileno())
        os.chmod(fusion_path, 0o600)
        os.chmod(world_path, 0o600)

        arrays_path = staging / "transient-perception.npz"
        np.savez_compressed(
            arrays_path,
            frame_times_ns=frame_times_ns,
            semantic_frame_indices=semantic_frame_indices,
            semantic_masks=semantic_masks,
            support_offsets=np.asarray(support_offsets, dtype=np.int64),
            support_points=np.concatenate(support_points)
            if support_points
            else np.empty((0, 3), dtype=np.float32),
            support_colors=np.concatenate(support_colors)
            if support_colors
            else np.empty((0, 3), dtype=np.uint8),
            box_offsets=np.asarray(box_offsets, dtype=np.int64),
            box_centers=np.asarray(box_centers, dtype=np.float32).reshape((-1, 3)),
            box_half_sizes=np.asarray(box_half_sizes, dtype=np.float32).reshape((-1, 3)),
            box_quaternions=np.asarray(box_quaternions, dtype=np.float32).reshape((-1, 4)),
            box_colors=np.asarray(box_colors, dtype=np.uint8).reshape((-1, 4)),
        )
        os.chmod(arrays_path, 0o600)
        with gpu_path.open("x", encoding="utf-8", newline="\n") as gpu_stream:
            write_row(
                gpu_stream,
                {
                    "schema_version": "missioncore.e19-refusion-telemetry/v1",
                    "gpu_used": False,
                    "parent_result_id": parent.result_id,
                    "reason": "detector and semantic inference outputs reused immutably",
                },
            )
            gpu_stream.flush()
            os.fsync(gpu_stream.fileno())
        os.chmod(gpu_path, 0o600)

        acceptance_profile = profile["acceptance"]
        checks = {
            "parent_result_accepted": parent.accepted,
            "frame_accounting": len(parent_fusion_rows) == parent.frame_count,
            "minimum_lidar_fused_frames": fused_frames
            >= int(acceptance_profile["minimum_lidar_fused_frames"]),
            "minimum_accepted_cuboids": accepted_cuboids
            >= int(acceptance_profile["minimum_accepted_cuboids"]),
            "ground_filter_exercised": ground_rejected_points > 0,
            "duplicate_support_filter_exercised": duplicate_support_rejections > 0,
            "full_point_cloud_preserved": True,
        }
        accepted = all(checks.values())
        created_at = datetime.now(UTC).isoformat(timespec="milliseconds").replace("+00:00", "Z")
        parent_metrics = parent_report.get("metrics", {})
        metrics = {
            "source_span_seconds": parent_metrics.get("source_span_seconds"),
            "detector": copy.deepcopy(parent_metrics.get("detector")),
            "semantic": copy.deepcopy(parent_metrics.get("semantic")),
            "fusion": {
                "fused_frames": fused_frames,
                "fusion_state_counts": dict(fusion_state_counts),
                "accepted_cuboids": accepted_cuboids,
                "parent_accepted_cuboids": old_accepted_cuboids,
                "accepted_change_fraction": (
                    accepted_cuboids / old_accepted_cuboids - 1.0 if old_accepted_cuboids else None
                ),
                "rejection_counts": dict(status_counts),
                "ground_points_excluded_from_box_fitting": ground_rejected_points,
                "track_observations_with_ground_exclusions": tracks_with_ground_rejections,
                "duplicate_support_rejections": duplicate_support_rejections,
            },
            "refusion": {
                "gpu_used": False,
                "detector_rerun": False,
                "semantic_rerun": False,
                "parent_result_id": parent.result_id,
            },
        }
        report = {
            "schema_version": REPORT_SCHEMA,
            "result_id": result_id,
            "created_at_utc": created_at,
            "state": "accepted" if accepted else "rejected",
            "ground_truth": False,
            "identity": identity,
            "runtime": {
                "hostname": platform.node(),
                "python": platform.python_version(),
                "numpy": np.__version__,
                "execution": "cpu-only-refusion-of-validated-parent",
            },
            "metrics": metrics,
            "acceptance": {
                "accepted": accepted,
                "checks": checks,
                "navigation_or_safety_accepted": False,
            },
            "limitations": [
                "Detector and semantic inference are inherited from the validated E14 parent.",
                "This is a recorded refusion result, not a physical live K1 gate.",
                "The complete point cloud is an immutable input and is not modified.",
                "Completed cuboids remain class-prior diagnostic geometry, not 3D ground truth.",
                "Navigation and safety authority remain disabled.",
            ],
        }
        report_path = staging / "run-report.json"
        write_json(report_path, report)
        artifacts = [
            artifact(
                semantic_path,
                "e10-semantic-frames",
                "application/x-ndjson",
                SEMANTIC_SCHEMA,
            ),
            artifact(
                fusion_path,
                "e10-fusion-frames",
                "application/x-ndjson",
                FUSION_SCHEMA,
            ),
            artifact(
                world_path,
                "e10-world-state",
                "application/x-ndjson",
                WORLD_SCHEMA,
            ),
            artifact(arrays_path, "e10-transient-perception", "application/x-npz"),
            artifact(gpu_path, "worker-gpu-telemetry", "application/x-ndjson"),
            artifact(
                report_path,
                "e10-run-report",
                "application/json",
                REPORT_SCHEMA,
            ),
        ]
        result = {
            "schema_version": RESULT_SCHEMA,
            "result_id": result_id,
            "identity_sha256": identity_sha256,
            "identity": identity,
            "created_at_utc": created_at,
            "acceptance_state": report["state"],
            "ground_truth": False,
            "publication_scope": "recorded-integrated-realtime-qualification-only",
            "frames_processed": parent.frame_count,
            "artifacts": artifacts,
        }
        write_json(staging / "result.json", result)
        os.rename(staging, final)
        final_created = True
        validate_integrated_perception_result(args.job, final, args.lidar_packs_root)
        return final
    except BaseException:
        if staging.exists():
            shutil.rmtree(staging)
        if final_created and final.exists():
            shutil.rmtree(final)
        raise
    finally:
        lidar.close()


def main() -> int:
    result = materialize(arguments())
    report = read_object(result / "run-report.json")
    print(
        json.dumps(
            {
                "result_root": str(result),
                "result_id": result.name,
                "accepted": report["acceptance"]["accepted"],
                "fusion": report["metrics"]["fusion"],
            },
            ensure_ascii=False,
            sort_keys=True,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
