lattice-common 2026.1.203

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use uuid::Uuid;

// ─── Identifiers ────────────────────────────────────────────

pub type AllocId = Uuid;
pub type NodeId = String; // xname-style: x1000c0s0b0n0
pub type TenantId = String;
pub type VClusterId = String;
pub type GroupId = u32; // dragonfly group index
pub type UserId = String; // OIDC subject
pub type LaunchId = Uuid;

pub type SessionId = Uuid;

/// An interactive session attached to a running allocation.
/// Tracked in quorum state for global session limit enforcement (F20).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Session {
    pub id: SessionId,
    pub allocation_id: AllocId,
    pub user: UserId,
    pub created_at: DateTime<Utc>,
}

// ─── Allocation ─────────────────────────────────────────────

/// The universal work unit. Replaces both Slurm jobs and K8s pods.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Allocation {
    pub id: AllocId,
    pub tenant: TenantId,
    pub project: String,
    pub vcluster: VClusterId,
    pub user: UserId,
    pub tags: HashMap<String, String>,

    // Type (single or task group)
    pub allocation_type: AllocationType,

    // What to run
    pub environment: Environment,
    pub entrypoint: String,

    // Resources
    pub resources: ResourceRequest,

    // Lifecycle
    pub lifecycle: Lifecycle,

    // Requeue
    pub requeue_policy: RequeuePolicy,
    pub max_requeue: u32,

    // Data
    pub data: DataRequirements,

    // Networking
    pub connectivity: Connectivity,

    // Dependencies
    pub depends_on: Vec<Dependency>,

    // Checkpointing
    pub checkpoint: CheckpointStrategy,

    // Telemetry
    pub telemetry_mode: TelemetryMode,

    // Liveness probe (for Unbounded/Reactive service allocations)
    pub liveness_probe: Option<LivenessProbe>,

    // State (managed by scheduler, not set by user)
    pub state: AllocationState,
    pub created_at: DateTime<Utc>,
    pub started_at: Option<DateTime<Utc>>,
    pub completed_at: Option<DateTime<Utc>>,
    pub assigned_nodes: Vec<NodeId>,

    /// DAG this allocation belongs to (set by scheduler on DAG submission)
    pub dag_id: Option<String>,
    /// Process exit code (set on completion/failure)
    pub exit_code: Option<i32>,
    /// Human-readable status message (failure reason, cancellation note, etc.)
    pub message: Option<String>,
    /// Number of times this allocation has been requeued
    pub requeue_count: u32,
    /// Number of times this allocation has been preempted
    pub preempted_count: u32,
    /// Whether to resume from checkpoint on next scheduling
    pub resume_from_checkpoint: bool,
    /// Whether this allocation requires sensitive workload isolation.
    /// Sensitive allocations get dedicated nodes, encrypted storage, audit logging,
    /// and single-session attach limits (INV-C2).
    #[serde(default)]
    pub sensitive: bool,

    // ─── Dispatch fields (INV-D6, INV-D8, INV-D11) ─────────────────
    /// Monotonic version counter for optimistic concurrency (INV-D6).
    /// Incremented on every state-changing Raft command for this allocation
    /// (SubmitAllocation, AssignNodes, UpdateAllocationState, ApplyCompletionReport,
    /// RollbackDispatch, RequeueAllocation). Separate from `requeue_count`.
    #[serde(default)]
    pub state_version: u64,
    /// Number of Dispatch Failures this allocation has experienced across
    /// potentially different nodes. Bounded by `max_dispatch_retries`; when
    /// reached, the allocation transitions to Failed rather than Pending
    /// on next RollbackDispatch (INV-D11).
    #[serde(default)]
    pub dispatch_retry_count: u32,
    /// Timestamp of the most recent ApplyCompletionReport that advanced this
    /// allocation's phase. Used by INV-D8 extended silent-sweep to detect
    /// ghosting agents (heartbeating but no progress).
    #[serde(default)]
    pub last_completion_report_at: Option<DateTime<Utc>>,
    /// Timestamp when the scheduler first assigned nodes to this allocation
    /// (i.e., when it moved from Pending to Running-with-nodes). Set by the
    /// `AssignNodes` apply-step. Used by INV-D8 silent-sweep to distinguish
    /// freshly-placed allocations from genuinely stuck ones
    /// (D-ADV-ARCH-07 fresh-allocation exemption).
    #[serde(default)]
    pub assigned_at: Option<DateTime<Utc>>,
    /// Per-node phase tracking for DEC-DISP-11 conservative multi-node
    /// aggregation. Written by the `ApplyCompletionReport` apply-step from
    /// each Completion Report's (node_id, phase). The global
    /// `AllocationState` is the conservative aggregate over these
    /// per-node phases.
    #[serde(default)]
    pub per_node_phase: HashMap<NodeId, CompletionPhase>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum AllocationState {
    /// Submitted, waiting in queue
    Pending,
    /// Data being staged
    Staging,
    /// Running on assigned nodes
    Running,
    /// Checkpoint in progress (preemption pending)
    Checkpointing,
    /// Suspended (checkpointed, waiting for resources)
    Suspended,
    /// Completed successfully
    Completed,
    /// Failed
    Failed,
    /// Cancelled by user
    Cancelled,
}

// ─── Completion Reports (INV-D4, INV-D7, INV-D12, INV-D13) ──────────

/// Local allocation phase reported by the agent (mirrors proto
/// `CompletionPhase`). The agent's per-allocation state progresses
/// monotonically through these phases; INV-D7 rejects regressions at
/// the quorum. INV-D13 uses latest-wins semantics in the agent buffer.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum CompletionPhase {
    /// Prologue in progress on this node; runtime prepared but no process yet.
    Staging,
    /// Workload Process spawned and running on this node.
    Running,
    /// Workload Process exited with exit_code 0 and epilogue complete.
    Completed,
    /// Workload Process exited non-zero, liveness probe failed, or runtime
    /// error. `reason` on the report carries the cause.
    Failed,
}

/// A per-allocation state-change record carried on a heartbeat (IP-03).
/// Delivered by the node agent to the quorum at `heartbeat_interval`
/// cadence, subject to INV-D4 idempotency, INV-D7 monotonicity, INV-D12
/// source-auth, and INV-D13 latest-wins buffering.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompletionReport {
    pub allocation_id: AllocId,
    pub phase: CompletionPhase,
    /// Set on the first Running report for this allocation on this node.
    pub pid: Option<u32>,
    /// Set for Completed/Failed.
    pub exit_code: Option<i32>,
    /// Free-form for Failed; short annotation for others.
    pub reason: Option<String>,
}

impl CompletionReport {
    /// Map this agent-side phase to a global allocation state component.
    /// Multi-node aggregation per DEC-DISP-11 is applied at the quorum;
    /// this mapping gives the per-node contribution.
    pub fn to_per_node_allocation_state(&self) -> AllocationState {
        match self.phase {
            CompletionPhase::Staging => AllocationState::Staging,
            CompletionPhase::Running => AllocationState::Running,
            CompletionPhase::Completed => AllocationState::Completed,
            CompletionPhase::Failed => AllocationState::Failed,
        }
    }
}

/// Refusal reasons from agent to dispatcher (DEC-DISP-05). Mirrors proto
/// `RefusalReason`. Dispatcher must treat unknown variants as
/// `MalformedRequest` per D-ADV-ARCH-10 resolution.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum RefusalReason {
    /// Agent temporarily saturated; retry within attempt budget.
    Busy,
    /// Agent cannot fulfil this allocation shape; rollback to Pending so
    /// the scheduler re-places on a capable node.
    UnsupportedCapability,
    /// Allocation spec is bad; transition directly to Failed.
    MalformedRequest,
    /// Agent already has this allocation (INV-D3 short-circuit).
    AlreadyRunning,
}

/// Runtime variants for executing the allocation entrypoint.
/// One Runtime is selected per allocation; they are not composed.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum RuntimeVariant {
    /// Bare-Process Runtime: spawns entrypoint directly in a cgroup scope,
    /// no mount namespace. Used when neither `uenv` nor `image` is set.
    BareProcess,
    /// Uenv Runtime: mounts squashfs and spawns via nsenter.
    Uenv,
    /// Podman Runtime: pulls OCI image and spawns via podman run + nsenter.
    Podman,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AllocationType {
    /// Single allocation
    Single,
    /// Task group (job array equivalent)
    TaskGroup {
        range_start: u32,
        range_end: u32,
        step: u32,
        max_concurrent: u32,
    },
}

