// Graph serialization support for unified data+compute graphs
// Phase 6: Self-Optimizing Platform

use indexmap::IndexMap;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

use crate::{Graph, GraphError, NodeId, Result};

/// Serializable graph format for persistence and interchange
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SerializableGraph {
    /// Graph format version for compatibility
    pub version: String,
    /// Graph metadata
    pub metadata: IndexMap<String, String>,
    /// All nodes in the graph
    pub nodes: Vec<SerializableNode>,
    /// All edges in the graph
    pub edges: Vec<SerializableEdge>,
    /// Timestamp of creation/serialization
    pub timestamp: chrono::DateTime<chrono::Utc>,
}

/// Serializable node representation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SerializableNode {
    /// Node identifier
    pub id: NodeId,
    /// Operation type and configuration
    pub operation: crate::OperationType,
    /// Node metadata
    pub metadata: IndexMap<String, String>,
}

/// Serializable edge representation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SerializableEdge {
    /// Source node ID
    pub from: NodeId,
    /// Target node ID
    pub to: NodeId,
    /// Edge metadata
    pub metadata: IndexMap<String, String>,
}

impl SerializableGraph {
    /// Current format version
    pub const CURRENT_VERSION: &'static str = "1.0.0";

    /// Create from graph
    pub fn from_graph(graph: &Graph) -> Self {
        let nodes: Vec<SerializableNode> = graph
            .nodes()
            .map(|node| SerializableNode {
                id: node.id,
                operation: node.operation.clone(),
                metadata: node.metadata.clone(),
            })
            .collect();

        let mut edges = Vec::new();
        for node in graph.nodes() {
            if let Ok(dependencies) = graph.dependencies(node.id) {
                for dep_id in dependencies {
                    edges.push(SerializableEdge {
                        from: dep_id,
                        to: node.id,
                        metadata: IndexMap::new(), // Could be extended with edge metadata
                    });
                }
            }
        }

        Self {
            version: Self::CURRENT_VERSION.to_string(),
            metadata: graph.metadata().clone(),
            nodes,
            edges,
            timestamp: chrono::Utc::now(),
        }
    }

    /// Convert to graph
    pub fn to_graph(&self) -> Result<Graph> {
        // Validate version compatibility
        if !self.is_compatible_version() {
            return Err(GraphError::serialization_error(format!(
                "Incompatible graph version: {}",
                self.version
            )));
        }

        let mut builder = crate::GraphBuilder::new();
        let mut node_mapping = HashMap::new();

        // Add all nodes
        for serializable_node in &self.nodes {
            let new_node_id = match &serializable_node.operation {
                crate::OperationType::Data(data_op) => builder.add_data_op(data_op.clone())?,
                crate::OperationType::Compute(compute_op) => {
                    // Need to remap node references
                    let remapped_op =
                        self.remap_compute_op_for_deserialization(compute_op, &node_mapping)?;
                    builder.add_compute_op(remapped_op)?
                }
            };
            node_mapping.insert(serializable_node.id, new_node_id);
        }

        // Add all edges
        for edge in &self.edges {
            let from_node = node_mapping.get(&edge.from).ok_or_else(|| {
                GraphError::serialization_error(format!("Missing source node: {}", edge.from))
            })?;
            let to_node = node_mapping.get(&edge.to).ok_or_else(|| {
                GraphError::serialization_error(format!("Missing target node: {}", edge.to))
            })?;

            builder.add_dependency(*to_node, *from_node)?;
        }

        let mut graph = builder.build()?;

        // Copy metadata
        for (key, value) in &self.metadata {
            graph.add_metadata(key.clone(), value.clone());
        }

        Ok(graph)
    }

    /// Check if version is compatible with current implementation
    fn is_compatible_version(&self) -> bool {
        // Simple version check - in real implementation would have more sophisticated logic
        let version_parts: Vec<&str> = self.version.split('.').collect();
        let current_parts: Vec<&str> = Self::CURRENT_VERSION.split('.').collect();

        if version_parts.len() >= 2 && current_parts.len() >= 2 {
            // Major version must match, minor version can be different
            version_parts[0] == current_parts[0]
        } else {
            false
        }
    }

