//! Task definitions and decomposition.

use crate::error::ComputeError;
use crate::processor::{Operation, OperationType, Precision, ProcessorType};
use crate::{ComputeJob, JobType};
use serde::{Deserialize, Serialize};
use std::time::Duration;

/// Unique task identifier.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct TaskId(pub u64);

impl TaskId {
    /// Creates a new task ID.
    pub fn new() -> Self {
        use rand::Rng;
        TaskId(rand::thread_rng().gen())
    }
}

impl Default for TaskId {
    fn default() -> Self {
        Self::new()
    }
}

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

/// Task priority levels.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub enum TaskPriority {
    /// Background, can be preempted.
    Background = 0,
    /// Normal priority.
    Normal = 1,
    /// High priority.
    High = 2,
    /// Critical, must complete.
    Critical = 3,
}

impl Default for TaskPriority {
    fn default() -> Self {
        TaskPriority::Normal
    }
}

/// Task execution status.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum TaskStatus {
    /// Waiting to be scheduled.
    Pending,
    /// Queued for execution.
    Queued,
    /// Currently executing.
    Running,
    /// Completed successfully.
    Completed,
    /// Failed.
    Failed,
    /// Cancelled.
    Cancelled,
}

/// A schedulable task.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Task {
    /// Task ID.
    pub id: TaskId,
    /// Operation to execute.
    pub operation: Operation,
    /// Priority level.
    pub priority: TaskPriority,
    /// Dependencies (tasks that must complete first).
    pub dependencies: Vec<TaskId>,
    /// Current status.
    pub status: TaskStatus,
    /// Deadline (optional).
    pub deadline: Option<u64>,
}

impl Task {
    /// Creates a new task.
    pub fn new(operation: Operation) -> Self {
        Self {
            id: TaskId::new(),
            operation,
            priority: TaskPriority::Normal,
            dependencies: Vec::new(),
            status: TaskStatus::Pending,
            deadline: None,
        }
    }

    /// Sets the priority.
    pub fn with_priority(mut self, priority: TaskPriority) -> Self {
        self.priority = priority;
        self
    }

    /// Adds dependencies.
    pub fn with_dependencies(mut self, deps: Vec<TaskId>) -> Self {
        self.dependencies = deps;
        self
    }

    /// Sets deadline.
    pub fn with_deadline(mut self, deadline: u64) -> Self {
        self.deadline = Some(deadline);
        self
    }

    /// Checks if task is compatible with a processor type.
    pub fn is_compatible_with(&self, proc_type: ProcessorType) -> bool {
        // Check based on operation type
        let op_type = self.operation.op_type();

        match proc_type {
            ProcessorType::Cpu(_) => {
                // CPUs can do most things, but slowly
                true
            }
            ProcessorType::Gpu(_) => {
                // GPUs are good for parallel operations
                matches!(
                    op_type,
                    OperationType::MatMul
                        | OperationType::Conv2d
                        | OperationType::SelfAttention
                        | OperationType::FlashAttention
                        | OperationType::Embedding
                        | OperationType::Add
                        | OperationType::Mul
                        | OperationType::Softmax
                )
            }
            ProcessorType::Tpu(_) => {
                // TPUs are good for large matrix ops
                matches!(
                    op_type,
                    OperationType::MatMul
                        | OperationType::Conv2d
                        | OperationType::SelfAttention
                        | OperationType::FlashAttention
                )
            }
            ProcessorType::Lpu => {
                // LPUs are good for sequential inference
                matches!(
                    op_type,
                    OperationType::MatMul
                        | OperationType::SelfAttention
                        | OperationType::KVCache
                        | OperationType::Sampling
                )
            }
            ProcessorType::Npu(_) => {
                // NPUs are good for inference
                matches!(
                    op_type,
                    OperationType::MatMul
                        | OperationType::Conv2d
                        | OperationType::Add
                        | OperationType::Softmax
                )
            }
            _ => true, // Default to compatible
        }
    }
}

/// Result of task execution.
#[derive(Clone, Debug)]
pub struct TaskResult {
    /// Task ID.
    pub task_id: TaskId,
    /// Output data.
    pub output: Vec<u8>,
    /// Execution duration.
    pub duration: Duration,
    /// Energy consumed (Joules).
    pub energy: f64,
}

/// Compute task for job execution.
#[derive(Clone, Debug)]
pub struct ComputeTask {
    /// Task.
    pub task: Task,
    /// Resource requirements.
    pub requirements: TaskRequirements,
    /// Preferred processor type.
    pub preferred_processor: Option<ProcessorType>,
    /// Fallback processor type.
    pub fallback_processor: Option<ProcessorType>,
}