// ─── Requeue Policy ─────────────────────────────────────────

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum RequeuePolicy {
    /// Allocation fails permanently on any node failure.
    Never,
    /// Requeue only on node-side failures (hardware, agent crash, partition).
    OnNodeFailure,
    /// Requeue on any failure including application crash.
    Always,
}

// ─── Software Delivery ──────────────────────────────────────

/// Image type: uenv (SquashFS) or OCI container.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
pub enum ImageType {
    #[default]
    Uenv,
    Oci,
}

/// A content-addressed reference to an image (uenv or OCI).
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ImageRef {
    /// Original user-provided spec (e.g., "prgenv-gnu/24.11:v1").
    pub spec: String,
    /// Whether this is a uenv or OCI image.
    pub image_type: ImageType,
    /// Registry URL (e.g., "jfrog.cscs.ch/uenv" or "nvcr.io").
    pub registry: String,
    /// Image name (e.g., "prgenv-gnu").
    pub name: String,
    /// Image version (e.g., "24.11").
    pub version: String,
    /// Original tag preserved for audit (INV-SD1).
    pub original_tag: String,
    /// Content hash (empty if deferred resolution).
    pub sha256: String,
    /// Image size in bytes.
    pub size_bytes: u64,
    /// Mount point (e.g., "/user-environment" for uenv, "/" for container).
    pub mount_point: String,
    /// If true, resolution is deferred to scheduler time (INV-SD4).
    pub resolve_on_schedule: bool,
}

/// Operation to apply to an environment variable.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
pub enum EnvOp {
    Prepend,
    Append,
    #[default]
    Set,
    Unset,
}

/// A single patch to an environment variable (view activation).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnvPatch {
    /// Variable name (e.g., "PATH").
    pub variable: String,
    /// Operation to apply.
    pub op: EnvOp,
    /// Value to use (ignored for Unset).
    pub value: String,
    /// Separator for Prepend/Append (default ":").
    pub separator: String,
}

impl Default for EnvPatch {
    fn default() -> Self {
        Self {
            variable: String::new(),
            op: EnvOp::Set,
            value: String::new(),
            separator: ":".to_string(),
        }
    }
}

/// A named view definition from image metadata.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ViewDef {
    pub name: String,
    pub description: String,
    pub patches: Vec<EnvPatch>,
}

/// Metadata extracted from a resolved image.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ImageMetadata {
    pub name: String,
    pub description: String,
    pub mount_point: String,
    pub views: Vec<ViewDef>,
    pub default_view: Option<String>,
}

/// OCI container specification (when running via Sarus/Podman).
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ContainerSpec {
    /// EDF inheritance chain (base environment names).
    pub base_environments: Vec<String>,
    /// Additional bind mounts.
    pub mounts: Vec<MountSpec>,
    /// CDI device specs (e.g., "nvidia.com/gpu=all").
    pub devices: Vec<String>,
    /// Working directory inside the container.
    pub workdir: String,
    /// Whether the container rootfs is writable.
    pub writable: bool,
    /// Additional environment variables.
    pub env: Vec<(String, String)>,
    /// OCI annotations.
    pub annotations: Vec<(String, String)>,
}

/// A bind mount specification.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct MountSpec {
    pub source: String,
    pub target: String,
    /// Mount options (e.g., "ro", "rw").
    pub options: String,
}

// ─── Environment ────────────────────────────────────────────

/// Software environment for an allocation. Replaces the legacy flat struct
/// with content-addressed images, env patches from views, and optional
/// OCI container spec.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Environment {
    /// Resolved image references (uenv and/or OCI).
    pub images: Vec<ImageRef>,
    /// Environment variable patches (from view activation / EDF).
    pub env_patches: Vec<EnvPatch>,
    /// CDI device specs (e.g., "nvidia.com/gpu=all").
    pub devices: Vec<String>,
    /// Additional bind mounts.
    pub mounts: Vec<MountSpec>,
    /// Full container spec (set when running OCI images).
    pub container: Option<ContainerSpec>,
    /// Whether the container rootfs is writable.
    pub writable: bool,
    /// For sensitive: require signed images.
    pub sign_required: bool,
    /// Require vulnerability scan before scheduling (sensitive-workloads).
    pub scan_required: bool,
    /// Restrict to approved base images only (sensitive-workloads).
    pub approved_bases_only: bool,
}

