//! Processor abstractions for heterogeneous compute.
//!
//! Supports all processor types:
//! - CPU (x86_64, ARM64, RISC-V)
//! - GPU (NVIDIA CUDA, AMD ROCm, Intel OneAPI, Apple Metal)
//! - TPU (Google TPU v2-v5)
//! - NPU (Apple Neural Engine, Qualcomm Hexagon, Intel VPU)
//! - LPU (Groq Language Processing Unit)
//! - FPGA (Xilinx, Intel/Altera)
//! - DSP (Digital Signal Processors)
//! - Custom accelerators

mod capabilities;
mod operation;
mod profiles;
mod types;

pub use capabilities::{ComputeThroughput, MemorySpecs, MemoryType, ProcessorCapabilities};
pub use operation::{Operation, OperationType};
pub use profiles::ProcessorProfiles;
pub use types::*;

use crate::error::ComputeError;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::time::Duration;

/// Unique processor identifier (within a node).
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ProcessorId(pub u64);

impl std::fmt::Display for ProcessorId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "proc_{}", self.0)
    }
}

/// Unified abstraction for any processor type.
#[async_trait]
pub trait Processor: Send + Sync {
    /// Get processor ID.
    fn id(&self) -> ProcessorId;

    /// Get processor type.
    fn processor_type(&self) -> ProcessorType;

    /// Get capabilities.
    fn capabilities(&self) -> &ProcessorCapabilities;

    /// Check if processor can execute operation.
    fn can_execute(&self, op: &Operation) -> bool;

    /// Estimate execution time for operation.
    fn estimate_time(&self, op: &Operation) -> Duration;

    /// Estimate energy consumption for operation (Joules).
    fn estimate_energy(&self, op: &Operation) -> f64;

    /// Execute operation.
    async fn execute(&self, op: Operation) -> Result<OperationResult, ComputeError>;

    /// Current utilization (0.0 - 1.0).
    fn utilization(&self) -> f64;

    /// Available memory (bytes).
    fn available_memory(&self) -> u64;

    /// Check if this processor shares memory with another type.
    fn shares_memory_with(&self, other: &ProcessorType) -> bool {
        // By default, processors don't share memory
        // Override for unified memory architectures (Apple Silicon, AMD APUs)
        self.processor_type() == *other
    }
}

/// Result of an operation execution.
#[derive(Clone, Debug)]
pub struct OperationResult {
    /// Output data.
    pub output: Vec<u8>,
    /// Execution time.
    pub duration: Duration,
    /// Energy consumed (Joules).
    pub energy: f64,
    /// Peak memory used (bytes).
    pub peak_memory: u64,
}

/// Generic processor implementation for simulation/testing.
pub struct GenericProcessor {
    id: ProcessorId,
    processor_type: ProcessorType,
    capabilities: ProcessorCapabilities,
    utilization: std::sync::atomic::AtomicU64,
    available_memory: std::sync::atomic::AtomicU64,
}

impl GenericProcessor {
    /// Creates a new generic processor.
    pub fn new(
        id: ProcessorId,
        processor_type: ProcessorType,
        capabilities: ProcessorCapabilities,
    ) -> Self {
        let available_memory = capabilities.memory.capacity_bytes;
        Self {
            id,
            processor_type,
            capabilities,
            utilization: std::sync::atomic::AtomicU64::new(0),
            available_memory: std::sync::atomic::AtomicU64::new(available_memory),
        }
    }

    /// Creates a CPU processor.
    pub fn cpu(id: ProcessorId, variant: CpuVariant) -> Self {
        Self::new(
            id,
            ProcessorType::Cpu(variant),
            ProcessorProfiles::cpu_default(),
        )
    }

    /// Creates an NVIDIA GPU processor.
    pub fn nvidia_gpu(id: ProcessorId, compute_capability: (u8, u8)) -> Self {
        let capabilities = match compute_capability {
            (9, 0) => ProcessorProfiles::nvidia_h100(),
            (8, 9) => ProcessorProfiles::nvidia_rtx_4090(),
            (8, 6) => ProcessorProfiles::nvidia_rtx_3090(),
            _ => ProcessorProfiles::nvidia_default(),
        };
        Self::new(
            id,
            ProcessorType::Gpu(GpuVariant::NvidiaCuda { compute_capability }),
            capabilities,
        )
    }

    /// Creates a TPU processor.
    pub fn tpu(id: ProcessorId, version: TpuVersion) -> Self {
        let capabilities = match version {
            TpuVersion::V5p => ProcessorProfiles::google_tpu_v5p(),
            TpuVersion::V4 => ProcessorProfiles::google_tpu_v4(),
            _ => ProcessorProfiles::google_tpu_default(),
        };
        Self::new(id, ProcessorType::Tpu(version), capabilities)
    }

    /// Creates a Groq LPU processor.
    pub fn lpu(id: ProcessorId) -> Self {
        Self::new(id, ProcessorType::Lpu, ProcessorProfiles::groq_lpu())
    }

    /// Creates an Apple Neural Engine processor.
    pub fn apple_neural_engine(id: ProcessorId, cores: u32) -> Self {
        Self::new(
            id,
            ProcessorType::Npu(NpuVariant::AppleNeuralEngine { cores }),
            ProcessorProfiles::apple_neural_engine(cores),
        )
    }
}

#[async_trait]
impl Processor for GenericProcessor {
    fn id(&self) -> ProcessorId {
        self.id
    }

