//! Graph IR parsing implementations.

use anyhow::Result;
use serde_json::Value;

use crate::aot::{self, GraphOperation};
use crate::templates;

use super::SynthesisEngine;

impl SynthesisEngine {
    /// Advanced graph IR parser supporting multiple formats (JSON, ONNX-like, custom)
    pub(crate) fn parse_graph_ir(&self, graph_ir: &str) -> Result<Vec<GraphOperation>> {
        tracing::debug!("Parsing graph IR: {} characters", graph_ir.len());

        // Try to parse as JSON first
        if let Ok(json_value) = serde_json::from_str::<Value>(graph_ir) {
            return self.parse_json_ir(&json_value);
        }

        // Try to parse as YAML
        #[cfg(feature = "yaml-parsing")]
        if let Ok(yaml_value) = serde_yaml::from_str::<Value>(graph_ir) {
            return self.parse_json_ir(&yaml_value); // Reuse JSON parser for YAML
        }

        // Try custom DSL format
        if graph_ir.trim_start().starts_with("graph") {
            return self.parse_custom_dsl(graph_ir);
        }

        // Try ONNX protobuf format detection
        if graph_ir.starts_with("ir_version") || graph_ir.contains("ModelProto") {
            return self.parse_onnx_like(graph_ir);
        }

        // Fallback: return error for unrecognized format
        Err(anyhow::anyhow!(
            "Unrecognized graph IR format. Supported formats: JSON, YAML, Custom DSL, ONNX-like"
        ))
    }

    /// Parse JSON/YAML format graph IR
    pub(crate) fn parse_json_ir(&self, json: &Value) -> Result<Vec<GraphOperation>> {
        let mut operations = Vec::new();

        // Handle different JSON structures
        if let Some(graph) = json.get("graph") {
            return self.parse_json_graph_object(graph);
        }

        if let Some(nodes) = json.get("nodes").and_then(|v| v.as_array()) {
            for node in nodes {
                if let Some(op) = self.parse_json_node(node)? {
                    operations.push(op);
                }
            }
        } else if json.is_array() {
            // Direct array of operations
            let nodes = json.as_array().unwrap();
            for node in nodes {
                if let Some(op) = self.parse_json_node(node)? {
                    operations.push(op);
                }
            }
        } else {
            // Single operation
            if let Some(op) = self.parse_json_node(json)? {
                operations.push(op);
            }
        }

        Ok(operations)
    }

    /// Parse a graph object from JSON
    pub(crate) fn parse_json_graph_object(&self, graph: &Value) -> Result<Vec<GraphOperation>> {
        let mut operations = Vec::new();

        if let Some(nodes) = graph.get("nodes").and_then(|v| v.as_array()) {
            for node in nodes {
                if let Some(op) = self.parse_json_node(node)? {
                    operations.push(op);
                }
            }
        }

        // Handle edges/connections if present
        if let Some(_edges) = graph.get("edges") {
            // In a real implementation, would use edges to set up input/output connections
            tracing::debug!(
                "Graph contains edge information (connection parsing not fully implemented)"
            );
        }

        Ok(operations)
    }

    /// Parse individual JSON node into GraphOperation
    pub(crate) fn parse_json_node(&self, node: &Value) -> Result<Option<GraphOperation>> {
        let id = node
            .get("id")
            .or_else(|| node.get("name"))
            .and_then(|v| v.as_str())
            .unwrap_or(&format!("op_{}", uuid::Uuid::new_v4()))
            .to_string();

        let op_type = node
            .get("op_type")
            .or_else(|| node.get("type"))
            .or_else(|| node.get("operation"))
            .and_then(|v| v.as_str());

        let inputs = node.get("inputs").and_then(|v| v.as_array()).map_or_else(
            || vec!["input".to_string()],
            |arr| {
                arr.iter()
                    .filter_map(|v| v.as_str().map(std::string::ToString::to_string))
                    .collect()
            },
        );

        let outputs = node.get("outputs").and_then(|v| v.as_array()).map_or_else(
            || vec!["output".to_string()],
            |arr| {
                arr.iter()
                    .filter_map(|v| v.as_str().map(std::string::ToString::to_string))
                    .collect()
            },
        );

        if let Some(op_type_str) = op_type {
            let kernel_op = self.parse_operation_type(op_type_str, node)?;

            let fusable = node
                .get("fusable")
                .and_then(serde_json::Value::as_bool)
                .unwrap_or(true);

            let memory_layout = self.parse_memory_layout(node)?;

            return Ok(Some(GraphOperation {
                id,
                operation: kernel_op,
                inputs,
                outputs,
                memory_layout,
                fusable,
            }));
        }

        tracing::warn!("Skipping node without operation type: {}", id);
        Ok(None)
    }