// ─── Resources ──────────────────────────────────────────────

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceRequest {
    /// Number of nodes (exact or range)
    pub nodes: NodeCount,
    /// Hardware constraints
    pub constraints: ResourceConstraints,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum NodeCount {
    Exact(u32),
    /// Elastic node range. Invariants: min >= 1, max >= min.
    Range {
        min: u32,
        max: u32,
    },
}

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ResourceConstraints {
    pub gpu_type: Option<String>,
    pub features: Vec<String>,
    /// Topology hint: "tight" (fewest groups), "spread" (max bandwidth), "any"
    pub topology: Option<TopologyHint>,
    /// Feature count requirements (e.g., at least 4 nodes with "nvme_scratch")
    pub feature_counts: HashMap<String, u32>,
    /// Require nodes with unified memory (GH200/MI300A)
    #[serde(default)]
    pub require_unified_memory: bool,
    /// Prefer allocating CPUs+GPUs within the same NUMA domain
    #[serde(default)]
    pub prefer_same_numa: bool,
    /// Allow CXL-attached memory to satisfy capacity requirements
    #[serde(default)]
    pub allow_cxl_memory: bool,
    /// Memory binding policy (numactl)
    pub memory_policy: Option<MemoryPolicy>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum TopologyHint {
    Tight,  // pack into fewest dragonfly groups
    Spread, // spread across groups for max bisection bandwidth
    Any,    // no preference
}

// ─── Lifecycle ──────────────────────────────────────────────

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Lifecycle {
    pub lifecycle_type: LifecycleType,
    pub preemption_class: u8, // 0 = lowest priority, higher = harder to preempt
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum LifecycleType {
    /// Traditional batch job with walltime
    Bounded { walltime: chrono::Duration },
    /// Long-running service (inference endpoint, monitoring)
    Unbounded,
    /// Autoscaling based on metrics
    Reactive {
        min_nodes: u32,
        max_nodes: u32,
        metric: String,
        target: String,
    },
}

// ─── Data ───────────────────────────────────────────────────

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct DataRequirements {
    pub mounts: Vec<DataMount>,
    /// Use sane defaults (home dir, scratch, output dir)
    pub use_defaults: bool,
    /// Per-node scratch size hint
    pub scratch_per_node: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DataMount {
    pub source: String, // s3://... or nfs://...
    pub target: String, // mount point inside allocation
    pub access: DataAccess,
    pub tier_hint: Option<StorageTier>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum DataAccess {
    ReadOnly,
    ReadWrite,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum StorageTier {
    Hot,
    Warm,
    Cold,
}

// ─── Connectivity ───────────────────────────────────────────

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Connectivity {
    /// Shared network domain name — allocations in the same domain can communicate
    pub network_domain: Option<String>,
    /// Exposed endpoints (for services)
    pub expose: Vec<ServiceEndpoint>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServiceEndpoint {
    pub name: String,
    pub port: u16,
    pub protocol: Option<String>,
}

/// A registered service endpoint — published into the service registry
/// when an allocation with `expose` endpoints transitions to Running.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RegisteredEndpoint {
    pub allocation_id: AllocId,
    pub tenant: TenantId,
    pub nodes: Vec<NodeId>,
    pub port: u16,
    pub protocol: Option<String>,
}

/// A service registry entry — groups all endpoints for a named service.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ServiceRegistryEntry {
    pub endpoints: Vec<RegisteredEndpoint>,
}

// ─── Liveness Probe ─────────────────────────────────────────

/// A liveness probe checks whether a service allocation's workload is healthy.
/// If the probe fails `failure_threshold` consecutive times, the allocation is
/// marked Failed (and the reconciler can requeue it).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LivenessProbe {
    /// The type of probe to run.
    pub probe_type: ProbeType,
    /// How often (in seconds) to run the probe.
    pub period_secs: u32,
    /// Seconds to wait before the first probe (gives the workload time to start).
    pub initial_delay_secs: u32,
    /// Number of consecutive failures before marking the allocation as Failed.
    pub failure_threshold: u32,
    /// Seconds to wait for the probe to succeed before counting as a failure.
    pub timeout_secs: u32,
}

impl Default for LivenessProbe {
    fn default() -> Self {
        Self {
            probe_type: ProbeType::Tcp { port: 8080 },
            period_secs: 30,
            initial_delay_secs: 10,
            failure_threshold: 3,
            timeout_secs: 5,
        }
    }
}

/// The mechanism used to check workload liveness.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ProbeType {
    /// Open a TCP connection to the given port on the allocation's node.
    Tcp { port: u16 },
    /// Send an HTTP GET request and expect a 2xx response.
    Http { port: u16, path: String },
}

// ─── Network Domain ────────────────────────────────────────

/// A network domain groups allocations that need L3 reachability via Slingshot VNI.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NetworkDomain {
    pub name: String,
    pub tenant: TenantId,
    /// Assigned Slingshot VNI
    pub vni: u32,
    pub state: NetworkDomainState,
    /// Allocation IDs currently in this domain
    pub member_allocations: Vec<AllocId>,
    pub created_at: DateTime<Utc>,
    /// Grace deadline after last member completes (for DAG domain persistence)
    pub grace_deadline: Option<DateTime<Utc>>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum NetworkDomainState {
    /// Active, accepting new members
    Active,
    /// Last member completed, grace timer running
    Draining,
    /// VNI released back to pool
    Released,
}

// ─── Dependencies ───────────────────────────────────────────

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Dependency {
    /// Reference to another allocation (by ID or name)
    pub ref_id: String,
    /// Condition for dependency satisfaction
    pub condition: DependencyCondition,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum DependencyCondition {
    /// afterok — only if dependency succeeded
    Success,
    /// afternotok — only if dependency failed
    Failure,
    /// afterany — regardless of outcome
    Any,
    /// aftercorr — corresponding task in task group
    Corresponding,
    /// singleton — only one allocation with this name runs at a time
    Mutex,
}

// ─── Checkpointing ─────────────────────────────────────────

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum CheckpointStrategy {
    /// Scheduler decides based on cost function
    Auto,
    /// Application manages its own checkpointing
    Manual,
    /// Non-checkpointable (treated as non-preemptible)
    None,
}

// ─── Telemetry ──────────────────────────────────────────────

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub enum TelemetryMode {
    /// 30s aggregation, bicubic smoothing, low overhead
    #[default]
    Prod,
    /// 1s raw streams, full profiling, higher overhead
    Debug { duration_seconds: u64 },
    /// Access logging for compliance, moderate overhead
    Audit,
}

// ─── Observability: Attach ─────────────────────────────────

/// An active interactive terminal session attached to a running allocation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AttachSession {
    pub session_id: Uuid,
    pub allocation_id: AllocId,
    pub node_id: NodeId,
    pub user: UserId,
    pub command: String,
    pub started_at: DateTime<Utc>,
    pub ended_at: Option<DateTime<Utc>>,
}

// ─── Observability: Logs ───────────────────────────────────

/// A single log entry from an allocation's stdout/stderr.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LogEntry {
    pub node_id: NodeId,
    pub stream: LogStream,
    pub data: Vec<u8>,
    pub timestamp: DateTime<Utc>,
}

/// Which output stream a log entry came from.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum LogStream {
    Stdout,
    Stderr,
}

/// Configuration for log capture on a running allocation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LogConfig {
    /// Size of the per-allocation ring buffer on each node (bytes).
    pub ring_buffer_size: u64,
    /// Whether to persist logs to S3.
    pub s3_persistence: bool,
    /// Retention duration for persisted logs.
    pub retention: Option<chrono::Duration>,
}

impl Default for LogConfig {
    fn default() -> Self {
        Self {
            ring_buffer_size: 64 * 1024 * 1024, // 64 MB
            s3_persistence: true,
            retention: None, // use system default
        }
    }
}

// ─── Observability: Metrics ────────────────────────────────

/// Per-node metrics snapshot.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeMetricsSnapshot {
    pub node_id: NodeId,
    pub timestamp: DateTime<Utc>,
    pub cpu_utilization: f64,
    pub memory_used_bytes: u64,
    pub memory_total_bytes: u64,
    pub network_tx_bytes_per_sec: f64,
    pub network_rx_bytes_per_sec: f64,
    pub io_read_bytes_per_sec: f64,
    pub io_write_bytes_per_sec: f64,
    pub io_latency_p99_us: f64,
    pub gpus: Vec<GpuMetricsSnapshot>,
}

/// Per-GPU metrics snapshot.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuMetricsSnapshot {
    pub index: u32,
    pub utilization: f64,
    pub memory_used_bytes: u64,
    pub memory_total_bytes: u64,
    pub power_draw_watts: f64,
    pub temperature_celsius: f64,
    pub ecc_errors: u64,
}

/// Aggregated metrics across all nodes in an allocation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AllocationMetricsSummary {
    pub allocation_id: AllocId,
    pub timestamp: DateTime<Utc>,
    pub gpu_utilization_mean: f64,
    pub cpu_utilization_mean: f64,
    pub memory_used_bytes: u64,
    pub memory_total_bytes: u64,
    pub gpu_memory_used_bytes: u64,
    pub gpu_memory_total_bytes: u64,
    pub network_tx_bytes_per_sec: f64,
    pub network_rx_bytes_per_sec: f64,
    pub io_read_bytes_per_sec: f64,
    pub io_write_bytes_per_sec: f64,
    pub io_latency_p99_us: f64,
}

// ─── Observability: Diagnostics ────────────────────────────

/// Network diagnostics for a running allocation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NetworkDiagnostics {
    pub group_span: u32,
    pub groups: Vec<GroupId>,
    pub csig_congestion_avg: f64,
    pub inter_node_bandwidth_gbps: f64,
    pub target_bandwidth_gbps: f64,
    pub node_pairs: Vec<NodePairBandwidth>,
    pub nvlink_throughput_gbps: f64,
    pub network_errors: u64,
}

/// Measured bandwidth between a pair of nodes.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodePairBandwidth {
    pub source_node: NodeId,
    pub target_node: NodeId,
    pub bandwidth_gbps: f64,
    pub latency_us: f64,
}

/// Storage diagnostics for a running allocation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StorageDiagnostics {
    pub mounts: Vec<MountDiagnostics>,
}

/// Per-mount storage diagnostics.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MountDiagnostics {
    pub mount_path: String,
    pub mount_type: String,
    pub read_throughput_gbps: f64,
    pub write_throughput_gbps: f64,
    pub qos_floor_gbps: f64,
    pub latency_p50_us: f64,
    pub latency_p95_us: f64,
    pub latency_p99_us: f64,
    pub iops_read: f64,
    pub iops_write: f64,
    pub health: String,
}

// ─── Observability: Alerts ─────────────────────────────────

/// A threshold alert generated by a node agent.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MetricAlert {
    pub node_id: NodeId,
    pub metric_name: String,
    pub current_value: f64,
    pub threshold: f64,
    pub severity: AlertSeverity,
    pub message: String,
    pub timestamp: DateTime<Utc>,
}

/// Severity level for metric alerts.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AlertSeverity {
    Info,
    Warning,
    Critical,
}