    /// Remap compute operation node references for deserialization
    fn remap_compute_op_for_deserialization(
        &self,
        op: &crate::ComputeOp,
        node_mapping: &HashMap<NodeId, NodeId>,
    ) -> Result<crate::ComputeOp> {
        use crate::ComputeOp;

        let remap = |id: NodeId| -> Result<NodeId> {
            node_mapping.get(&id).copied().ok_or_else(|| {
                GraphError::serialization_error(format!("Missing node mapping for {id}"))
            })
        };

        let remapped_op = match op {
            ComputeOp::TensorLoad {
                data_node,
                shape,
                dtype,
            } => ComputeOp::TensorLoad {
                data_node: remap(*data_node)?,
                shape: shape.clone(),
                dtype: dtype.clone(),
            },
            ComputeOp::Add { lhs, rhs } => ComputeOp::Add {
                lhs: remap(*lhs)?,
                rhs: remap(*rhs)?,
            },
            ComputeOp::Mul { lhs, rhs } => ComputeOp::Mul {
                lhs: remap(*lhs)?,
                rhs: remap(*rhs)?,
            },
            ComputeOp::MatMul { lhs, rhs } => ComputeOp::MatMul {
                lhs: remap(*lhs)?,
                rhs: remap(*rhs)?,
            },
            ComputeOp::Conv2d {
                input,
                weight,
                bias,
                stride,
                padding,
                dilation,
                groups,
            } => ComputeOp::Conv2d {
                input: remap(*input)?,
                weight: remap(*weight)?,
                bias: bias.map(remap).transpose()?,
                stride: stride.clone(),
                padding: padding.clone(),
                dilation: dilation.clone(),
                groups: *groups,
            },
            ComputeOp::Activation {
                input,
                activation_type,
            } => ComputeOp::Activation {
                input: remap(*input)?,
                activation_type: activation_type.clone(),
            },
            ComputeOp::Pool2d {
                input,
                pool_type,
                kernel_size,
                stride,
            } => ComputeOp::Pool2d {
                input: remap(*input)?,
                pool_type: pool_type.clone(),
                kernel_size: kernel_size.clone(),
                stride: stride.clone(),
            },
            ComputeOp::BatchNorm {
                input,
                running_mean,
                running_var,
                weight,
                bias,
                training,
            } => ComputeOp::BatchNorm {
                input: remap(*input)?,
                running_mean: remap(*running_mean)?,
                running_var: remap(*running_var)?,
                weight: weight.map(remap).transpose()?,
                bias: bias.map(remap).transpose()?,
                training: *training,
            },
            ComputeOp::Reshape { input, shape } => ComputeOp::Reshape {
                input: remap(*input)?,
                shape: shape.clone(),
            },
            ComputeOp::Transpose { input, dims } => ComputeOp::Transpose {
                input: remap(*input)?,
                dims: dims.clone(),
            },
        };

        Ok(remapped_op)
    }

    /// Serialize to JSON string
    pub fn to_json(&self) -> Result<String> {
        serde_json::to_string_pretty(self)
            .map_err(|e| GraphError::serialization_error(e.to_string()))
    }

    /// Deserialize from JSON string
    pub fn from_json(json: &str) -> Result<Self> {
        serde_json::from_str(json).map_err(|e| GraphError::serialization_error(e.to_string()))
    }

    /// Serialize to binary format (MessagePack)
    pub fn to_binary(&self) -> Result<Vec<u8>> {
        // For now, use bincode as binary format
        bincode::serialize(self).map_err(|e| GraphError::serialization_error(e.to_string()))
    }

    /// Deserialize from binary format
    pub fn from_binary(data: &[u8]) -> Result<Self> {
        bincode::deserialize(data).map_err(|e| GraphError::serialization_error(e.to_string()))
    }

