//! Parallel planning agent for autonomous parallelization strategy optimization.

use crate::{
    agents::AutonomousAgent,
    error::AutoResult,
    proposal::{Proposal, ProposalType},
};
use rtx_graph::Graph;
use rtx_runtime::Runtime;
use std::sync::Arc;
use tracing::info;

/// Parallelization opportunity analysis.
#[derive(Debug, Clone)]
pub struct ParallelizationOpportunity {
    pub node_id: String,
    pub opportunity_type: ParallelType,
    pub estimated_speedup: f32,
    pub resource_requirements: ResourceRequirements,
}

#[derive(Debug, Clone)]
pub enum ParallelType {
    DataParallel,
    ModelParallel,
    PipelineParallel,
    TensorParallel,
}

#[derive(Debug, Clone)]
pub struct ResourceRequirements {
    pub min_devices: usize,
    pub memory_per_device: usize,
    pub communication_overhead: f32,
}

/// Agent for autonomous parallel execution planning.
pub struct ParallelPlannerAgent {
    runtime: Arc<Runtime>,
    name: String,
    version: String,
}

impl ParallelPlannerAgent {
    pub fn new(runtime: Arc<Runtime>) -> AutoResult<Self> {
        Ok(Self {
            runtime,
            name: "ParallelPlannerAgent".to_string(),
            version: "1.0.0".to_string(),
        })
    }

