//! Optimization passes for GPU kernels

use crate::ir::{IRGraph, NodeId, OperationType};
use anyhow::Result;
use tracing::debug;

/// Optimization level
#[derive(Debug, Clone, Copy)]
pub enum OptLevel {
    /// No optimizations
    None,
    /// Basic optimizations
    Basic,
    /// Aggressive optimizations
    Aggressive,
}

/// Statistics from optimization passes
#[derive(Debug, Clone, Default)]
pub struct OptimizationStats {
    pub passes_applied: u32,
    pub nodes_eliminated: u32,
    pub operations_fused: u32,
    pub memory_saved: u64,
}

/// Trait for optimization passes
pub trait OptimizationPass {
    /// Apply the optimization pass to the graph
    fn apply(&self, graph: &mut IRGraph) -> Result<OptimizationStats>;

    /// Get the name of this optimization pass
    fn name(&self) -> &str;
}

/// Dead code elimination pass
pub struct DeadCodeEliminationPass;

impl OptimizationPass for DeadCodeEliminationPass {
    fn apply(&self, graph: &mut IRGraph) -> Result<OptimizationStats> {
        let mut stats = OptimizationStats::default();
        let mut dead_nodes = Vec::new();

        // Find nodes with no outputs and not in graph outputs
        for (node_id, node) in &graph.nodes {
            if node.outputs.is_empty() && !graph.outputs.contains(node_id) {
                // Check if this node has side effects
                match node.operation {
                    OperationType::Store | OperationType::Call => {
                        // These have side effects, keep them
                    }
                    _ => {
                        dead_nodes.push(*node_id);
                    }
                }
            }
        }

        // Remove dead nodes
        for node_id in dead_nodes {
            graph.remove_node(node_id);
            stats.nodes_eliminated += 1;
        }

        stats.passes_applied = 1;
        Ok(stats)
    }

    fn name(&self) -> &'static str {
        "DeadCodeElimination"
    }
}

/// Constant folding pass
pub struct ConstantFoldingPass;

impl OptimizationPass for ConstantFoldingPass {
    fn apply(&self, graph: &mut IRGraph) -> Result<OptimizationStats> {
        let mut stats = OptimizationStats::default();

        // Find constant operations that can be folded
        let mut to_replace = Vec::new();

        for (node_id, node) in &graph.nodes {
            match node.operation {
                OperationType::Add
                | OperationType::Sub
                | OperationType::Mul
                | OperationType::Div => {
                    // Check if both inputs are constants
                    let input_nodes: Vec<_> = node
                        .inputs
                        .iter()
                        .filter_map(|id| graph.get_node(*id))
                        .collect();

                    if input_nodes.len() == 2
                        && input_nodes
                            .iter()
                            .all(|n| matches!(n.operation, OperationType::Constant))
                    {
                        // This could be constant folded
                        to_replace.push(*node_id);
                    }
                }
                _ => {}
            }
        }

        // Replace with constant nodes
        for node_id in to_replace {
            if let Some(node) = graph.get_node_mut(node_id) {
                node.operation = OperationType::Constant;
                node.inputs.clear();
                stats.nodes_eliminated += 1;
            }
        }

        stats.passes_applied = 1;
        Ok(stats)
    }

    fn name(&self) -> &'static str {
        "ConstantFolding"
    }
}

/// Operation fusion pass
pub struct OperationFusionPass;

impl OptimizationPass for OperationFusionPass {
    fn apply(&self, graph: &mut IRGraph) -> Result<OptimizationStats> {
        let mut stats = OptimizationStats::default();

        // Look for fusable patterns
        let mut fusion_groups = Vec::new();

        for (node_id, node) in &graph.nodes {
            // Look for simple fusion patterns
            if matches!(node.operation, OperationType::Add | OperationType::Mul)
                && let Some(next_nodes) = self.find_fusable_sequence(*node_id, graph)
            {
                fusion_groups.push(next_nodes);
            }
        }

        // Apply fusions
        for group in fusion_groups {
            if group.len() > 1 {
                self.fuse_operations(graph, &group);
                stats.operations_fused += (group.len() - 1) as u32;
            }
        }

        stats.passes_applied = 1;
        Ok(stats)
    }