    /// Parse operation type from string and attributes
    pub(crate) fn parse_operation_type(
        &self,
        op_type: &str,
        node: &Value,
    ) -> Result<templates::KernelOperation> {
        let attrs = node.get("attributes").or_else(|| node.get("attrs"));

        match op_type.to_lowercase().as_str() {
            "gemm" | "matmul" | "dense" => {
                let m = self.get_int_attr(attrs, "M", 128)? as usize;
                let n = self.get_int_attr(attrs, "N", 128)? as usize;
                let k = self.get_int_attr(attrs, "K", 128)? as usize;
                let transpose_a = self.get_bool_attr(attrs, "transpose_a", false)?;
                let transpose_b = self.get_bool_attr(attrs, "transpose_b", false)?;

                Ok(templates::KernelOperation::Gemm {
                    m: m as u32,
                    n: n as u32,
                    k: k as u32,
                    transpose_a,
                    transpose_b,
                })
            }

            "conv2d" | "convolution" => {
                let batch_size = self.get_int_attr(attrs, "batch_size", 1)? as usize;
                let in_channels = self.get_int_attr(attrs, "in_channels", 64)? as usize;
                let out_channels = self.get_int_attr(attrs, "out_channels", 64)? as usize;
                let height = self.get_int_attr(attrs, "height", 224)? as usize;
                let width = self.get_int_attr(attrs, "width", 224)? as usize;
                let kernel_size = self.get_int_attr(attrs, "kernel_size", 3)? as usize;

                Ok(templates::KernelOperation::Convolution {
                    batch_size: batch_size as u32,
                    in_channels: in_channels as u32,
                    out_channels: out_channels as u32,
                    height: height as u32,
                    width: width as u32,
                    kernel_size: kernel_size as u32,
                })
            }

            "attention" | "multiheadattention" => {
                let sequence_length = self.get_int_attr(attrs, "sequence_length", 512)? as usize;
                let head_dim = self.get_int_attr(attrs, "head_dim", 64)? as usize;
                let num_heads = self.get_int_attr(attrs, "num_heads", 8)? as usize;

                Ok(templates::KernelOperation::Attention {
                    sequence_length: sequence_length as u32,
                    head_dim: head_dim as u32,
                    num_heads: num_heads as u32,
                })
            }

            "add" | "sub" | "mul" | "div" | "relu" | "sigmoid" | "tanh" => {
                let size = self.get_int_attr(attrs, "size", 1024)? as usize;
                let elementwise_op = match op_type.to_lowercase().as_str() {
                    "add" => templates::ElementwiseOp::Add,
                    "sub" => templates::ElementwiseOp::Sub,
                    "mul" => templates::ElementwiseOp::Mul,
                    "div" => templates::ElementwiseOp::Div,
                    "relu" => templates::ElementwiseOp::ReLU,
                    "sigmoid" => templates::ElementwiseOp::Sigmoid,
                    "tanh" => templates::ElementwiseOp::Tanh,
                    _ => templates::ElementwiseOp::Add, // default
                };

                Ok(templates::KernelOperation::Elementwise {
                    operation: elementwise_op,
                    size: size as u32,
                })
            }

            "sum" | "mean" | "max" | "min" => {
                let input_size = self.get_int_attr(attrs, "input_size", 1024)? as usize;
                let axis = self.get_int_attr(attrs, "axis", 0)? as usize;
                let reduction_op = match op_type.to_lowercase().as_str() {
                    "sum" => templates::ReductionOp::Sum,
                    "mean" => templates::ReductionOp::Mean,
                    "max" => templates::ReductionOp::Max,
                    "min" => templates::ReductionOp::Min,
                    _ => templates::ReductionOp::Sum, // default
                };

                Ok(templates::KernelOperation::Reduction {
                    operation: reduction_op,
                    input_size: input_size as u32,
                    axis: axis as u32,
                })
            }

            _ => {
                tracing::warn!(
                    "Unknown operation type '{}', defaulting to elementwise add",
                    op_type
                );
                Ok(templates::KernelOperation::Elementwise {
                    operation: templates::ElementwiseOp::Add,
                    size: 1024,
                })
            }
        }
    }

    /// Parse memory layout from node attributes
    pub(crate) fn parse_memory_layout(&self, node: &Value) -> Result<aot::MemoryLayout> {
        let attrs = node.get("attributes").or_else(|| node.get("attrs"));

        let data_type = self.get_string_attr(attrs, "data_type", "f32")?;
        let alignment = self.get_int_attr(attrs, "alignment", 32)? as usize;

        Ok(aot::MemoryLayout {
            dtype: data_type,
            shape: vec![1],    // Default shape
            strides: vec![1],  // Default strides
            alignment: alignment as u32,
            contiguous: true,  // Default assumption
        })
    }

