from __future__ import annotations

import hashlib
import json
import math
import os
import re
import secrets
import shutil
import stat
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image

from k1link.artifacts import utc_now_iso, write_json_atomic
from k1link.device_plugins.xgrids_k1.analyze.calibrated_projection import (
    Kb4ProjectionProfile,
)
from k1link.device_plugins.xgrids_k1.calibration_schema import (
    MAIN_CAMERA_SLOT_BY_SOURCE,
    parse_k1_factory_calibration,
)
from k1link.device_plugins.xgrids_k1.calibration_snapshot import (
    CALIBRATION_SNAPSHOT_MANIFEST_VERSION,
)

VALID_FOV_MASK_SCHEMA = "missioncore.k1-valid-fov-mask/v1"
VALID_FOV_MASK_IDENTITY_SCHEMA = "missioncore.k1-valid-fov-mask-identity/v1"
DEFAULT_EDGE_MARGIN_PIXELS = 4.0
MAX_MANIFEST_BYTES = 4 * 1024 * 1024

_SHA256 = re.compile(r"^[a-f0-9]{64}$")
_SAFE_GENERATION = re.compile(r"^valid-fov-mask-[a-f0-9]{64}$")


class K1ValidFovMaskError(RuntimeError):
    """A K1 valid-FOV mask could not be derived or validated safely."""


@dataclass(frozen=True, slots=True)
class K1ValidFovMask:
    generation_id: str
    root: Path
    manifest_path: Path
    mask_path: Path
    source_id: str
    calibration_slot: str
    calibration_sha256: str
    width: int
    height: int
    center_xy: tuple[float, float]
    radius_pixels: float
    crop_xyxy: tuple[int, int, int, int]
    valid_pixel_count: int
    valid_fraction: float