    fn name(&self) -> &'static str {
        "OperationFusion"
    }
}

impl OperationFusionPass {
    fn find_fusable_sequence(&self, start_node: NodeId, graph: &IRGraph) -> Option<Vec<NodeId>> {
        let mut sequence = vec![start_node];
        let mut current = start_node;

        // Simple linear fusion for now
        while let Some(node) = graph.get_node(current) {
            if node.outputs.len() == 1 {
                let next_id = node.outputs[0];
                if let Some(next_node) = graph.get_node(next_id) {
                    if matches!(next_node.operation, OperationType::Add | OperationType::Mul) {
                        sequence.push(next_id);
                        current = next_id;
                    } else {
                        break;
                    }
                } else {
                    break;
                }
            } else {
                break;
            }
        }

        if sequence.len() > 1 {
            Some(sequence)
        } else {
            None
        }
    }

    fn fuse_operations(&self, graph: &mut IRGraph, nodes: &[NodeId]) {
        if nodes.is_empty() {
            return;
        }

        // Create a fused operation node
        let first_id = nodes[0];
        if let Some(first_node) = graph.get_node_mut(first_id) {
            // Mark as fused operation
            first_node
                .metadata
                .insert("fused".to_string(), "true".to_string());
            first_node
                .metadata
                .insert("fused_count".to_string(), nodes.len().to_string());
        }

        // Remove other nodes in the fusion group
        for &node_id in &nodes[1..] {
            graph.remove_node(node_id);
        }
    }
}

/// Memory optimization pass
pub struct MemoryOptimizationPass;

impl OptimizationPass for MemoryOptimizationPass {
    fn apply(&self, graph: &mut IRGraph) -> Result<OptimizationStats> {
        let mut stats = OptimizationStats::default();

        // Analyze memory access patterns
        let mut memory_ops = Vec::new();

        for (node_id, node) in &graph.nodes {
            match node.operation {
                OperationType::Load | OperationType::Store => {
                    memory_ops.push(*node_id);
                }
                _ => {}
            }
        }

        // Simple memory coalescing optimization
        for &node_id in &memory_ops {
            if let Some(node) = graph.get_node_mut(node_id) {
                node.metadata
                    .insert("memory_optimized".to_string(), "true".to_string());
            }
        }

        stats.memory_saved = (memory_ops.len() * 64) as u64; // Estimate
        stats.passes_applied = 1;
        Ok(stats)
    }

    fn name(&self) -> &'static str {
        "MemoryOptimization"
    }
}

/// Main optimizer that coordinates optimization passes
pub struct Optimizer {
    passes: Vec<Box<dyn OptimizationPass>>,
}

impl Optimizer {
    /// Create new optimizer
    pub fn new() -> Self {
        Self { passes: Vec::new() }
    }

    /// Add an optimization pass
    pub fn add_pass(&mut self, pass: Box<dyn OptimizationPass>) {
        self.passes.push(pass);
    }

    /// Optimize a graph
    pub fn optimize(&self, graph: &mut IRGraph) -> Result<OptimizationStats> {
        let mut total_stats = OptimizationStats::default();

        for pass in &self.passes {
            debug!("Applying optimization pass: {}", pass.name());
            let pass_stats = pass.apply(graph)?;

            total_stats.passes_applied += pass_stats.passes_applied;
            total_stats.nodes_eliminated += pass_stats.nodes_eliminated;
            total_stats.operations_fused += pass_stats.operations_fused;
            total_stats.memory_saved += pass_stats.memory_saved;
        }

        Ok(total_stats)
    }
}

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

/// Extended operation type for advanced optimizations
#[derive(Debug, Clone)]
pub enum ExtendedOperationType {
    Basic(OperationType),
    Fused {
        operations: Vec<NodeId>,
        pattern: String,
    },
    Optimized {
        original: Box<Self>,
        optimization: String,
    },
}