/// Task resource requirements.
#[derive(Clone, Debug, Default)]
pub struct TaskRequirements {
    /// Minimum memory (bytes).
    pub min_memory: u64,
    /// Minimum TFLOPS.
    pub min_tflops: f64,
    /// Maximum latency (ms).
    pub max_latency_ms: Option<u32>,
    /// Requires specific precision.
    pub precision: Option<Precision>,
}

/// Decomposed workload.
#[derive(Clone, Debug)]
pub struct DecomposedWorkload {
    /// All tasks.
    pub tasks: Vec<Task>,
    /// Total estimated FLOPS.
    pub estimated_flops: f64,
    /// Total estimated memory.
    pub estimated_memory: u64,
}

/// Task decomposer that breaks jobs into schedulable tasks.
pub struct TaskDecomposer {
    /// Default batch size for inference.
    inference_batch_size: usize,
    /// Default precision.
    default_precision: Precision,
}

impl TaskDecomposer {
    /// Creates a new task decomposer.
    pub fn new() -> Self {
        Self {
            inference_batch_size: 32,
            default_precision: Precision::Fp16,
        }
    }

    /// Decomposes a job into tasks.
    pub fn decompose(&self, job: &ComputeJob) -> Result<Vec<Task>, ComputeError> {
        match &job.job_type {
            JobType::Training { .. } => self.decompose_training(job),
            JobType::Inference { .. } => self.decompose_inference(job),
            JobType::Container { .. } => self.decompose_container(job),
            JobType::Serverless { .. } => self.decompose_serverless(job),
            JobType::Wasm { .. } => self.decompose_wasm(job),
        }
    }

    /// Decompose training job.
    fn decompose_training(&self, job: &ComputeJob) -> Result<Vec<Task>, ComputeError> {
        let mut tasks = Vec::new();

        if let JobType::Training {
            epochs: _,
            batch_size,
            ..
        } = &job.job_type
        {
            // Data loading task
            tasks.push(
                Task::new(Operation::DataLoad {
                    bytes: 1024 * 1024 * 100, // 100MB
                    async_: true,
                })
                .with_priority(TaskPriority::High),
            );

            let data_load_id = tasks[0].id;

            // Preprocessing task
            tasks.push(
                Task::new(Operation::DataPreprocess {
                    batch: *batch_size as usize,
                    transforms: vec!["normalize".to_string(), "augment".to_string()],
                })
                .with_dependencies(vec![data_load_id])
                .with_priority(TaskPriority::High),
            );

            let preprocess_id = tasks[1].id;

            // Forward pass (simplified as MatMul)
            tasks.push(
                Task::new(Operation::MatMul {
                    m: *batch_size as usize,
                    n: 4096,
                    k: 4096,
                    precision: self.default_precision,
                })
                .with_dependencies(vec![preprocess_id])
                .with_priority(TaskPriority::Critical),
            );

            let forward_id = tasks[2].id;

            // Backward pass
            tasks.push(
                Task::new(Operation::Backward {
                    forward_op: Box::new(Operation::MatMul {
                        m: *batch_size as usize,
                        n: 4096,
                        k: 4096,
                        precision: self.default_precision,
                    }),
                })
                .with_dependencies(vec![forward_id])
                .with_priority(TaskPriority::Critical),
            );

            let backward_id = tasks[3].id;

            // Optimizer step
            tasks.push(
                Task::new(Operation::OptimizerStep {
                    parameters: 1_000_000,
                    optimizer: "adamw".to_string(),
                    precision: self.default_precision,
                })
                .with_dependencies(vec![backward_id])
                .with_priority(TaskPriority::High),
            );
        }

        Ok(tasks)
    }

    /// Decompose inference job.
    fn decompose_inference(&self, job: &ComputeJob) -> Result<Vec<Task>, ComputeError> {
        let mut tasks = Vec::new();

        if let JobType::Inference { batch_size, .. } = &job.job_type {
            // Tokenization (CPU optimal)
            tasks.push(
                Task::new(Operation::Tokenization {
                    text_bytes: 4096,
                    vocab_size: 32000,
                })
                .with_priority(TaskPriority::High),
            );

            let token_id = tasks[0].id;

            // Embedding (GPU optimal)
            tasks.push(
                Task::new(Operation::Embedding {
                    batch: *batch_size as usize,
                    seq_len: 512,
                    vocab_size: 32000,
                    embed_dim: 4096,
                    precision: self.default_precision,
                })
                .with_dependencies(vec![token_id])
                .with_priority(TaskPriority::Critical),
            );

            let embed_id = tasks[1].id;

            // Self-attention (TPU/GPU optimal)
            tasks.push(
                Task::new(Operation::SelfAttention {
                    batch: *batch_size as usize,
                    seq_len: 512,
                    num_heads: 32,
                    head_dim: 128,
                    precision: self.default_precision,
                })
                .with_dependencies(vec![embed_id])
                .with_priority(TaskPriority::Critical),
            );

            let attention_id = tasks[2].id;

            // Sampling (LPU optimal)
            tasks.push(
                Task::new(Operation::Sampling {
                    batch: *batch_size as usize,
                    vocab_size: 32000,
                    temperature: 0.7,
                })
                .with_dependencies(vec![attention_id])
                .with_priority(TaskPriority::High),
            );

            let sample_id = tasks[3].id;

            // Detokenization (CPU optimal)
            tasks.push(
                Task::new(Operation::Detokenization {
                    tokens: 256,
                    vocab_size: 32000,
                })
                .with_dependencies(vec![sample_id])
                .with_priority(TaskPriority::Normal),
            );
        }

        Ok(tasks)
    }