def prepare_k1_valid_fov_mask(
    *,
    calibration_snapshot_root: Path,
    source_id: str,
    output_root: Path,
    edge_margin_pixels: float = DEFAULT_EDGE_MARGIN_PIXELS,
) -> K1ValidFovMask:
    """Seal one reusable circular valid-FOV mask bound to exact K1 calibration.

    The mask is generated once from the calibrated principal point and the
    admitted image transform. It is content-addressed by calibration, source,
    geometry and construction policy, so repeated experiments load the same
    PNG rather than estimating a lens boundary from every frame.
    """

    if source_id not in MAIN_CAMERA_SLOT_BY_SOURCE:
        raise K1ValidFovMaskError("source is not an admitted K1 main camera")
    if (
        isinstance(edge_margin_pixels, bool)
        or not isinstance(edge_margin_pixels, (int, float))
        or not math.isfinite(float(edge_margin_pixels))
        or not 0.0 <= float(edge_margin_pixels) <= 64.0
    ):
        raise K1ValidFovMaskError("edge margin must be finite and between 0 and 64 pixels")

    snapshot_root = calibration_snapshot_root.expanduser().resolve(strict=True)
    snapshot = _validated_snapshot(snapshot_root)
    calibration = parse_k1_factory_calibration(
        (snapshot_root / "camera.yaml").read_bytes(),
        (snapshot_root / "extrinsic_camera_lidar.yaml").read_bytes(),
    )
    normalized = calibration.normalized_profile()
    if snapshot.get("normalized_calibration") != normalized:
        raise K1ValidFovMaskError("calibration snapshot normalized profile changed")
    profile = Kb4ProjectionProfile.from_factory_calibration(calibration, source_id)
    stream_bindings = normalized.get("stream_bindings")
    if not isinstance(stream_bindings, dict):
        raise K1ValidFovMaskError("calibration stream bindings are unavailable")
    stream_binding = stream_bindings.get(source_id)
    if not isinstance(stream_binding, dict):
        raise K1ValidFovMaskError("calibration source binding is unavailable")

    fx, fy, cx, cy = profile.intrinsic_fx_fy_cx_cy
    image_edge_radius = min(
        cx,
        float(profile.width - 1) - cx,
        cy,
        float(profile.height - 1) - cy,
    )
    radius = image_edge_radius - float(edge_margin_pixels)
    if not math.isfinite(radius) or radius < 32.0:
        raise K1ValidFovMaskError("valid-FOV construction leaves no usable image circle")

    identity = {
        "schema_version": VALID_FOV_MASK_IDENTITY_SCHEMA,
        "calibration_sha256": snapshot["content_identity_sha256"],
        "source_id": source_id,
        "calibration_slot": profile.calibration_slot,
        "camera_model": "kb4",
        "admitted_resolution": [profile.width, profile.height],
        "admitted_intrinsic_fx_fy_cx_cy": [fx, fy, cx, cy],
        "distortion_kb4": list(profile.distortion_kb4),
        "image_transform": stream_binding.get("image_transform"),
        "construction": {
            "kind": "calibrated-principal-point-inscribed-circle",
            "pixel_coordinate_convention": "integer-pixel-centers",
            "edge_radius_policy": "minimum-distance-to-admitted-image-edge",
            "edge_margin_pixels": float(edge_margin_pixels),
            "outside_value": 0,
            "inside_value": 255,
        },
    }
    identity_sha256 = hashlib.sha256(_canonical_json(identity)).hexdigest()
    generation_id = f"valid-fov-mask-{identity_sha256}"
    root = _prepare_private_directory(output_root)
    final = root / generation_id
    if final.exists():
        existing = validate_k1_valid_fov_mask(final)
        if existing.calibration_sha256 != snapshot["content_identity_sha256"]:
            raise K1ValidFovMaskError("valid-FOV generation collides with another calibration")
        return existing

    staging = root / f".{generation_id}.{secrets.token_hex(12)}.incomplete"
    published = False
    try:
        staging.mkdir(mode=0o700)
        mask = _circle_mask(profile.width, profile.height, cx, cy, radius)
        valid_y, valid_x = np.nonzero(mask)
        if valid_x.size == 0 or valid_y.size == 0:
            raise K1ValidFovMaskError("valid-FOV mask is empty")
        crop = (
            int(valid_x.min()),
            int(valid_y.min()),
            int(valid_x.max()) + 1,
            int(valid_y.max()) + 1,
        )
        mask_path = staging / "mask.png"
        Image.fromarray(mask, mode="L").save(mask_path, format="PNG", optimize=False)
        _fsync_file(mask_path)
        os.chmod(mask_path, 0o600)
        mask_sha256 = _sha256_file(mask_path)
        valid_pixel_count = int(np.count_nonzero(mask))
        valid_fraction = valid_pixel_count / int(mask.size)
        manifest = {
            "schema_version": VALID_FOV_MASK_SCHEMA,
            "generation_id": generation_id,
            "identity_sha256": identity_sha256,
            "identity": identity,
            "created_at_utc": utc_now_iso(),
            "geometry": {
                "center_xy": [cx, cy],
                "image_edge_radius_pixels": image_edge_radius,
                "radius_pixels": radius,
                "crop_xyxy_exclusive": list(crop),
                "valid_pixel_count": valid_pixel_count,
                "total_pixel_count": int(mask.size),
                "valid_fraction": valid_fraction,
            },
            "artifact": {
                "path": "mask.png",
                "media_type": "image/png",
                "mode": "L",
                "inside_value": 255,
                "outside_value": 0,
                "byte_length": mask_path.stat().st_size,
                "sha256": mask_sha256,
            },
            "usage": {
                "quality": "set invalid lens exterior to the model profile's fixed fill value",
                "speed": (
                    "apply crop_xyxy_exclusive before model preprocessing and map outputs back"
                ),
                "dense_compute_warning": (
                    "multiplying an unchanged 800x600 tensor by this mask alone does not reduce "
                    "dense neural-network FLOPs"
                ),
            },
        }
        write_json_atomic(staging / "manifest.json", manifest)
        os.chmod(staging / "manifest.json", 0o600)
        _fsync_directory(staging)
        os.replace(staging, final)
        _fsync_directory(root)
        published = True
    finally:
        if not published and staging.exists():
            shutil.rmtree(staging)

    return validate_k1_valid_fov_mask(final)


