//! The executor trait.
//!
//! An *executor* is a vendor-specific compute backend bound to one or
//! more devices. A1 ships with a single CUDA executor wrapping
//! `ggml-cuda` against one Tesla P40. A2 adds tensor-split across two
//! P40s under the same executor. A4 introduces a second executor
//! (AMD via `ggml-rocm`) and exercises the cross-vendor activation
//! handoff path — the same path that already exists in the trait
//! from day one, just with no second implementer to call it.
//!
//! The shape of this trait is load-bearing for the C-ABI module
//! boundary that lands in Thread I. Decisions taken here propagate
//! into the stable `everycycle-api` surface; revisit before adding a
//! method that cannot be flattened to a `repr(C)` call.

use everycycle_hal::PciBusAddress;

/// Vendor of the kernel collection this executor wraps. The
/// supervisor uses this for routing and for picking a same-vendor
/// fast path when one is available.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub enum Vendor {
    Cuda,
    Rocm,
    Vulkan,
    Cpu,
}

/// Numeric type of the values in a tensor. Mirrors the precisions
/// the appraise probe measures; the runtime asks the executor only
/// for dtypes it knows the executor declared support for.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub enum Dtype {
    Fp32,
    Fp16,
    Bf16,
    Fp8E4M3,
    Fp8E5M2,
    Int8,
    Int4,
}

/// Logical tensor shape. Row-major; rank up to eight is enough for
/// every workload in the B-thread roadmap.
#[derive(Clone, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub struct TensorShape {
    pub dims: Vec<u64>,
}

impl TensorShape {
    #[must_use]
    pub fn element_count(&self) -> u64 {
        self.dims.iter().product()
    }
}

/// One device the executor is bound to. The bus address is the
/// stable identity; the kind tag is informational and matches the
/// HAL inventory.
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct BoundDevice {
    pub bus_address: PciBusAddress,
    pub display_name: String,
}

/// Opaque handle to an activation tensor that lives inside one
/// executor's address space. The handle has meaning *only* to the
/// executor that produced it; passing it to a different executor is
/// a programming error. Cross-executor transfers go through the
/// portable form.
///
/// Lifetime is reference-counted at the supervisor; the executor
/// must hold the underlying memory live until `release_activation`
/// is called.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct ActivationHandle(pub u64);

/// Wire-portable activation. The supervisor uses this to bridge
/// between executors of different vendors and to checkpoint state
/// across crashes.
///
/// The bytes are laid out in row-major order at the declared dtype.
/// Endianness is little. Quantized blocks (Q4_K, IQ4_XS, ...) are
/// not represented here — activations are dequantized to a real
/// dtype before export. Cross-vendor weight handoff is a different
/// problem and does not flow through this path.
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct PortableActivation {
    pub shape: TensorShape,
    pub dtype: Dtype,
    pub bytes: Vec<u8>,
}

/// A unit of compute the planner has assigned to one executor.
///
/// Opaque at this layer: the planner emits a vendor-specific plan
/// the executor knows how to decode (e.g. a serialized ggml
/// computation graph fragment, or a CUDA-graph descriptor). The
/// supervisor never inspects this — it just routes it.
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct SubGraphPlan {
    pub vendor: Vendor,
    pub payload: Vec<u8>,
}

/// Errors an executor call can return. Variants will grow with
/// experience; the closed set today is enough for A1.
#[derive(Debug)]
pub enum ExecutorError {
    /// Requested dtype is not implemented on this executor.
    UnsupportedDtype(Dtype),
    /// Plan vendor does not match this executor's vendor.
    VendorMismatch { plan: Vendor, executor: Vendor },
    /// Allocation failed — usually VRAM exhaustion.
    OutOfMemory,
    /// Handle was issued by a different executor instance.
    UnknownHandle(ActivationHandle),
    /// Vendor SDK returned an error string. We pass it through
    /// rather than try to enumerate every possible vendor error.
    Backend(String),
}

impl core::fmt::Display for ExecutorError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::UnsupportedDtype(d) => write!(f, "executor does not support dtype {d:?}"),
            Self::VendorMismatch { plan, executor } => {
                write!(
                    f,
                    "plan vendor {plan:?} does not match executor vendor {executor:?}"
                )
            }
            Self::OutOfMemory => f.write_str("executor allocation failed"),
            Self::UnknownHandle(h) => write!(f, "unknown activation handle {h:?}"),
            Self::Backend(s) => write!(f, "backend error: {s}"),
        }
    }
}