    /// Get format statistics
    pub fn stats(&self) -> SerializationStats {
        SerializationStats {
            node_count: self.nodes.len(),
            edge_count: self.edges.len(),
            data_ops: self
                .nodes
                .iter()
                .filter(|node| node.operation.is_data_op())
                .count(),
            compute_ops: self
                .nodes
                .iter()
                .filter(|node| node.operation.is_compute_op())
                .count(),
            version: self.version.clone(),
            timestamp: self.timestamp,
        }
    }
}

/// Statistics about serialized graph
#[derive(Debug, Clone)]
pub struct SerializationStats {
    /// Number of nodes
    pub node_count: usize,
    /// Number of edges
    pub edge_count: usize,
    /// Number of data operations
    pub data_ops: usize,
    /// Number of compute operations
    pub compute_ops: usize,
    /// Graph format version
    pub version: String,
    /// Serialization timestamp
    pub timestamp: chrono::DateTime<chrono::Utc>,
}

impl SerializationStats {
    /// Get total operation count
    pub fn total_ops(&self) -> usize {
        self.data_ops + self.compute_ops
    }

    /// Get data to compute ratio
    pub fn data_compute_ratio(&self) -> f64 {
        if self.compute_ops == 0 {
            f64::INFINITY
        } else {
            self.data_ops as f64 / self.compute_ops as f64
        }
    }
}

/// Graph export formats
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ExportFormat {
    /// JSON format (human readable)
    Json,
    /// Binary format (compact)
    Binary,
    /// GraphViz DOT format (visualization)
    Dot,
}

/// Graph exporter for different formats
pub struct GraphExporter;

impl GraphExporter {
    /// Export graph to specified format
    pub fn export(graph: &Graph, format: ExportFormat) -> Result<Vec<u8>> {
        match format {
            ExportFormat::Json => {
                let json = graph.to_json()?;
                Ok(json.into_bytes())
            }
            ExportFormat::Binary => {
                let serializable = SerializableGraph::from_graph(graph);
                serializable.to_binary()
            }
            ExportFormat::Dot => Self::export_dot(graph),
        }
    }

    /// Export to GraphViz DOT format for visualization
    fn export_dot(graph: &Graph) -> Result<Vec<u8>> {
        let mut dot = String::new();
        dot.push_str("digraph RustyTorchGraph {\n");
        dot.push_str("  rankdir=TB;\n");
        dot.push_str("  node [shape=box];\n");

        // Add nodes
        for node in graph.nodes() {
            let label = format!("{}\\n{}", node.id, node.operation.name());
            let color = if node.operation.is_data_op() {
                "lightblue"
            } else {
                "lightgreen"
            };

            dot.push_str(&format!(
                "  \"{}\" [label=\"{}\" fillcolor=\"{}\" style=\"filled\"];\n",
                node.id, label, color
            ));
        }

        // Add edges
        for node in graph.nodes() {
            if let Ok(dependencies) = graph.dependencies(node.id) {
                for dep_id in dependencies {
                    dot.push_str(&format!("  \"{}\" -> \"{}\";\n", dep_id, node.id));
                }
            }
        }

        dot.push_str("}\n");
        Ok(dot.into_bytes())
    }