def validate_k1_valid_fov_mask(mask_root: Path) -> K1ValidFovMask:
    root = mask_root.expanduser().resolve(strict=True)
    if not root.is_dir() or _SAFE_GENERATION.fullmatch(root.name) is None:
        raise K1ValidFovMaskError("valid-FOV root is invalid")
    manifest = _read_json_object(root / "manifest.json", root, MAX_MANIFEST_BYTES)
    identity = manifest.get("identity")
    identity_sha256 = manifest.get("identity_sha256")
    if (
        manifest.get("schema_version") != VALID_FOV_MASK_SCHEMA
        or manifest.get("generation_id") != root.name
        or not isinstance(identity, dict)
        or identity.get("schema_version") != VALID_FOV_MASK_IDENTITY_SCHEMA
        or not isinstance(identity_sha256, str)
        or _SHA256.fullmatch(identity_sha256) is None
        or root.name != f"valid-fov-mask-{identity_sha256}"
        or hashlib.sha256(_canonical_json(identity)).hexdigest() != identity_sha256
    ):
        raise K1ValidFovMaskError("valid-FOV identity is inconsistent")

    artifact = manifest.get("artifact")
    geometry = manifest.get("geometry")
    resolution = identity.get("admitted_resolution")
    intrinsic = identity.get("admitted_intrinsic_fx_fy_cx_cy")
    if (
        not isinstance(artifact, dict)
        or artifact.get("path") != "mask.png"
        or artifact.get("media_type") != "image/png"
        or artifact.get("mode") != "L"
        or artifact.get("inside_value") != 255
        or artifact.get("outside_value") != 0
        or not isinstance(geometry, dict)
        or not isinstance(resolution, list)
        or len(resolution) != 2
        or not all(isinstance(value, int) and not isinstance(value, bool) for value in resolution)
        or not isinstance(intrinsic, list)
        or len(intrinsic) != 4
    ):
        raise K1ValidFovMaskError("valid-FOV manifest contract is invalid")
    width, height = int(resolution[0]), int(resolution[1])
    if width < 1 or height < 1:
        raise K1ValidFovMaskError("valid-FOV resolution is invalid")
    mask_path = root / "mask.png"
    metadata = _confined_regular_file(mask_path, root)
    if (
        artifact.get("byte_length") != metadata.st_size
        or not isinstance(artifact.get("sha256"), str)
        or _SHA256.fullmatch(str(artifact["sha256"])) is None
        or _sha256_file(mask_path) != artifact["sha256"]
    ):
        raise K1ValidFovMaskError("valid-FOV mask artifact changed")
    try:
        with Image.open(mask_path) as opened:
            if opened.mode != "L" or opened.size != (width, height):
                raise K1ValidFovMaskError("valid-FOV PNG dimensions or mode changed")
            mask = np.asarray(opened, dtype=np.uint8)
    except (OSError, ValueError) as exc:
        raise K1ValidFovMaskError("valid-FOV PNG is unavailable") from exc
    if not np.isin(mask, (0, 255)).all():
        raise K1ValidFovMaskError("valid-FOV PNG is not binary")
    valid_y, valid_x = np.nonzero(mask)
    if valid_x.size == 0 or valid_y.size == 0:
        raise K1ValidFovMaskError("valid-FOV PNG is empty")
    crop = (
        int(valid_x.min()),
        int(valid_y.min()),
        int(valid_x.max()) + 1,
        int(valid_y.max()) + 1,
    )
    count = int(valid_x.size)
    center = geometry.get("center_xy")
    radius = geometry.get("radius_pixels")
    expected_crop = geometry.get("crop_xyxy_exclusive")
    expected_count = geometry.get("valid_pixel_count")
    expected_total = geometry.get("total_pixel_count")
    expected_fraction = geometry.get("valid_fraction")
    if (
        not isinstance(center, list)
        or len(center) != 2
        or not all(isinstance(value, (int, float)) for value in center)
        or not isinstance(radius, (int, float))
        or not math.isfinite(float(radius))
        or list(crop) != expected_crop
        or count != expected_count
        or mask.size != expected_total
        or not isinstance(expected_fraction, (int, float))
        or not math.isclose(count / mask.size, float(expected_fraction), abs_tol=1e-12)
    ):
        raise K1ValidFovMaskError("valid-FOV geometry changed")
    source_id = identity.get("source_id")
    calibration_slot = identity.get("calibration_slot")
    calibration_sha256 = identity.get("calibration_sha256")
    if (
        not isinstance(source_id, str)
        or source_id not in MAIN_CAMERA_SLOT_BY_SOURCE
        or not isinstance(calibration_slot, str)
        or calibration_slot != MAIN_CAMERA_SLOT_BY_SOURCE[source_id]
        or not isinstance(calibration_sha256, str)
        or _SHA256.fullmatch(calibration_sha256) is None
    ):
        raise K1ValidFovMaskError("valid-FOV calibration binding is invalid")
    return K1ValidFovMask(
        generation_id=root.name,
        root=root,
        manifest_path=root / "manifest.json",
        mask_path=mask_path,
        source_id=source_id,
        calibration_slot=calibration_slot,
        calibration_sha256=calibration_sha256,
        width=width,
        height=height,
        center_xy=(float(center[0]), float(center[1])),
        radius_pixels=float(radius),
        crop_xyxy=crop,
        valid_pixel_count=count,
        valid_fraction=count / mask.size,
    )