// ─── Node ───────────────────────────────────────────────────

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Node {
    pub id: NodeId,
    pub group: GroupId,
    pub capabilities: NodeCapabilities,
    pub state: NodeState,
    pub owner: Option<NodeOwnership>,
    /// Conformance fingerprint (hash of OS image, kernel, driver versions)
    pub conformance_fingerprint: Option<String>,
    /// Last heartbeat received from this node's agent
    pub last_heartbeat: Option<DateTime<Utc>>,
    /// Monotonic version counter for ownership changes.
    /// Incremented on every ClaimNode/ReleaseNode. Heartbeats must carry
    /// the current owner_version to be accepted (ADV-06).
    #[serde(default)]
    pub owner_version: u64,

    // ─── Dispatch fields (INV-D1, INV-D11, INV-D14, INV-D5) ─────────
    /// Reachable gRPC endpoint for this node's agent (host:port).
    /// Required for the node to participate in scheduling (INV-D1).
    /// Must match a SAN in the agent's workload certificate (INV-D14).
    /// Empty string means "unregistered" — node is scheduler-invisible (INV-D2).
    #[serde(default)]
    pub agent_address: String,
    /// Consecutive dispatch failures attributed to this node. Incremented
    /// atomically with RollbackDispatch (INV-D11). Reset to 0 on any
    /// successful dispatch on this node. Threshold `max_node_dispatch_failures`
    /// transitions the node to Degraded (subject to cluster-wide ratio guard).
    #[serde(default)]
    pub consecutive_dispatch_failures: u32,
    /// Timestamp when the node entered Degraded via INV-D11. Used by
    /// DEC-DISP-01 auto-probe-recovery — after `degraded_probe_interval`
    /// since this timestamp, the node is probed back to Ready with halved counter.
    #[serde(default)]
    pub degraded_at: Option<DateTime<Utc>>,
    /// Agent has completed boot but is still reattaching to surviving
    /// Workload Processes (DEC-DISP-10). Raft-committed via extended
    /// Command::RecordHeartbeat. Silent-sweep (INV-D8) honours this flag
    /// subject to the lifecycle rules in INV-D5.
    #[serde(default)]
    pub reattach_in_progress: bool,
    /// Timestamp of the first heartbeat that set `reattach_in_progress = true`
    /// after the most recent re-registration. Used by INV-D5 lifecycle:
    /// the grace timer is absolute from this timestamp; it does not reset
    /// on subsequent true-valued heartbeats.
    #[serde(default)]
    pub reattach_first_set_at: Option<DateTime<Utc>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeCapabilities {
    pub gpu_type: Option<String>,
    pub gpu_count: u32,
    pub cpu_cores: u32,
    pub memory_gb: u64,
    pub features: Vec<String>,
    /// Detailed GPU topology (interconnect, NIC affinity)
    pub gpu_topology: Option<GpuTopology>,
    /// Memory topology (NUMA, CXL, unified memory)
    pub memory_topology: Option<MemoryTopology>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum NodeState {
    /// Node exists in inventory but has never reported
    Unknown,
    /// OpenCHAMI booting/reimaging the node
    Booting,
    /// Healthy, agent reporting, available for scheduling
    Ready,
    /// Heartbeat missed or minor issue detected
    Degraded { reason: String },
    /// Confirmed failure, grace period expired
    Down { reason: String },
    /// Operator or scheduler requested drain, waiting for allocations to finish
    Draining,
    /// All allocations completed after drain
    Drained,
    /// Boot failure or unrecoverable hardware error
    Failed { reason: String },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeOwnership {
    pub tenant: TenantId,
    pub vcluster: VClusterId,
    pub allocation: AllocId,
    /// For sensitive: the specific user who claimed this node
    pub claimed_by: Option<UserId>,
    /// Home vCluster (permanent assignment) vs current (may be borrowed)
    pub is_borrowed: bool,
}

// ─── Tenant ─────────────────────────────────────────────────

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Tenant {
    pub id: TenantId,
    pub name: String,
    pub quota: TenantQuota,
    pub isolation_level: IsolationLevel,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TenantQuota {
    /// Maximum nodes this tenant can use simultaneously
    pub max_nodes: u32,
    /// Target fair-share fraction (0.0 - 1.0)
    pub fair_share_target: f64,
    /// GPU-hours budget (None = unlimited)
    pub gpu_hours_budget: Option<f64>,
    /// Node-hours budget (None = unlimited). Universal metric that works
    /// for both GPU and CPU-only nodes. When both gpu_hours_budget and
    /// node_hours_budget are set, the worse utilization drives the penalty.
    #[serde(default)]
    pub node_hours_budget: Option<f64>,
    /// Maximum concurrent allocations (hard limit per quota-enforcement)
    pub max_concurrent_allocations: Option<u32>,
    /// Burst allowance multiplier (e.g. 1.5 = allow up to 150% of fair share
    /// when resources are idle). Burst allocations are first to be preempted.
    /// None means no bursting beyond fair_share_target.
    #[serde(default)]
    pub burst_allowance: Option<f64>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum IsolationLevel {
    /// Standard HPC: shared nodes possible, elastic borrowing
    Standard,
    /// Strict: dedicated nodes, no sharing, audit logging
    Strict,
}

// ─── VCluster ───────────────────────────────────────────────

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VCluster {
    pub id: VClusterId,
    pub name: String,
    pub tenant: TenantId,
    pub scheduler_type: SchedulerType,
    pub cost_weights: CostWeights,
    /// Nodes permanently assigned to this vCluster
    pub dedicated_nodes: Vec<NodeId>,
    /// Can borrow from other vClusters' idle nodes?
    pub allow_borrowing: bool,
    /// Can lend idle nodes to other vClusters?
    pub allow_lending: bool,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum SchedulerType {
    /// Priority + backfill + topology-aware (traditional HPC)
    HpcBackfill,
    /// Bin-packing + autoscale (services)
    ServiceBinPack,
    /// User-claim reservation (sensitive)
    SensitiveReservation,
    /// FIFO, short-lived (interactive sessions)
    InteractiveFifo,
}

/// Weights for the composite cost function. Weights are relative;
/// normalization (sum to 1.0) is optional — the solver normalizes internally.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CostWeights {
    pub priority: f64,
    pub wait_time: f64,
    pub fair_share: f64,
    pub topology: f64,
    pub data_readiness: f64,
    pub backlog: f64,
    pub energy: f64,
    pub checkpoint_efficiency: f64,
    pub conformance: f64,
}

impl Default for CostWeights {
    fn default() -> Self {
        // Default: balanced HPC profile
        Self {
            priority: 0.20,
            wait_time: 0.20,
            fair_share: 0.20,
            topology: 0.15,
            data_readiness: 0.10,
            backlog: 0.05,
            energy: 0.00,
            checkpoint_efficiency: 0.00,
            conformance: 0.10,
        }
    }
}

// ─── GPU Topology ───────────────────────────────────────────

/// Intra-node GPU topology: devices, interconnect links, and NIC affinity.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuTopology {
    pub devices: Vec<GpuDevice>,
    /// Mapping of GPU index → closest NIC index (for Slingshot/UE affinity)
    pub nic_affinity: HashMap<u32, u32>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuDevice {
    pub index: u32,
    pub vendor: GpuVendor,
    pub model: String,
    pub memory_bytes: u64,
    pub links: Vec<GpuLink>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuLink {
    pub peer_index: u32,
    pub link_type: GpuLinkType,
    pub bandwidth_gbps: f64,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum GpuLinkType {
    NVLink,
    NVSwitch,
    InfinityFabric,
    PCIe,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum GpuVendor {
    Nvidia,
    Amd,
}

// ─── Memory Topology ────────────────────────────────────────

/// Intra-node memory topology: NUMA domains, CXL tiers, and unified memory.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryTopology {
    pub domains: Vec<MemoryDomain>,
    pub interconnects: Vec<MemoryInterconnect>,
    pub total_capacity_bytes: u64,
}

/// A memory domain (NUMA node, HBM bank, CXL-attached pool, or unified).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryDomain {
    pub id: u32,
    pub domain_type: MemoryDomainType,
    pub capacity_bytes: u64,
    pub numa_node: Option<u32>,
    pub attached_cpus: Vec<u32>,
    pub attached_gpus: Vec<u32>,
}

/// Interconnect link between two memory domains.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryInterconnect {
    pub domain_a: u32,
    pub domain_b: u32,
    pub link_type: MemoryLinkType,
    pub bandwidth_gbps: f64,
    pub latency_ns: u64,
}

/// Type of memory in a domain.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum MemoryDomainType {
    Dram,
    Hbm,
    CxlAttached,
    Unified,
}

/// Type of link between memory domains.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum MemoryLinkType {
    NumaLink,
    CxlSwitch,
    CoherentFabric,
}

/// Memory binding policy for numactl.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum MemoryPolicy {
    Local,
    Interleave,
    Preferred,
    Bind,
}

// ─── Topology ───────────────────────────────────────────────

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TopologyModel {
    pub groups: Vec<TopologyGroup>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TopologyGroup {
    pub id: GroupId,
    pub nodes: Vec<NodeId>,
    /// Adjacent groups (lower inter-group hop cost)
    pub adjacent_groups: Vec<GroupId>,
}

// ─── State Machine Validation ──────────────────────────────

impl AllocationState {
    /// Returns true if transitioning from `self` to `target` is a valid state change.
    ///
    /// Valid transitions:
    /// - Pending → Staging, Running, Cancelled, Failed
    /// - Staging → Running, Failed, Cancelled
    /// - Running → Checkpointing, Completed, Failed, Cancelled
    /// - Checkpointing → Suspended, Failed, Cancelled
    /// - Suspended → Pending (requeue), Cancelled, Failed
    /// - Terminal states (Completed, Failed, Cancelled) → nothing
    pub fn can_transition_to(&self, target: &AllocationState) -> bool {
        matches!(
            (self, target),
            (AllocationState::Pending, AllocationState::Staging)
                | (AllocationState::Pending, AllocationState::Running)
                | (AllocationState::Pending, AllocationState::Cancelled)
                | (AllocationState::Pending, AllocationState::Failed)
                | (AllocationState::Staging, AllocationState::Running)
                | (AllocationState::Staging, AllocationState::Failed)
                | (AllocationState::Staging, AllocationState::Cancelled)
                | (AllocationState::Running, AllocationState::Checkpointing)
                | (AllocationState::Running, AllocationState::Completed)
                | (AllocationState::Running, AllocationState::Failed)
                | (AllocationState::Running, AllocationState::Cancelled)
                | (AllocationState::Checkpointing, AllocationState::Suspended)
                | (AllocationState::Checkpointing, AllocationState::Failed)
                | (AllocationState::Checkpointing, AllocationState::Cancelled)
                | (AllocationState::Suspended, AllocationState::Pending)
                | (AllocationState::Suspended, AllocationState::Cancelled)
                | (AllocationState::Suspended, AllocationState::Failed)
                // Service reconciliation: failed services can be requeued
                | (AllocationState::Failed, AllocationState::Pending)
        )
    }

    /// Returns true if this is a terminal state (no further transitions possible).
    pub fn is_terminal(&self) -> bool {
        matches!(
            self,
            AllocationState::Completed | AllocationState::Failed | AllocationState::Cancelled
        )
    }
}

impl NodeState {
    /// Returns true if transitioning from `self` to `target` is a valid state change.
    ///
    /// Valid transitions:
    /// - Unknown → Booting, Failed
    /// - Booting → Ready, Failed
    /// - Ready → Degraded, Draining, Down
    /// - Degraded → Ready, Down, Draining
    /// - Down → Booting (reimaging), Failed
    /// - Draining → Drained
    /// - Drained → Ready (undrain), Booting (reimage)
    /// - Failed → Booting (reimage)
    pub fn can_transition_to(&self, target: &NodeState) -> bool {
        matches!(
            (self, target),
            (NodeState::Unknown, NodeState::Booting)
                | (NodeState::Unknown, NodeState::Failed { .. })
                | (NodeState::Booting, NodeState::Ready)
                | (NodeState::Booting, NodeState::Failed { .. })
                | (NodeState::Ready, NodeState::Degraded { .. })
                | (NodeState::Ready, NodeState::Draining)
                | (NodeState::Ready, NodeState::Down { .. })
                | (NodeState::Degraded { .. }, NodeState::Ready)
                | (NodeState::Degraded { .. }, NodeState::Down { .. })
                | (NodeState::Degraded { .. }, NodeState::Draining)
                | (NodeState::Down { .. }, NodeState::Ready)
                | (NodeState::Down { .. }, NodeState::Booting)
                | (NodeState::Down { .. }, NodeState::Failed { .. })
                | (NodeState::Draining, NodeState::Drained)
                | (NodeState::Drained, NodeState::Ready)
                | (NodeState::Drained, NodeState::Booting)
                | (NodeState::Failed { .. }, NodeState::Booting)
        )
    }

    /// Returns true if the node can accept workloads.
    pub fn is_operational(&self) -> bool {
        matches!(self, NodeState::Ready | NodeState::Degraded { .. })
    }
}

impl NetworkDomainState {
    /// Returns true if transitioning from `self` to `target` is a valid state change.
    ///
    /// Valid transitions:
    /// - Active → Draining
    /// - Draining → Active (new member joins), Released (grace expired)
    /// - Released → nothing (terminal)
    pub fn can_transition_to(&self, target: &NetworkDomainState) -> bool {
        matches!(
            (self, target),
            (NetworkDomainState::Active, NetworkDomainState::Draining)
                | (NetworkDomainState::Draining, NetworkDomainState::Active)
                | (NetworkDomainState::Draining, NetworkDomainState::Released)
        )
    }
}

// ─── MPI Process Management ─────────────────────────────────

/// PMI mode for MPI process launch.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum PmiMode {
    /// Native PMI-2 wire protocol (default, no external dependencies).
    #[default]
    Pmi2,
    /// OpenPMIx sidecar (requires `pmix` feature on node agent).
    Pmix,
}

/// CXI (Cassini eXtended Interface) credentials for Slingshot fabric access.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CxiCredentials {
    pub vni: u32,
    pub auth_key: Vec<u8>,
    pub svc_id: u32,
}

/// Peer node agent info for cross-node MPI coordination.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PeerInfo {
    pub node_id: NodeId,
    pub grpc_address: String,
    pub first_rank: u32,
    pub num_ranks: u32,
}