    /// Parse custom DSL format
    pub(crate) fn parse_custom_dsl(&self, dsl: &str) -> Result<Vec<GraphOperation>> {
        let mut operations = Vec::new();
        let lines: Vec<&str> = dsl.lines().collect();

        let mut current_op: Option<GraphOperation> = None;

        for line in lines {
            let line = line.trim();
            if line.is_empty() || line.starts_with('#') {
                continue;
            }

            if line.starts_with("op ") {
                // Save previous operation
                if let Some(op) = current_op.take() {
                    operations.push(op);
                }

                // Parse operation declaration: "op operation_name type"
                let parts: Vec<&str> = line.split_whitespace().collect();
                if parts.len() >= 3 {
                    let op_id = parts[1].to_string();
                    let op_type = parts[2];

                    let operation = match op_type {
                        "gemm" => templates::KernelOperation::Gemm {
                            m: 128,
                            n: 128,
                            k: 128,
                            transpose_a: false,
                            transpose_b: false,
                        },
                        "conv2d" => templates::KernelOperation::Convolution {
                            batch_size: 1,
                            in_channels: 64,
                            out_channels: 64,
                            height: 224,
                            width: 224,
                            kernel_size: 3,
                        },
                        "add" => templates::KernelOperation::Elementwise {
                            operation: templates::ElementwiseOp::Add,
                            size: 1024,
                        },
                        _ => templates::KernelOperation::Elementwise {
                            operation: templates::ElementwiseOp::Add,
                            size: 1024,
                        },
                    };

                    current_op = Some(GraphOperation {
                        id: op_id,
                        operation,
                        inputs: vec!["input".to_string()],
                        outputs: vec!["output".to_string()],
                        memory_layout: aot::MemoryLayout::default(),
                        fusable: true,
                    });
                }
            } else if line.starts_with("  ") && current_op.is_some() {
                // Parse operation attributes
                let attr_line = line.trim();
                if let Some(op) = current_op.as_mut() {
                    self.parse_dsl_attribute(op, attr_line)?;
                }
            }
        }

        // Save final operation
        if let Some(op) = current_op {
            operations.push(op);
        }

        Ok(operations)
    }

    /// Parse DSL attribute line
    pub(crate) fn parse_dsl_attribute(&self, op: &mut GraphOperation, attr_line: &str) -> Result<()> {
        if let Some((key, value)) = attr_line.split_once(':') {
            let key = key.trim();
            let value = value.trim();

            match key {
                "inputs" => {
                    op.inputs = value.split(',').map(|s| s.trim().to_string()).collect();
                }
                "outputs" => {
                    op.outputs = value.split(',').map(|s| s.trim().to_string()).collect();
                }
                "fusable" => {
                    op.fusable = value.parse().unwrap_or(true);
                }
                _ => {
                    tracing::debug!("Unknown DSL attribute: {} = {}", key, value);
                }
            }
        }
        Ok(())
    }

    /// Parse ONNX-like format (simplified)
    pub(crate) fn parse_onnx_like(&self, _onnx_ir: &str) -> Result<Vec<GraphOperation>> {
        // This would require a proper ONNX protobuf parser
        // For now, return a placeholder implementation
        tracing::warn!("ONNX parsing not fully implemented, returning default operation");

        Ok(vec![GraphOperation {
            id: "onnx_op".to_string(),
            operation: templates::KernelOperation::Elementwise {
                operation: templates::ElementwiseOp::Add,
                size: 1024,
            },
            inputs: vec!["onnx_input".to_string()],
            outputs: vec!["onnx_output".to_string()],
            memory_layout: aot::MemoryLayout::default(),
            fusable: true,
        }])
    }

    // Helper methods for attribute parsing
    pub(crate) fn get_int_attr(&self, attrs: Option<&Value>, key: &str, default: i64) -> Result<i64> {
        if let Some(attrs) = attrs
            && let Some(value) = attrs.get(key)
        {
            if let Some(num) = value.as_i64() {
                return Ok(num);
            }
            if let Some(s) = value.as_str() {
                return s
                    .parse()
                    .map_err(|_| anyhow::anyhow!("Invalid integer value for {key}: {s}"));
            }
        }
        Ok(default)
    }

    pub(crate) fn get_bool_attr(&self, attrs: Option<&Value>, key: &str, default: bool) -> Result<bool> {
        if let Some(attrs) = attrs
            && let Some(value) = attrs.get(key)
        {
            if let Some(b) = value.as_bool() {
                return Ok(b);
            }
            if let Some(s) = value.as_str() {
                return s
                    .parse()
                    .map_err(|_| anyhow::anyhow!("Invalid boolean value for {key}: {s}"));
            }
        }
        Ok(default)
    }

    pub(crate) fn get_string_attr(&self, attrs: Option<&Value>, key: &str, default: &str) -> Result<String> {
        if let Some(attrs) = attrs
            && let Some(value) = attrs.get(key)
            && let Some(s) = value.as_str()
        {
            return Ok(s.to_string());
        }
        Ok(default.to_string())
    }
}