def _validated_snapshot(root: Path) -> dict[str, Any]:
    manifest = _read_json_object(root / "manifest.json", root, MAX_MANIFEST_BYTES)
    artifacts = manifest.get("artifacts")
    device = manifest.get("device")
    identity_sha256 = manifest.get("content_identity_sha256")
    if (
        manifest.get("schema_version") != CALIBRATION_SNAPSHOT_MANIFEST_VERSION
        or not isinstance(artifacts, list)
        or len(artifacts) != 2
        or not isinstance(device, dict)
        or not isinstance(identity_sha256, str)
        or _SHA256.fullmatch(identity_sha256) is None
    ):
        raise K1ValidFovMaskError("calibration snapshot manifest is incompatible")
    expected_names = {"camera.yaml", "extrinsic_camera_lidar.yaml"}
    by_name: dict[str, dict[str, Any]] = {}
    by_source: dict[str, str] = {}
    for artifact in artifacts:
        if not isinstance(artifact, dict):
            raise K1ValidFovMaskError("calibration snapshot artifact is invalid")
        name = artifact.get("artifact_name")
        source_path = artifact.get("source_path")
        digest = artifact.get("sha256")
        byte_length = artifact.get("bytes")
        if (
            not isinstance(name, str)
            or name not in expected_names
            or name in by_name
            or not isinstance(source_path, str)
            or not source_path.startswith("/mnt/system/factory-data/config/")
            or not isinstance(digest, str)
            or _SHA256.fullmatch(digest) is None
            or not isinstance(byte_length, int)
            or isinstance(byte_length, bool)
            or byte_length < 1
        ):
            raise K1ValidFovMaskError("calibration snapshot artifact descriptor is invalid")
        path = root / name
        metadata = _confined_regular_file(path, root)
        if metadata.st_size != byte_length or _sha256_file(path) != digest:
            raise K1ValidFovMaskError("calibration snapshot artifact changed")
        by_name[name] = artifact
        by_source[source_path] = digest
    if set(by_name) != expected_names or len(by_source) != 2:
        raise K1ValidFovMaskError("calibration snapshot artifact set changed")
    vendor_device_id = device.get("vendor_device_id")
    device_serial = device.get("device_serial")
    if not isinstance(vendor_device_id, str) or not isinstance(device_serial, str):
        raise K1ValidFovMaskError("calibration snapshot device binding is unavailable")
    digest = hashlib.sha256()
    try:
        digest.update(vendor_device_id.encode("ascii"))
        digest.update(b"\x00")
        digest.update(device_serial.encode("ascii"))
    except UnicodeEncodeError as exc:
        raise K1ValidFovMaskError("calibration snapshot device binding is invalid") from exc
    for source_path in sorted(by_source):
        digest.update(b"\x00")
        digest.update(source_path.encode("utf-8"))
        digest.update(bytes.fromhex(by_source[source_path]))
    if digest.hexdigest() != identity_sha256:
        raise K1ValidFovMaskError("calibration snapshot content identity changed")
    return manifest


def _circle_mask(width: int, height: int, cx: float, cy: float, radius: float) -> np.ndarray:
    y, x = np.ogrid[:height, :width]
    inside = (x - cx) ** 2 + (y - cy) ** 2 <= radius**2
    return np.where(inside, 255, 0).astype(np.uint8)


def _prepare_private_directory(path: Path) -> Path:
    candidate = path.expanduser()
    candidate.mkdir(mode=0o700, parents=True, exist_ok=True)
    metadata = candidate.lstat()
    if stat.S_ISLNK(metadata.st_mode) or not stat.S_ISDIR(metadata.st_mode):
        raise K1ValidFovMaskError("valid-FOV output root must be a real directory")
    root = candidate.resolve(strict=True)
    os.chmod(root, 0o700)
    return root


def _read_json_object(path: Path, root: Path, maximum_bytes: int) -> dict[str, Any]:
    metadata = _confined_regular_file(path, root)
    if not 0 < metadata.st_size <= maximum_bytes:
        raise K1ValidFovMaskError("JSON manifest is outside bounds")
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise K1ValidFovMaskError("JSON manifest is unavailable") from exc
    if not isinstance(value, dict):
        raise K1ValidFovMaskError("JSON manifest is not an object")
    return value


def _confined_regular_file(path: Path, root: Path) -> os.stat_result:
    try:
        resolved_root = root.resolve(strict=True)
        resolved = path.resolve(strict=True)
        metadata = path.lstat()
    except OSError as exc:
        raise K1ValidFovMaskError("artifact is unavailable") from exc
    if (
        stat.S_ISLNK(metadata.st_mode)
        or not stat.S_ISREG(metadata.st_mode)
        or not resolved.is_relative_to(resolved_root)
    ):
        raise K1ValidFovMaskError("artifact is not a confined regular file")
    return metadata


def _sha256_file(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 _canonical_json(value: object) -> bytes:
    return json.dumps(
        value,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
        allow_nan=False,
    ).encode("utf-8")


def _fsync_file(path: Path) -> None:
    with path.open("rb") as stream:
        os.fsync(stream.fileno())


def _fsync_directory(path: Path) -> None:
    flags = os.O_RDONLY
    if hasattr(os, "O_DIRECTORY"):
        flags |= os.O_DIRECTORY
    descriptor = os.open(path, flags)
    try:
        os.fsync(descriptor)
    finally:
        os.close(descriptor)