/// Assignment of ranks to a single node.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeRankAssignment {
    pub node_id: NodeId,
    pub first_rank: u32,
    pub num_ranks: u32,
}

/// Layout of MPI ranks across nodes.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RankLayout {
    pub total_ranks: u32,
    pub tasks_per_node: u32,
    pub node_assignments: Vec<NodeRankAssignment>,
}

impl RankLayout {
    /// Compute a rank layout from a node list and tasks-per-node.
    pub fn compute(nodes: &[NodeId], tasks_per_node: u32) -> Self {
        let mut assignments = Vec::with_capacity(nodes.len());
        let mut rank = 0u32;
        for node_id in nodes {
            assignments.push(NodeRankAssignment {
                node_id: node_id.clone(),
                first_rank: rank,
                num_ranks: tasks_per_node,
            });
            rank += tasks_per_node;
        }
        RankLayout {
            total_ranks: rank,
            tasks_per_node,
            node_assignments: assignments,
        }
    }
}

/// Exit status of a single MPI rank.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RankExitStatus {
    pub rank: u32,
    pub exit_code: Option<i32>,
    pub signal: Option<i32>,
}

/// Full specification for an MPI launch, computed by the API server before fan-out.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LaunchSpec {
    pub launch_id: LaunchId,
    pub allocation_id: AllocId,
    pub entrypoint: String,
    pub args: Vec<String>,
    pub env: HashMap<String, String>,
    pub rank_layout: RankLayout,
    pub pmi_mode: PmiMode,
    pub cxi_credentials: Option<CxiCredentials>,
}

// ─── Software Delivery Utilities ────────────────────────────

/// Apply a list of environment patches to an env var map, in declaration order (INV-SD7).
///
/// - Prepend: `var = value + separator + existing` (or just `value` if unset)
/// - Append: `var = existing + separator + value` (or just `value` if unset)
/// - Set: `var = value`
/// - Unset: remove `var`
pub fn apply_env_patches(patches: &[EnvPatch], env: &mut HashMap<String, String>) {
    for patch in patches {
        match patch.op {
            EnvOp::Set => {
                env.insert(patch.variable.clone(), patch.value.clone());
            }
            EnvOp::Unset => {
                env.remove(&patch.variable);
            }
            EnvOp::Prepend => {
                let existing = env.get(&patch.variable).cloned().unwrap_or_default();
                let new_val = if existing.is_empty() {
                    patch.value.clone()
                } else {
                    format!("{}{}{}", patch.value, patch.separator, existing)
                };
                env.insert(patch.variable.clone(), new_val);
            }
            EnvOp::Append => {
                let existing = env.get(&patch.variable).cloned().unwrap_or_default();
                let new_val = if existing.is_empty() {
                    patch.value.clone()
                } else {
                    format!("{}{}{}", existing, patch.separator, patch.value)
                };
                env.insert(patch.variable.clone(), new_val);
            }
        }
    }
}