    /// Import graph from specified format
    pub fn import(data: &[u8], format: ExportFormat) -> Result<Graph> {
        match format {
            ExportFormat::Json => {
                let json = String::from_utf8(data.to_vec())
                    .map_err(|e| GraphError::serialization_error(e.to_string()))?;
                Graph::from_json(&json)
            }
            ExportFormat::Binary => {
                let serializable = SerializableGraph::from_binary(data)?;
                serializable.to_graph()
            }
            ExportFormat::Dot => Err(GraphError::serialization_error(
                "DOT format import not supported".to_string(),
            )),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{ComputeOp, DataOp, GraphBuilder};

    #[test]
    fn test_graph_serialization() {
        let mut builder = GraphBuilder::new();

        let read_node = builder
            .add_data_op(DataOp::Read {
                source: "test.csv".to_string(),
                format: "csv".to_string(),
                schema: None,
            })
            .expect("Failed to add read node");

        let tensor_node = builder
            .add_compute_op(ComputeOp::TensorLoad {
                data_node: read_node,
                shape: vec![100, 50],
                dtype: "f32".to_string(),
            })
            .expect("Failed to add tensor node");

        builder
            .add_dependency(tensor_node, read_node)
            .expect("Failed to add dependency");

        let graph = builder.build().expect("Failed to build graph");

        // Test JSON serialization
        let json = graph.to_json().expect("Failed to serialize to JSON");
        assert!(json.contains("Read"));
        assert!(json.contains("TensorLoad"));

        let deserialized = Graph::from_json(&json).expect("Failed to deserialize from JSON");
        assert_eq!(deserialized.node_count(), graph.node_count());
        assert_eq!(deserialized.edge_count(), graph.edge_count());
    }

    #[test]
    fn test_serializable_graph() {
        let mut builder = GraphBuilder::new();

        let _read_node = builder
            .add_data_op(DataOp::Read {
                source: "data.parquet".to_string(),
                format: "parquet".to_string(),
                schema: Some("id:i32,value:f64".to_string()),
            })
            .expect("Failed to add read node");

        let graph = builder.build().expect("Failed to build graph");

        let serializable = SerializableGraph::from_graph(&graph);
        assert_eq!(serializable.nodes.len(), 1);
        assert_eq!(serializable.edges.len(), 0);
        assert_eq!(serializable.version, SerializableGraph::CURRENT_VERSION);

        let restored = serializable.to_graph().expect("Failed to restore graph");
        assert_eq!(restored.node_count(), graph.node_count());
    }

    #[test]
    fn test_export_formats() {
        let mut builder = GraphBuilder::new();

        let _read_node = builder
            .add_data_op(DataOp::Read {
                source: "test.csv".to_string(),
                format: "csv".to_string(),
                schema: None,
            })
            .expect("Failed to add read node");

        let graph = builder.build().expect("Failed to build graph");

        // Test JSON export
        let json_data =
            GraphExporter::export(&graph, ExportFormat::Json).expect("Failed to export to JSON");
        assert!(!json_data.is_empty());

        // Test binary export
        let binary_data = GraphExporter::export(&graph, ExportFormat::Binary)
            .expect("Failed to export to binary");
        assert!(!binary_data.is_empty());

        // Test DOT export
        let dot_data =
            GraphExporter::export(&graph, ExportFormat::Dot).expect("Failed to export to DOT");
        let dot_str = String::from_utf8(dot_data).expect("Invalid DOT format");
        assert!(dot_str.contains("digraph"));
        assert!(dot_str.contains("Read"));
    }

    #[test]
    fn test_version_compatibility() {
        let serializable = SerializableGraph {
            version: "1.0.0".to_string(),
            metadata: IndexMap::new(),
            nodes: Vec::new(),
            edges: Vec::new(),
            timestamp: chrono::Utc::now(),
        };

        assert!(serializable.is_compatible_version());

        let incompatible = SerializableGraph {
            version: "2.0.0".to_string(),
            ..serializable.clone()
        };

        assert!(!incompatible.is_compatible_version());
    }

    #[test]
    fn test_serialization_stats() {
        let mut builder = GraphBuilder::new();

        let read_node = builder
            .add_data_op(DataOp::Read {
                source: "test.csv".to_string(),
                format: "csv".to_string(),
                schema: None,
            })
            .expect("Failed to add read node");

        let tensor_node = builder
            .add_compute_op(ComputeOp::TensorLoad {
                data_node: read_node,
                shape: vec![10, 10],
                dtype: "f32".to_string(),
            })
            .expect("Failed to add tensor node");

        builder
            .add_dependency(tensor_node, read_node)
            .expect("Failed to add dependency");

        let graph = builder.build().expect("Failed to build graph");
        let serializable = SerializableGraph::from_graph(&graph);
        let stats = serializable.stats();

        assert_eq!(stats.node_count, 2);
        assert_eq!(stats.edge_count, 1);
        assert_eq!(stats.data_ops, 1);
        assert_eq!(stats.compute_ops, 1);
        assert_eq!(stats.total_ops(), 2);
        assert_eq!(stats.data_compute_ratio(), 1.0);
    }
}