#!/usr/bin/env python3
"""Build the minimal hash-addressed K1 runtime imported by the E15 worker."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import uuid
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

SCHEMA = "missioncore.e15-worker-package/v1"
COPIED_FILES = (
    "k1link/__init__.py",
    "k1link/compute/inline_temporal.py",
    "k1link/compute/live_perception.py",
    "k1link/data_plane/__init__.py",
    "k1link/data_plane/views.py",
    "k1link/device_plugins/__init__.py",
    "k1link/device_plugins/xgrids_k1/protocol/__init__.py",
    "k1link/device_plugins/xgrids_k1/protocol/normalizer.py",
    "k1link/device_plugins/xgrids_k1/protocol/protobuf_wire.py",
    "k1link/device_plugins/xgrids_k1/protocol/streams.py",
)
GENERATED_FILES = {
    "k1link/compute/__init__.py": (
        '"""Minimal E15 worker projection; import live_perception explicitly."""\n'
    ),
    "k1link/device_plugins/xgrids_k1/__init__.py": (
        '"""Minimal E15 K1 decoder projection; no device-runtime side effects."""\n'
    ),
}


def _canonical(value: object) -> bytes:
    return json.dumps(
        value,
        sort_keys=True,
        separators=(",", ":"),
        allow_nan=False,
    ).encode()


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        while chunk := stream.read(1024 * 1024):
            digest.update(chunk)
    return digest.hexdigest()


def _identity(source_root: Path) -> dict[str, Any]:
    source_files = []
    for relative in COPIED_FILES:
        source = source_root / relative
        if not source.is_file() or source.is_symlink():
            raise RuntimeError(f"E15 worker source is not a regular file: {relative}")
        source_files.append({"path": relative, "sha256": _sha256(source)})
    return {
        "schema_version": SCHEMA,
        "classification": "minimal-live-worker-import-projection",
        "source_files": source_files,
        "generated_initializers": {
            path: hashlib.sha256(content.encode()).hexdigest()
            for path, content in sorted(GENERATED_FILES.items())
        },
        "imports": [
            "k1link.compute.inline_temporal.TemporalStabilizer",
            "k1link.compute.inline_temporal.StreamingSemanticStabilizer",
            "k1link.compute.live_perception.LiveSensorSynchronizer",
            "k1link.data_plane.DecodedPointCloudView",
            "k1link.data_plane.DecodedPoseView",
            "k1link.device_plugins.xgrids_k1.protocol.normalizer.normalize_k1_message",
        ],
    }


def validate_worker_package(root: Path, *, allow_staging: bool = False) -> dict[str, Any]:
    resolved = root.resolve(strict=True)
    manifest_path = resolved / "manifest.json"
    manifest = json.loads(manifest_path.read_text(encoding="utf-8-sig"))
    identity = manifest.get("identity")
    identity_sha256 = manifest.get("identity_sha256")
    artifacts = manifest.get("artifacts")
    expected_paths = {*COPIED_FILES, *GENERATED_FILES, "manifest.json"}
    actual_paths = {
        path.relative_to(resolved).as_posix() for path in resolved.rglob("*") if path.is_file()
    }
    if (
        manifest.get("schema_version") != SCHEMA
        or not isinstance(identity, dict)
        or identity.get("schema_version") != SCHEMA
        or not isinstance(identity_sha256, str)
        or hashlib.sha256(_canonical(identity)).hexdigest() != identity_sha256
        or manifest.get("package_id") != f"e15-worker-package-{identity_sha256}"
        or (
            resolved.name != manifest.get("package_id")
            and not (
                allow_staging
                and resolved.name.startswith(f".{manifest.get('package_id')}.")
                and resolved.name.endswith(".tmp")
            )
        )
        or not isinstance(artifacts, list)
        or actual_paths != expected_paths
    ):
        raise RuntimeError("E15 worker package identity is invalid")
    artifact_paths = set()
    for artifact in artifacts:
        if not isinstance(artifact, dict):
            raise RuntimeError("E15 worker package artifact is invalid")
        relative = str(artifact.get("path", ""))
        path = resolved / relative
        if (
            relative not in expected_paths - {"manifest.json"}
            or relative in artifact_paths
            or not path.is_file()
            or path.is_symlink()
            or artifact.get("byte_length") != path.stat().st_size
            or artifact.get("sha256") != _sha256(path)
        ):
            raise RuntimeError("E15 worker package artifact changed")
        artifact_paths.add(relative)
    if artifact_paths != expected_paths - {"manifest.json"}:
        raise RuntimeError("E15 worker package artifact set changed")
    return manifest


def build_worker_package(source_root: Path, output_root: Path) -> Path:
    source = source_root.resolve(strict=True)
    output = output_root.resolve()
    output.mkdir(mode=0o700, parents=True, exist_ok=True)
    identity = _identity(source)
    identity_sha256 = hashlib.sha256(_canonical(identity)).hexdigest()
    package_id = f"e15-worker-package-{identity_sha256}"
    destination = output / package_id
    if destination.exists():
        validate_worker_package(destination)
        return destination
    staging = output / f".{package_id}.{uuid.uuid4().hex}.tmp"
    staging.mkdir(mode=0o700)
    try:
        for relative in COPIED_FILES:
            target = staging / relative
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.copyfile(source / relative, target)
        for relative, content in GENERATED_FILES.items():
            target = staging / relative
            target.parent.mkdir(parents=True, exist_ok=True)
            target.write_text(content, encoding="utf-8")
        artifacts = []
        for relative in sorted((*COPIED_FILES, *GENERATED_FILES)):
            path = staging / relative
            artifacts.append(
                {
                    "path": relative,
                    "byte_length": path.stat().st_size,
                    "sha256": _sha256(path),
                }
            )
        manifest = {
            "schema_version": SCHEMA,
            "package_id": package_id,
            "identity_sha256": identity_sha256,
            "identity": identity,
            "created_at_utc": datetime.now(UTC)
            .isoformat(timespec="milliseconds")
            .replace("+00:00", "Z"),
            "artifacts": artifacts,
        }
        with (staging / "manifest.json").open("x", encoding="utf-8", newline="\n") as stream:
            json.dump(manifest, stream, indent=2, sort_keys=True)
            stream.write("\n")
            stream.flush()
            os.fsync(stream.fileno())
        validate_worker_package(staging, allow_staging=True)
        staging.replace(destination)
    finally:
        if staging.exists():
            shutil.rmtree(staging)
    validate_worker_package(destination)
    return destination


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source-root", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    args = parser.parse_args()
    destination = build_worker_package(args.source_root, args.output_root)
    print(destination)
    return 0


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