/// Check that no mount targets overlap (prefix-shadow each other, INV-SD3).
///
/// Two targets overlap when one is a prefix of the other followed by `/` or an
/// exact match. For example, `/opt` overlaps `/opt/env` (prefix shadowing).
pub fn check_mount_overlap(targets: &[&str]) -> Result<(), String> {
    let mut sorted: Vec<&str> = targets.to_vec();
    sorted.sort_by_key(|t| t.len());

    for i in 0..sorted.len() {
        for j in (i + 1)..sorted.len() {
            let short = sorted[i];
            let long = sorted[j];
            if short == long {
                return Err(format!("duplicate mount target: {short}"));
            }
            // Check prefix: /opt shadows /opt/env
            if let Some(rest) = long.strip_prefix(short) {
                if rest.starts_with('/') || short.ends_with('/') {
                    return Err(format!("mount target overlap: {short} shadows {long}"));
                }
            }
        }
    }
    Ok(())
}

// ─── Tests ─────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;

    // ── Software Delivery type tests ──

    #[test]
    fn environment_default_has_empty_vecs_and_false_bools() {
        let env = Environment::default();
        assert!(env.images.is_empty());
        assert!(env.env_patches.is_empty());
        assert!(env.devices.is_empty());
        assert!(env.mounts.is_empty());
        assert!(env.container.is_none());
        assert!(!env.writable);
        assert!(!env.sign_required);
        assert!(!env.scan_required);
        assert!(!env.approved_bases_only);
    }

    #[test]
    fn image_ref_default() {
        let img = ImageRef::default();
        assert!(img.spec.is_empty());
        assert_eq!(img.image_type, ImageType::Uenv);
        assert!(img.sha256.is_empty());
        assert_eq!(img.size_bytes, 0);
        assert!(!img.resolve_on_schedule);
    }

    #[test]
    fn image_ref_deferred_has_empty_sha256() {
        let img = ImageRef {
            spec: "prgenv-gnu/24.11:v1".into(),
            resolve_on_schedule: true,
            ..ImageRef::default()
        };
        assert!(img.sha256.is_empty());
        assert!(img.resolve_on_schedule);
    }

    #[test]
    fn env_patch_default_separator_is_colon() {
        let patch = EnvPatch::default();
        assert_eq!(patch.separator, ":");
        assert_eq!(patch.op, EnvOp::Set);
    }

    #[test]
    fn container_spec_default() {
        let cs = ContainerSpec::default();
        assert!(cs.base_environments.is_empty());
        assert!(cs.mounts.is_empty());
        assert!(cs.devices.is_empty());
        assert!(cs.workdir.is_empty());
        assert!(!cs.writable);
    }

    #[test]
    fn mount_spec_default() {
        let ms = MountSpec::default();
        assert!(ms.source.is_empty());
        assert!(ms.target.is_empty());
        assert!(ms.options.is_empty());
    }

    #[test]
    fn view_def_default() {
        let vd = ViewDef::default();
        assert!(vd.name.is_empty());
        assert!(vd.patches.is_empty());
    }

    #[test]
    fn image_metadata_default() {
        let im = ImageMetadata::default();
        assert!(im.views.is_empty());
        assert!(im.default_view.is_none());
    }

    #[test]
    fn apply_env_patches_set() {
        let patches = vec![EnvPatch {
            variable: "FOO".into(),
            op: EnvOp::Set,
            value: "bar".into(),
            separator: ":".into(),
        }];
        let mut env = HashMap::new();
        apply_env_patches(&patches, &mut env);
        assert_eq!(env.get("FOO").unwrap(), "bar");
    }

    #[test]
    fn apply_env_patches_unset() {
        let patches = vec![EnvPatch {
            variable: "FOO".into(),
            op: EnvOp::Unset,
            value: String::new(),
            separator: ":".into(),
        }];
        let mut env = HashMap::new();
        env.insert("FOO".into(), "bar".into());
        apply_env_patches(&patches, &mut env);
        assert!(!env.contains_key("FOO"));
    }

    #[test]
    fn apply_env_patches_prepend() {
        let patches = vec![EnvPatch {
            variable: "PATH".into(),
            op: EnvOp::Prepend,
            value: "/new/bin".into(),
            separator: ":".into(),
        }];
        let mut env = HashMap::new();
        env.insert("PATH".into(), "/usr/bin".into());
        apply_env_patches(&patches, &mut env);
        assert_eq!(env.get("PATH").unwrap(), "/new/bin:/usr/bin");
    }

    #[test]
    fn apply_env_patches_append() {
        let patches = vec![EnvPatch {
            variable: "PATH".into(),
            op: EnvOp::Append,
            value: "/extra/bin".into(),
            separator: ":".into(),
        }];
        let mut env = HashMap::new();
        env.insert("PATH".into(), "/usr/bin".into());
        apply_env_patches(&patches, &mut env);
        assert_eq!(env.get("PATH").unwrap(), "/usr/bin:/extra/bin");
    }

    #[test]
    fn apply_env_patches_prepend_to_empty() {
        let patches = vec![EnvPatch {
            variable: "NEW_VAR".into(),
            op: EnvOp::Prepend,
            value: "/opt/lib".into(),
            separator: ":".into(),
        }];
        let mut env = HashMap::new();
        apply_env_patches(&patches, &mut env);
        assert_eq!(env.get("NEW_VAR").unwrap(), "/opt/lib");
    }

    #[test]
    fn apply_env_patches_declaration_order() {
        let patches = vec![
            EnvPatch {
                variable: "X".into(),
                op: EnvOp::Set,
                value: "first".into(),
                separator: ":".into(),
            },
            EnvPatch {
                variable: "X".into(),
                op: EnvOp::Set,
                value: "second".into(),
                separator: ":".into(),
            },
        ];
        let mut env = HashMap::new();
        apply_env_patches(&patches, &mut env);
        assert_eq!(env.get("X").unwrap(), "second");
    }

    #[test]
    fn check_mount_overlap_no_overlap() {
        assert!(check_mount_overlap(&["/opt", "/usr", "/mnt/data"]).is_ok());
    }

    #[test]
    fn check_mount_overlap_prefix_shadow() {
        let result = check_mount_overlap(&["/opt", "/opt/env"]);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("shadows"));
    }

    #[test]
    fn check_mount_overlap_duplicate() {
        let result = check_mount_overlap(&["/opt", "/opt"]);
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("duplicate"));
    }

    #[test]
    fn check_mount_overlap_no_false_positive_on_shared_prefix() {
        // /opt and /optional should NOT overlap (no / separator)
        assert!(check_mount_overlap(&["/opt", "/optional"]).is_ok());
    }

    #[test]
    fn check_mount_overlap_empty_list() {
        assert!(check_mount_overlap(&[]).is_ok());
    }

    // ── AllocationState transition tests ──

    #[test]
    fn pending_can_transition_to_staging() {
        assert!(AllocationState::Pending.can_transition_to(&AllocationState::Staging));
    }

    #[test]
    fn pending_can_transition_to_running() {
        assert!(AllocationState::Pending.can_transition_to(&AllocationState::Running));
    }

    #[test]
    fn pending_can_transition_to_cancelled() {
        assert!(AllocationState::Pending.can_transition_to(&AllocationState::Cancelled));
    }

    #[test]
    fn pending_can_transition_to_failed() {
        assert!(AllocationState::Pending.can_transition_to(&AllocationState::Failed));
    }

    #[test]
    fn pending_cannot_transition_to_completed() {
        assert!(!AllocationState::Pending.can_transition_to(&AllocationState::Completed));
    }

    #[test]
    fn pending_cannot_transition_to_checkpointing() {
        assert!(!AllocationState::Pending.can_transition_to(&AllocationState::Checkpointing));
    }

    #[test]
    fn staging_can_transition_to_running() {
        assert!(AllocationState::Staging.can_transition_to(&AllocationState::Running));
    }

    #[test]
    fn staging_can_transition_to_failed() {
        assert!(AllocationState::Staging.can_transition_to(&AllocationState::Failed));
    }

    #[test]
    fn staging_can_transition_to_cancelled() {
        assert!(AllocationState::Staging.can_transition_to(&AllocationState::Cancelled));
    }

    #[test]
    fn running_can_transition_to_checkpointing() {
        assert!(AllocationState::Running.can_transition_to(&AllocationState::Checkpointing));
    }

    #[test]
    fn running_can_transition_to_completed() {
        assert!(AllocationState::Running.can_transition_to(&AllocationState::Completed));
    }

    #[test]
    fn running_can_transition_to_failed() {
        assert!(AllocationState::Running.can_transition_to(&AllocationState::Failed));
    }

    #[test]
    fn running_can_transition_to_cancelled() {
        assert!(AllocationState::Running.can_transition_to(&AllocationState::Cancelled));
    }

    #[test]
    fn running_cannot_transition_to_pending() {
        assert!(!AllocationState::Running.can_transition_to(&AllocationState::Pending));
    }

    #[test]
    fn checkpointing_can_transition_to_suspended() {
        assert!(AllocationState::Checkpointing.can_transition_to(&AllocationState::Suspended));
    }

    #[test]
    fn checkpointing_can_transition_to_failed() {
        assert!(AllocationState::Checkpointing.can_transition_to(&AllocationState::Failed));
    }

    #[test]
    fn checkpointing_can_transition_to_cancelled() {
        assert!(AllocationState::Checkpointing.can_transition_to(&AllocationState::Cancelled));
    }

    #[test]
    fn suspended_can_transition_to_pending() {
        assert!(AllocationState::Suspended.can_transition_to(&AllocationState::Pending));
    }

    #[test]
    fn suspended_can_transition_to_cancelled() {
        assert!(AllocationState::Suspended.can_transition_to(&AllocationState::Cancelled));
    }

    #[test]
    fn suspended_can_transition_to_failed() {
        assert!(AllocationState::Suspended.can_transition_to(&AllocationState::Failed));
    }

    #[test]
    fn completed_is_terminal() {
        assert!(AllocationState::Completed.is_terminal());
    }

    #[test]
    fn failed_is_terminal() {
        assert!(AllocationState::Failed.is_terminal());
    }

    #[test]
    fn cancelled_is_terminal() {
        assert!(AllocationState::Cancelled.is_terminal());
    }

    #[test]
    fn running_is_not_terminal() {
        assert!(!AllocationState::Running.is_terminal());
    }

    #[test]
    fn terminal_states_cannot_transition_to_anything_except_failed_to_pending() {
        let all_states = [
            AllocationState::Pending,
            AllocationState::Staging,
            AllocationState::Running,
            AllocationState::Checkpointing,
            AllocationState::Suspended,
            AllocationState::Completed,
            AllocationState::Failed,
            AllocationState::Cancelled,
        ];

        // Completed and Cancelled are fully terminal
        for terminal in &[AllocationState::Completed, AllocationState::Cancelled] {
            for target in &all_states {
                assert!(
                    !terminal.can_transition_to(target),
                    "{terminal:?} should not transition to {target:?}"
                );
            }
        }

        // Failed can only transition to Pending (service reconciliation)
        assert!(AllocationState::Failed.can_transition_to(&AllocationState::Pending));
        for target in &all_states {
            if *target != AllocationState::Pending {
                assert!(
                    !AllocationState::Failed.can_transition_to(target),
                    "Failed should not transition to {target:?}"
                );
            }
        }
    }

    // ── NodeState transition tests ──

    #[test]
    fn unknown_can_transition_to_booting() {
        assert!(NodeState::Unknown.can_transition_to(&NodeState::Booting));
    }

    #[test]
    fn unknown_can_transition_to_failed() {
        assert!(NodeState::Unknown.can_transition_to(&NodeState::Failed {
            reason: "hw error".into()
        }));
    }

    #[test]
    fn booting_can_transition_to_ready() {
        assert!(NodeState::Booting.can_transition_to(&NodeState::Ready));
    }

    #[test]
    fn booting_can_transition_to_failed() {
        assert!(NodeState::Booting.can_transition_to(&NodeState::Failed {
            reason: "boot failed".into()
        }));
    }

    #[test]
    fn ready_can_transition_to_degraded() {
        assert!(NodeState::Ready.can_transition_to(&NodeState::Degraded {
            reason: "heartbeat missed".into()
        }));
    }

    #[test]
    fn ready_can_transition_to_draining() {
        assert!(NodeState::Ready.can_transition_to(&NodeState::Draining));
    }

    #[test]
    fn ready_can_transition_to_down() {
        assert!(NodeState::Ready.can_transition_to(&NodeState::Down {
            reason: "operator".into()
        }));
    }

    #[test]
    fn degraded_can_transition_to_ready() {
        assert!(NodeState::Degraded {
            reason: "fixed".into()
        }
        .can_transition_to(&NodeState::Ready));
    }

    #[test]
    fn degraded_can_transition_to_down() {
        assert!(NodeState::Degraded {
            reason: "worsened".into()
        }
        .can_transition_to(&NodeState::Down {
            reason: "confirmed".into()
        }));
    }

    #[test]
    fn degraded_can_transition_to_draining() {
        assert!(NodeState::Degraded {
            reason: "draining".into()
        }
        .can_transition_to(&NodeState::Draining));
    }

    #[test]
    fn down_can_transition_to_ready() {
        assert!(NodeState::Down {
            reason: "operator disabled".into()
        }
        .can_transition_to(&NodeState::Ready));
    }

    #[test]
    fn down_can_transition_to_booting() {
        assert!(NodeState::Down {
            reason: "reimage".into()
        }
        .can_transition_to(&NodeState::Booting));
    }

    #[test]
    fn down_can_transition_to_failed() {
        assert!(NodeState::Down {
            reason: "unrecoverable".into()
        }
        .can_transition_to(&NodeState::Failed {
            reason: "hw".into()
        }));
    }

    #[test]
    fn draining_can_transition_to_drained() {
        assert!(NodeState::Draining.can_transition_to(&NodeState::Drained));
    }

    #[test]
    fn draining_cannot_transition_to_ready() {
        assert!(!NodeState::Draining.can_transition_to(&NodeState::Ready));
    }

    #[test]
    fn drained_can_transition_to_ready() {
        assert!(NodeState::Drained.can_transition_to(&NodeState::Ready));
    }

    #[test]
    fn drained_can_transition_to_booting() {
        assert!(NodeState::Drained.can_transition_to(&NodeState::Booting));
    }

    #[test]
    fn failed_node_can_transition_to_booting() {
        assert!(NodeState::Failed {
            reason: "reimage".into()
        }
        .can_transition_to(&NodeState::Booting));
    }

    #[test]
    fn ready_is_operational() {
        assert!(NodeState::Ready.is_operational());
    }

    #[test]
    fn degraded_is_operational() {
        assert!(NodeState::Degraded {
            reason: "minor".into()
        }
        .is_operational());
    }

    #[test]
    fn down_is_not_operational() {
        assert!(!NodeState::Down {
            reason: "down".into()
        }
        .is_operational());
    }

    #[test]
    fn draining_is_not_operational() {
        assert!(!NodeState::Draining.is_operational());
    }

    // ── NetworkDomainState transition tests ──

    #[test]
    fn active_can_transition_to_draining() {
        assert!(NetworkDomainState::Active.can_transition_to(&NetworkDomainState::Draining));
    }

    #[test]
    fn active_cannot_transition_to_released() {
        assert!(!NetworkDomainState::Active.can_transition_to(&NetworkDomainState::Released));
    }

    #[test]
    fn draining_can_transition_to_active() {
        assert!(NetworkDomainState::Draining.can_transition_to(&NetworkDomainState::Active));
    }

    #[test]
    fn draining_can_transition_to_released() {
        assert!(NetworkDomainState::Draining.can_transition_to(&NetworkDomainState::Released));
    }

    #[test]
    fn released_cannot_transition_to_anything() {
        assert!(!NetworkDomainState::Released.can_transition_to(&NetworkDomainState::Active));
        assert!(!NetworkDomainState::Released.can_transition_to(&NetworkDomainState::Draining));
        assert!(!NetworkDomainState::Released.can_transition_to(&NetworkDomainState::Released));
    }

    // ── Serialization round-trip tests ──

    #[test]
    fn allocation_state_serde_roundtrip() {
        let states = [
            AllocationState::Pending,
            AllocationState::Staging,
            AllocationState::Running,
            AllocationState::Checkpointing,
            AllocationState::Suspended,
            AllocationState::Completed,
            AllocationState::Failed,
            AllocationState::Cancelled,
        ];
        for state in &states {
            let json = serde_json::to_string(state).unwrap();
            let deser: AllocationState = serde_json::from_str(&json).unwrap();
            assert_eq!(*state, deser, "roundtrip failed for {state:?}");
        }
    }

    #[test]
    fn node_state_serde_roundtrip() {
        let states = [
            NodeState::Unknown,
            NodeState::Booting,
            NodeState::Ready,
            NodeState::Degraded {
                reason: "test".into(),
            },
            NodeState::Down {
                reason: "test".into(),
            },
            NodeState::Draining,
            NodeState::Drained,
            NodeState::Failed {
                reason: "test".into(),
            },
        ];
        for state in &states {
            let json = serde_json::to_string(state).unwrap();
            let deser: NodeState = serde_json::from_str(&json).unwrap();
            assert_eq!(*state, deser, "roundtrip failed for {state:?}");
        }
    }

    #[test]
    fn cost_weights_default_sums_to_one() {
        let w = CostWeights::default();
        let sum = w.priority
            + w.wait_time
            + w.fair_share
            + w.topology
            + w.data_readiness
            + w.backlog
            + w.energy
            + w.checkpoint_efficiency
            + w.conformance;
        assert!(
            (sum - 1.0).abs() < 1e-10,
            "default weights sum to {sum}, expected 1.0"
        );
    }

    #[test]
    fn cost_weights_serde_roundtrip() {
        let w = CostWeights::default();
        let json = serde_json::to_string(&w).unwrap();
        let deser: CostWeights = serde_json::from_str(&json).unwrap();
        assert!((deser.priority - w.priority).abs() < f64::EPSILON);
        assert!((deser.topology - w.topology).abs() < f64::EPSILON);
    }

    // ── Memory topology tests ──

    #[test]
    fn memory_domain_type_serde_roundtrip() {
        let types = [
            MemoryDomainType::Dram,
            MemoryDomainType::Hbm,
            MemoryDomainType::CxlAttached,
            MemoryDomainType::Unified,
        ];
        for t in &types {
            let json = serde_json::to_string(t).unwrap();
            let deser: MemoryDomainType = serde_json::from_str(&json).unwrap();
            assert_eq!(*t, deser, "roundtrip failed for {t:?}");
        }
    }

    #[test]
    fn memory_link_type_serde_roundtrip() {
        let types = [
            MemoryLinkType::NumaLink,
            MemoryLinkType::CxlSwitch,
            MemoryLinkType::CoherentFabric,
        ];
        for t in &types {
            let json = serde_json::to_string(t).unwrap();
            let deser: MemoryLinkType = serde_json::from_str(&json).unwrap();
            assert_eq!(*t, deser, "roundtrip failed for {t:?}");
        }
    }

    #[test]
    fn memory_policy_serde_roundtrip() {
        let policies = [
            MemoryPolicy::Local,
            MemoryPolicy::Interleave,
            MemoryPolicy::Preferred,
            MemoryPolicy::Bind,
        ];
        for p in &policies {
            let json = serde_json::to_string(p).unwrap();
            let deser: MemoryPolicy = serde_json::from_str(&json).unwrap();
            assert_eq!(*p, deser, "roundtrip failed for {p:?}");
        }
    }

    #[test]
    fn memory_topology_serde_roundtrip() {
        let topo = MemoryTopology {
            domains: vec![
                MemoryDomain {
                    id: 0,
                    domain_type: MemoryDomainType::Dram,
                    capacity_bytes: 128 * 1024 * 1024 * 1024,
                    numa_node: Some(0),
                    attached_cpus: vec![0, 1, 2, 3],
                    attached_gpus: vec![0, 1],
                },
                MemoryDomain {
                    id: 1,
                    domain_type: MemoryDomainType::CxlAttached,
                    capacity_bytes: 256 * 1024 * 1024 * 1024,
                    numa_node: None,
                    attached_cpus: vec![],
                    attached_gpus: vec![],
                },
            ],
            interconnects: vec![MemoryInterconnect {
                domain_a: 0,
                domain_b: 1,
                link_type: MemoryLinkType::CxlSwitch,
                bandwidth_gbps: 64.0,
                latency_ns: 200,
            }],
            total_capacity_bytes: 384 * 1024 * 1024 * 1024,
        };
        let json = serde_json::to_string(&topo).unwrap();
        let deser: MemoryTopology = serde_json::from_str(&json).unwrap();
        assert_eq!(deser.domains.len(), 2);
        assert_eq!(deser.interconnects.len(), 1);
        assert_eq!(deser.total_capacity_bytes, topo.total_capacity_bytes);
    }

    #[test]
    fn resource_constraints_default_has_memory_fields() {
        let rc = ResourceConstraints::default();
        assert!(!rc.require_unified_memory);
        assert!(!rc.prefer_same_numa);
        assert!(!rc.allow_cxl_memory);
        assert!(rc.memory_policy.is_none());
    }

    #[test]
    fn resource_constraints_memory_serde_roundtrip() {
        let rc = ResourceConstraints {
            require_unified_memory: true,
            prefer_same_numa: true,
            allow_cxl_memory: false,
            memory_policy: Some(MemoryPolicy::Interleave),
            ..Default::default()
        };
        let json = serde_json::to_string(&rc).unwrap();
        let deser: ResourceConstraints = serde_json::from_str(&json).unwrap();
        assert!(deser.require_unified_memory);
        assert!(deser.prefer_same_numa);
        assert!(!deser.allow_cxl_memory);
        assert_eq!(deser.memory_policy, Some(MemoryPolicy::Interleave));
    }

    // ── Property-based tests ──

    mod proptests {
        use super::*;
        use proptest::prelude::*;

        fn arb_allocation_state() -> impl Strategy<Value = AllocationState> {
            prop_oneof![
                Just(AllocationState::Pending),
                Just(AllocationState::Staging),
                Just(AllocationState::Running),
                Just(AllocationState::Checkpointing),
                Just(AllocationState::Suspended),
                Just(AllocationState::Completed),
                Just(AllocationState::Failed),
                Just(AllocationState::Cancelled),
            ]
        }

        fn arb_memory_domain_type() -> impl Strategy<Value = MemoryDomainType> {
            prop_oneof![
                Just(MemoryDomainType::Dram),
                Just(MemoryDomainType::Hbm),
                Just(MemoryDomainType::CxlAttached),
                Just(MemoryDomainType::Unified),
            ]
        }

        fn arb_memory_policy() -> impl Strategy<Value = MemoryPolicy> {
            prop_oneof![
                Just(MemoryPolicy::Local),
                Just(MemoryPolicy::Interleave),
                Just(MemoryPolicy::Preferred),
                Just(MemoryPolicy::Bind),
            ]
        }

        proptest! {
            #[test]
            fn memory_domain_type_roundtrip(dt in arb_memory_domain_type()) {
                let json = serde_json::to_string(&dt).unwrap();
                let back: MemoryDomainType = serde_json::from_str(&json).unwrap();
                prop_assert_eq!(dt, back);
            }

            #[test]
            fn memory_policy_roundtrip(p in arb_memory_policy()) {
                let json = serde_json::to_string(&p).unwrap();
                let back: MemoryPolicy = serde_json::from_str(&json).unwrap();
                prop_assert_eq!(p, back);
            }

            #[test]
            fn terminal_states_block_all_transitions(target in arb_allocation_state()) {
                // Completed and Cancelled are fully terminal
                for terminal in &[AllocationState::Completed, AllocationState::Cancelled] {
                    prop_assert!(!terminal.can_transition_to(&target));
                }
                // Failed → Pending is allowed (service reconciliation), everything else blocked
                if target != AllocationState::Pending {
                    prop_assert!(!AllocationState::Failed.can_transition_to(&target));
                }
            }

            #[test]
            fn no_self_transitions(state in arb_allocation_state()) {
                prop_assert!(!state.can_transition_to(&state));
            }

            #[test]
            fn node_count_range_min_le_max(min in 1u32..1000, max in 1u32..1000) {
                prop_assume!(min <= max);
                let _nc = NodeCount::Range { min, max };
                // Construction succeeds without panic
            }
        }
    }
}