    /// Decompose container job.
    fn decompose_container(&self, _job: &ComputeJob) -> Result<Vec<Task>, ComputeError> {
        // Container jobs are typically a single task
        Ok(vec![Task::new(Operation::Generic {
            op_type: OperationType::DataLoad,
            flops: 1e9,
            memory: 1024 * 1024 * 1024,
        })
        .with_priority(TaskPriority::Normal)])
    }

    /// Decompose serverless function.
    fn decompose_serverless(&self, _job: &ComputeJob) -> Result<Vec<Task>, ComputeError> {
        // Serverless is typically a single task
        Ok(vec![Task::new(Operation::Generic {
            op_type: OperationType::DataPreprocess,
            flops: 1e6,
            memory: 256 * 1024 * 1024,
        })
        .with_priority(TaskPriority::High)])
    }

    /// Decompose WASM job.
    fn decompose_wasm(&self, _job: &ComputeJob) -> Result<Vec<Task>, ComputeError> {
        // WASM is typically a single task
        Ok(vec![Task::new(Operation::Generic {
            op_type: OperationType::DataPreprocess,
            flops: 1e6,
            memory: 16 * 1024 * 1024,
        })
        .with_priority(TaskPriority::Normal)])
    }
}

impl Default for TaskDecomposer {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_task_creation() {
        let task = Task::new(Operation::MatMul {
            m: 1024,
            n: 1024,
            k: 1024,
            precision: Precision::Fp32,
        })
        .with_priority(TaskPriority::High);

        assert_eq!(task.priority, TaskPriority::High);
        assert!(task.dependencies.is_empty());
        assert_eq!(task.status, TaskStatus::Pending);
    }

    #[test]
    fn test_task_dependencies() {
        let task1 = Task::new(Operation::DataLoad {
            bytes: 1000,
            async_: true,
        });
        let task1_id = task1.id;

        let task2 = Task::new(Operation::MatMul {
            m: 1024,
            n: 1024,
            k: 1024,
            precision: Precision::Fp32,
        })
        .with_dependencies(vec![task1_id]);

        assert_eq!(task2.dependencies, vec![task1_id]);
    }

    #[test]
    fn test_task_compatibility() {
        let matmul_task = Task::new(Operation::MatMul {
            m: 1024,
            n: 1024,
            k: 1024,
            precision: Precision::Fp32,
        });

        // MatMul should be compatible with GPU and TPU
        assert!(matmul_task.is_compatible_with(ProcessorType::Gpu(
            crate::processor::GpuVariant::NvidiaCuda {
                compute_capability: (8, 0)
            }
        )));
        assert!(matmul_task.is_compatible_with(ProcessorType::Tpu(
            crate::processor::TpuVersion::V5p
        )));

        let data_load_task = Task::new(Operation::DataLoad {
            bytes: 1000,
            async_: true,
        });

        // DataLoad should be compatible with CPU
        assert!(data_load_task.is_compatible_with(ProcessorType::Cpu(
            crate::processor::CpuVariant::default()
        )));
    }

    #[test]
    fn test_task_decomposer() {
        let decomposer = TaskDecomposer::new();

        let job = ComputeJob {
            id: crate::JobId::new(),
            owner: [0u8; 32],
            job_type: JobType::Inference {
                model_cid: "model".to_string(),
                input_format: "json".to_string(),
                batch_size: 1,
            },
            resources: crate::ResourceRequirements::default(),
            input_cid: None,
            max_budget: 1_000_000,
            priority: crate::JobPriority::Normal,
            created_at: 0,
            deadline: None,
        };

        let tasks = decomposer.decompose(&job).unwrap();
        assert!(!tasks.is_empty());

        // Check dependencies form a chain
        for (i, task) in tasks.iter().enumerate() {
            if i > 0 {
                assert!(!task.dependencies.is_empty());
            }
        }
    }
}