    fn processor_type(&self) -> ProcessorType {
        self.processor_type.clone()
    }

    fn capabilities(&self) -> &ProcessorCapabilities {
        &self.capabilities
    }

    fn can_execute(&self, op: &Operation) -> bool {
        self.capabilities.operations.contains(&op.op_type())
    }

    fn estimate_time(&self, op: &Operation) -> Duration {
        // Estimate based on FLOPS and operation complexity
        let flops_needed = op.estimated_flops();
        let throughput = match op.precision() {
            Precision::Fp32 => self.capabilities.compute.fp32_tflops,
            Precision::Fp16 => self.capabilities.compute.fp16_tflops,
            Precision::Bf16 => self.capabilities.compute.bf16_tflops,
            Precision::Int8 => self.capabilities.compute.int8_tops,
            Precision::Int4 => self.capabilities.compute.int4_tops,
            Precision::Fp64 => self.capabilities.compute.fp64_tflops,
        };

        if throughput > 0.0 {
            let tflops = throughput;
            let flops_per_second = tflops * 1e12;
            let seconds = flops_needed / flops_per_second;
            Duration::from_secs_f64(seconds)
        } else {
            Duration::from_secs(1) // Fallback
        }
    }

    fn estimate_energy(&self, op: &Operation) -> f64 {
        // Estimate based on TDP and execution time
        let duration = self.estimate_time(op);
        let watts = self.capabilities.power.tdp_watts as f64;
        let efficiency = self.capabilities.power.efficiency;
        watts * duration.as_secs_f64() * efficiency
    }

    async fn execute(&self, op: Operation) -> Result<OperationResult, ComputeError> {
        // Check if we can execute
        if !self.can_execute(&op) {
            return Err(ComputeError::OperationNotSupported(
                self.processor_type.clone(),
                format!("{:?}", op.op_type()),
            ));
        }

        // Simulate execution
        let duration = self.estimate_time(&op);
        let energy = self.estimate_energy(&op);

        // Update utilization
        self.utilization
            .store(50, std::sync::atomic::Ordering::Relaxed);

        // Simulate work
        tokio::time::sleep(Duration::from_micros(100)).await;

        // Reset utilization
        self.utilization
            .store(0, std::sync::atomic::Ordering::Relaxed);

        Ok(OperationResult {
            output: vec![],
            duration,
            energy,
            peak_memory: op.estimated_memory(),
        })
    }

    fn utilization(&self) -> f64 {
        self.utilization.load(std::sync::atomic::Ordering::Relaxed) as f64 / 100.0
    }

    fn available_memory(&self) -> u64 {
        self.available_memory
            .load(std::sync::atomic::Ordering::Relaxed)
    }

    fn shares_memory_with(&self, other: &ProcessorType) -> bool {
        match (&self.processor_type, other) {
            // Apple Silicon has unified memory
            (ProcessorType::Cpu(CpuVariant::Arm64 { .. }), ProcessorType::Gpu(GpuVariant::AppleMetal))
            | (ProcessorType::Gpu(GpuVariant::AppleMetal), ProcessorType::Cpu(CpuVariant::Arm64 { .. }))
            | (ProcessorType::Npu(NpuVariant::AppleNeuralEngine { .. }), ProcessorType::Cpu(CpuVariant::Arm64 { .. }))
            | (ProcessorType::Npu(NpuVariant::AppleNeuralEngine { .. }), ProcessorType::Gpu(GpuVariant::AppleMetal)) => true,
            // Same type always shares
            (a, b) if a == b => true,
            _ => false,
        }
    }
}

/// Precision for operations.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum Precision {
    Fp64,
    Fp32,
    Fp16,
    Bf16,
    Int8,
    Int4,
}

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

    #[test]
    fn test_processor_creation() {
        let cpu = GenericProcessor::cpu(
            ProcessorId(0),
            CpuVariant::X86_64 {
                avx: AvxSupport::Avx512,
            },
        );

        assert_eq!(cpu.id(), ProcessorId(0));
        assert!(matches!(cpu.processor_type(), ProcessorType::Cpu(_)));
    }

    #[test]
    fn test_gpu_creation() {
        let gpu = GenericProcessor::nvidia_gpu(ProcessorId(1), (9, 0));

        assert_eq!(gpu.id(), ProcessorId(1));
        assert!(matches!(
            gpu.processor_type(),
            ProcessorType::Gpu(GpuVariant::NvidiaCuda { .. })
        ));
    }

    #[test]
    fn test_unified_memory() {
        let apple_cpu = GenericProcessor::new(
            ProcessorId(0),
            ProcessorType::Cpu(CpuVariant::Arm64 { sve: false }),
            ProcessorCapabilities::default(),
        );

        assert!(apple_cpu.shares_memory_with(&ProcessorType::Gpu(GpuVariant::AppleMetal)));
    }

    #[tokio::test]
    async fn test_operation_execution() {
        let cpu = GenericProcessor::cpu(
            ProcessorId(0),
            CpuVariant::X86_64 {
                avx: AvxSupport::Avx512,
            },
        );

        let op = Operation::MatMul {
            m: 1024,
            n: 1024,
            k: 1024,
            precision: Precision::Fp32,
        };

        // CPU might not support all ops depending on capabilities
        // This is testing the infrastructure
        let result = cpu.execute(op).await;
        // Result depends on capabilities
        assert!(result.is_ok() || result.is_err());
    }
}