impl std::error::Error for ExecutorError {}

/// The vendor backend boundary.
///
/// Every method is callable across threads; implementations must be
/// internally synchronized. The supervisor will issue concurrent
/// calls into one executor from multiple client tasks, and the
/// multi-client design (Thread C) assumes this from A1.
pub trait Executor: Send + Sync {
    /// Vendor of the wrapped kernel collection.
    fn vendor(&self) -> Vendor;

    /// Human-readable identifier for log lines and the audit screen.
    fn name(&self) -> &str;

    /// Devices this executor owns. A given physical device belongs
    /// to exactly one executor at any time.
    fn devices(&self) -> &[BoundDevice];

    /// Dtypes this executor advertises support for. The planner
    /// will not ask for a dtype not in this list.
    fn supported_dtypes(&self) -> &[Dtype];

    /// Allocate an activation tensor of the given shape and dtype
    /// on this executor's device(s).
    ///
    /// # Errors
    ///
    /// `UnsupportedDtype` if the dtype is not in `supported_dtypes`;
    /// `OutOfMemory` if allocation fails.
    fn alloc_activation(
        &self,
        shape: &TensorShape,
        dtype: Dtype,
    ) -> Result<ActivationHandle, ExecutorError>;

    /// Release an activation previously returned by `alloc_activation`
    /// or `execute`. Idempotent: re-releasing a handle is not an error.
    fn release_activation(&self, handle: ActivationHandle);

    /// Execute one planner-assigned sub-graph. The inputs are
    /// activations local to this executor; the outputs are too.
    ///
    /// # Errors
    ///
    /// `VendorMismatch` if `plan.vendor` is not this executor's
    /// vendor; `UnknownHandle` if any input was not issued by this
    /// executor; `Backend` for vendor-SDK failures.
    fn execute(
        &self,
        plan: &SubGraphPlan,
        inputs: &[ActivationHandle],
    ) -> Result<Vec<ActivationHandle>, ExecutorError>;

    /// Export an activation as a portable buffer.
    ///
    /// Forward-compat for Thread A4 (cross-vendor): the supervisor
    /// calls this on the producing executor, then calls
    /// `import_activation` on the consuming executor of a different
    /// vendor. Same-vendor pipelines should never hit this path —
    /// the supervisor short-circuits to a direct handle hand-off
    /// when source and destination executors are the same instance.
    ///
    /// Implementations are required from day one even when there is
    /// no second vendor to hand off to. The minimum contract is a
    /// device-to-host staging copy.
    ///
    /// # Errors
    ///
    /// `UnknownHandle` if the handle was not issued by this
    /// executor; `Backend` for SDK failures during the staging copy.
    fn export_activation(
        &self,
        handle: ActivationHandle,
    ) -> Result<PortableActivation, ExecutorError>;

    /// Import an activation produced by another executor.
    ///
    /// The byte layout contract is the one in `PortableActivation`'s
    /// documentation: row-major, little-endian, dequantized to the
    /// declared dtype. Vendor-specific re-quantization on import is
    /// the executor's choice, not the caller's concern.
    ///
    /// # Errors
    ///
    /// `UnsupportedDtype` if the portable activation's dtype is not
    /// in `supported_dtypes`; `OutOfMemory` if allocation fails;
    /// `Backend` for SDK failures during the host-to-device copy.
    fn import_activation(
        &self,
        portable: &PortableActivation,
    ) -> Result<ActivationHandle, ExecutorError>;
}

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

    #[test]
    fn shape_element_count() {
        let s = TensorShape {
            dims: vec![2, 3, 4],
        };
        assert_eq!(s.element_count(), 24);
    }

    #[test]
    fn error_display_renders() {
        let e = ExecutorError::VendorMismatch {
            plan: Vendor::Cuda,
            executor: Vendor::Rocm,
        };
        let rendered = e.to_string();
        assert!(rendered.contains("Cuda"));
        assert!(rendered.contains("Rocm"));
    }
}