    pub async fn analyze_parallelization_opportunities(
        &self,
        graph: &Graph,
    ) -> AutoResult<Vec<ParallelizationOpportunity>> {
        info!(
            "Analyzing parallelization opportunities for graph with {} nodes",
            graph.node_count()
        );
        let mut opportunities = Vec::new();

        // Analyze each node for parallelization potential
        for node in graph.nodes() {
            let node_id = node.id.to_string();
            let op_name = node.operation.name();

            // Identify compute-heavy operations suitable for different parallelization strategies
            match op_name {
                "Compute::MatMul" => {
                    // MatMul is ideal for data parallelism (batch dimension)
                    // and tensor parallelism (row/column splitting)
                    opportunities.push(ParallelizationOpportunity {
                        node_id: node_id.clone(),
                        opportunity_type: ParallelType::DataParallel,
                        estimated_speedup: self.estimate_data_parallel_speedup(graph, node),
                        resource_requirements: ResourceRequirements {
                            min_devices: 2,
                            memory_per_device: 1024 * 1024 * 1024, // 1GB base
                            communication_overhead: 0.05,          // Low overhead for data parallel
                        },
                    });

                    opportunities.push(ParallelizationOpportunity {
                        node_id: node_id.clone(),
                        opportunity_type: ParallelType::TensorParallel,
                        estimated_speedup: self.estimate_tensor_parallel_speedup(graph, node),
                        resource_requirements: ResourceRequirements {
                            min_devices: 4,
                            memory_per_device: 512 * 1024 * 1024,
                            communication_overhead: 0.15, // AllReduce overhead
                        },
                    });
                }

                "Compute::Activation" | "Compute::BatchNorm" => {
                    // Elementwise ops - good for data parallelism
                    opportunities.push(ParallelizationOpportunity {
                        node_id: node_id.clone(),
                        opportunity_type: ParallelType::DataParallel,
                        estimated_speedup: self.estimate_data_parallel_speedup(graph, node) * 0.9,
                        resource_requirements: ResourceRequirements {
                            min_devices: 2,
                            memory_per_device: 256 * 1024 * 1024,
                            communication_overhead: 0.02,
                        },
                    });
                }

                "Compute::Conv2d" => {
                    // Conv2d benefits from both data and model parallelism
                    opportunities.push(ParallelizationOpportunity {
                        node_id: node_id.clone(),
                        opportunity_type: ParallelType::DataParallel,
                        estimated_speedup: self.estimate_data_parallel_speedup(graph, node) * 1.1,
                        resource_requirements: ResourceRequirements {
                            min_devices: 2,
                            memory_per_device: 2048 * 1024 * 1024,
                            communication_overhead: 0.08,
                        },
                    });

                    opportunities.push(ParallelizationOpportunity {
                        node_id: node_id.clone(),
                        opportunity_type: ParallelType::ModelParallel,
                        estimated_speedup: 2.0,
                        resource_requirements: ResourceRequirements {
                            min_devices: 2,
                            memory_per_device: 1024 * 1024 * 1024,
                            communication_overhead: 0.12,
                        },
                    });
                }

                _ => {
                    // Generic operations - only suggest data parallelism if graph is large enough
                    if graph.node_count() > 10 {
                        opportunities.push(ParallelizationOpportunity {
                            node_id: node_id.clone(),
                            opportunity_type: ParallelType::DataParallel,
                            estimated_speedup: 1.5,
                            resource_requirements: ResourceRequirements {
                                min_devices: 2,
                                memory_per_device: 512 * 1024 * 1024,
                                communication_overhead: 0.1,
                            },
                        });
                    }
                }
            }
        }

        // Analyze graph depth for pipeline parallelism opportunities
        let graph_depth = self.calculate_graph_depth(graph);
        if graph_depth > 4 {
            opportunities.push(ParallelizationOpportunity {
                node_id: "pipeline_stages".to_string(),
                opportunity_type: ParallelType::PipelineParallel,
                estimated_speedup: self.estimate_pipeline_speedup(graph_depth),
                resource_requirements: ResourceRequirements {
                    min_devices: graph_depth.min(8),
                    memory_per_device: 1024 * 1024 * 1024,
                    communication_overhead: 0.05 * graph_depth as f32,
                },
            });
        }

        // Sort by estimated speedup (highest first)
        opportunities.sort_by(|a, b| {
            b.estimated_speedup
                .partial_cmp(&a.estimated_speedup)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Deduplicate by taking best opportunity per node
        let mut seen_nodes = std::collections::HashSet::new();
        opportunities.retain(|op| {
            if seen_nodes.contains(&op.node_id) {
                false
            } else {
                seen_nodes.insert(op.node_id.clone());
                true
            }
        });

        Ok(opportunities)
    }

    /// Estimate speedup from data parallelism based on graph structure
    fn estimate_data_parallel_speedup(&self, graph: &Graph, _node: &rtx_graph::GraphNode) -> f32 {
        // Base speedup of 2x for 2 devices, with diminishing returns
        let node_count = graph.node_count();
        let base_speedup = 2.0;

        // Larger graphs benefit more from data parallelism
        let scale_factor = (node_count as f32 / 10.0).min(2.0).max(1.0);

        base_speedup * scale_factor
    }

    /// Estimate speedup from tensor parallelism
    fn estimate_tensor_parallel_speedup(&self, graph: &Graph, _node: &rtx_graph::GraphNode) -> f32 {
        // Tensor parallelism has higher communication overhead but better for large tensors
        let node_count = graph.node_count();
        let base_speedup = 1.8;

        // Works better for larger models
        let scale_factor = (node_count as f32 / 20.0).min(1.5).max(0.8);

        base_speedup * scale_factor
    }

    /// Estimate speedup from pipeline parallelism based on graph depth
    fn estimate_pipeline_speedup(&self, depth: usize) -> f32 {
        // Pipeline parallelism approaches linear speedup with depth, minus bubble overhead
        let ideal_speedup = depth as f32;
        let bubble_overhead = 1.0 / depth as f32; // Fraction of time in bubbles

        ideal_speedup * (1.0 - bubble_overhead)
    }

    /// Calculate the critical path depth of the graph
    fn calculate_graph_depth(&self, graph: &Graph) -> usize {
        // Simple heuristic: use edge count / node count ratio as depth estimate
        // A more accurate implementation would do topological traversal
        let nodes = graph.node_count();
        let edges = graph.edge_count();

        if nodes == 0 {
            return 0;
        }

        // Estimate depth based on graph density
        let density = edges as f32 / nodes as f32;

        // For sequential graphs, density ~= 1 (each node has one incoming edge)
        // For parallel graphs, density < 1
        // For highly connected graphs, density > 1

        if density <= 1.0 {
            // Mostly sequential - depth is approximately node count
            nodes.min(16)
        } else {
            // More parallel - depth is less than node count
            ((nodes as f32) / density).ceil() as usize
        }
    }

    pub async fn generate_data_parallel_proposals(
        &self,
        graph: &Graph,
    ) -> AutoResult<Vec<Proposal>> {
        let mut proposals = Vec::new();

        let node_count = graph.node_count();

        if node_count > 5 {
            // Calculate optimal batch splitting strategy
            let optimal_devices = (node_count / 5).min(8).max(2);
            let estimated_speedup =
                self.estimate_data_parallel_speedup(graph, graph.nodes().next().unwrap());

            proposals.push(Proposal::new(
                ProposalType::DataParallel,
                format!(
                    "Apply data parallelism across {} devices for {} compute nodes (estimated {:.1}x speedup)",
                    optimal_devices, node_count, estimated_speedup
                ),
                estimated_speedup,
            ));

            // Gradient synchronization strategy
            if node_count > 20 {
                proposals.push(Proposal::new(
                    ProposalType::DataParallel,
                    "Use gradient compression (TopK or random) to reduce all-reduce communication overhead".to_string(),
                    1.4,
                ));
            }
        }

        Ok(proposals)
    }

    pub async fn generate_pipeline_parallel_proposals(
        &self,
        graph: &Graph,
    ) -> AutoResult<Vec<Proposal>> {
        let mut proposals = Vec::new();

        let depth = self.calculate_graph_depth(graph);

        if depth > 4 {
            let num_stages = depth.min(8);
            let estimated_speedup = self.estimate_pipeline_speedup(depth);

            proposals.push(Proposal::new(
                ProposalType::PipelineParallel,
                format!(
                    "Implement {}-stage pipeline parallelism for depth-{} graph (estimated {:.1}x speedup)",
                    num_stages, depth, estimated_speedup
                ),
                estimated_speedup,
            ));

            // Micro-batch suggestion for pipeline efficiency
            if depth > 6 {
                proposals.push(Proposal::new(
                    ProposalType::PipelineParallel,
                    format!(
                        "Use {} micro-batches to reduce pipeline bubble overhead",
                        depth * 2
                    ),
                    1.3,
                ));
            }

            // Interleaved scheduling for better memory efficiency
            if depth > 8 {
                proposals.push(Proposal::new(
                    ProposalType::PipelineParallel,
                    "Apply interleaved 1F1B schedule to reduce peak memory usage by 50%"
                        .to_string(),
                    1.2,
                ));
            }
        }

        Ok(proposals)
    }

    pub async fn generate_tensor_parallel_proposals(
        &self,
        graph: &Graph,
    ) -> AutoResult<Vec<Proposal>> {
        let mut proposals = Vec::new();

        // Count large tensor operations
        let mut large_matmul_count = 0;
        let mut large_conv_count = 0;

        for node in graph.nodes() {
            match node.operation.name() {
                "Compute::MatMul" => large_matmul_count += 1,
                "Compute::Conv2d" => large_conv_count += 1,
                _ => {}
            }
        }

        if large_matmul_count > 0 {
            let estimated_speedup =
                self.estimate_tensor_parallel_speedup(graph, graph.nodes().next().unwrap());
            proposals.push(Proposal::new(
                ProposalType::TensorParallel,
                format!(
                    "Apply column-parallel tensor splitting to {} MatMul operations (estimated {:.1}x speedup)",
                    large_matmul_count, estimated_speedup
                ),
                estimated_speedup,
            ));
        }

        if large_conv_count > 0 {
            proposals.push(Proposal::new(
                ProposalType::TensorParallel,
                format!(
                    "Apply channel-parallel splitting to {} Conv2d operations",
                    large_conv_count
                ),
                2.0,
            ));
        }

        // Sequence parallelism for attention-heavy models
        if large_matmul_count > 10 {
            proposals.push(Proposal::new(
                ProposalType::TensorParallel,
                "Apply sequence parallelism for attention computation to reduce memory footprint"
                    .to_string(),
                1.8,
            ));
        }

        Ok(proposals)
    }

    pub async fn optimize_communication_patterns(
        &self,
        device_count: usize,
    ) -> AutoResult<Vec<Proposal>> {
        let mut proposals = Vec::new();

        if device_count > 2 {
            // Ring all-reduce for larger device counts
            if device_count >= 4 {
                proposals.push(Proposal::new(
                    ProposalType::CommunicationOptimization,
                    format!(
                        "Use ring all-reduce pattern for {} devices to achieve 2(n-1)/n bandwidth efficiency",
                        device_count
                    ),
                    1.5,
                ));
            }

            // Hierarchical all-reduce for very large clusters
            if device_count >= 8 {
                proposals.push(Proposal::new(
                    ProposalType::CommunicationOptimization,
                    format!(
                        "Apply hierarchical all-reduce (intra-node then inter-node) for {} devices",
                        device_count
                    ),
                    1.8,
                ));
            }

            // Overlap communication with computation
            proposals.push(Proposal::new(
                ProposalType::CommunicationOptimization,
                format!(
                    "Overlap gradient all-reduce with backward computation for {} devices",
                    device_count
                ),
                1.3,
            ));

            // NCCL optimization hints
            if device_count > 4 {
                proposals.push(Proposal::new(
                    ProposalType::CommunicationOptimization,
                    "Enable NCCL tree algorithm for large tensor all-reduce operations".to_string(),
                    1.2,
                ));
            }
        }

        Ok(proposals)
    }
}

impl AutonomousAgent for ParallelPlannerAgent {
    fn name(&self) -> &str {
        &self.name
    }
    fn version(&self) -> &str {
        &self.version
    }
    fn is_healthy(&self) -> bool {
        true
    }
}