use std::collections::HashMap;
use std::io::Write;

use petgraph::stable_graph::EdgeIndex;
use petgraph::visit::{EdgeRef, IntoEdgeReferences};
use serde::Serialize;

use crate::graph::types::*;

#[derive(Serialize)]
struct JsonGraph {
    nodes: Vec<JsonNode>,
    edges: Vec<JsonEdge>,
}

#[derive(Serialize)]
struct JsonNode {
    unique_id: String,
    label: String,
    node_type: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    file_path: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    description: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    materialization: Option<String>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    tags: Vec<String>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    columns: Vec<String>,
}

#[derive(Serialize)]
struct JsonEdge {
    source: String,
    target: String,
    edge_type: String,
    /// Number of intermediate nodes folded into this edge by `--collapse`.
    /// Omitted when 0 or when collapse was not used.
    #[serde(skip_serializing_if = "Option::is_none")]
    via_hops: Option<usize>,
}

/// Render the lineage graph as JSON to stdout
pub fn render_json(graph: &LineageGraph) {
    render_json_to_writer(graph, None, &mut std::io::stdout().lock());
}

/// Render with optional per-edge `via_hops` annotations (from collapse).
pub fn render_json_with_via(graph: &LineageGraph, via_hops: &HashMap<EdgeIndex, usize>) {
    render_json_to_writer(graph, Some(via_hops), &mut std::io::stdout().lock());
}

fn render_json_to_writer<W: Write>(
    graph: &LineageGraph,
    via_hops: Option<&HashMap<EdgeIndex, usize>>,
    w: &mut W,
) {
    let nodes: Vec<JsonNode> = graph
        .node_indices()
        .map(|idx| {
            let node = &graph[idx];
            JsonNode {
                unique_id: node.unique_id.clone(),
                label: node.label.clone(),
                node_type: node.node_type.label().to_string(),
                file_path: node.file_path.as_ref().map(|p| p.to_string_lossy().into()),
                description: node.description.clone(),
                materialization: node.materialization.clone(),
                tags: node.tags.clone(),
                columns: node.columns.clone(),
            }
        })
        .collect();

    let edges: Vec<JsonEdge> = graph
        .edge_references()
        .map(|edge| {
            let source = &graph[edge.source()];
            let target = &graph[edge.target()];
            let via = via_hops
                .and_then(|m| m.get(&edge.id()).copied())
                .filter(|&n| n > 0);
            JsonEdge {
                source: source.unique_id.clone(),
                target: target.unique_id.clone(),
                edge_type: edge_type_label(edge.weight().edge_type),
                via_hops: via,
            }
        })
        .collect();

    let json_graph = JsonGraph { nodes, edges };
    serde_json::to_writer_pretty(&mut *w, &json_graph).unwrap();
    writeln!(w).unwrap();
}

fn edge_type_label(edge_type: EdgeType) -> String {
    match edge_type {
        EdgeType::Ref => "ref",
        EdgeType::Source => "source",
        EdgeType::Test => "test",
        EdgeType::Exposure => "exposure",
    }
    .to_string()
}

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

    fn make_node(unique_id: &str, label: &str, node_type: NodeType) -> NodeData {
        NodeData {
            unique_id: unique_id.into(),
            label: label.into(),
            node_type,
            file_path: None,
            description: None,
            materialization: None,
            tags: vec![],
            columns: vec![],
        }
    }

    fn render_to_string(graph: &LineageGraph) -> String {
        let mut buf = Vec::new();
        render_json_to_writer(graph, None, &mut buf);
        String::from_utf8(buf).unwrap()
    }

    fn render_with_via(graph: &LineageGraph, via: &HashMap<EdgeIndex, usize>) -> String {
        let mut buf = Vec::new();
        render_json_to_writer(graph, Some(via), &mut buf);
        String::from_utf8(buf).unwrap()
    }

    #[test]
    fn test_empty_graph() {
        let graph = LineageGraph::new();
        let output = render_to_string(&graph);
        let parsed: serde_json::Value = serde_json::from_str(&output).unwrap();
        assert_eq!(parsed["nodes"].as_array().unwrap().len(), 0);
        assert_eq!(parsed["edges"].as_array().unwrap().len(), 0);
    }

    #[test]
    fn test_edge_emits_via_hops_when_set() {
        let mut graph = LineageGraph::new();
        let a = graph.add_node(make_node("source.a", "a", NodeType::Source));
        let b = graph.add_node(make_node("model.b", "b", NodeType::Model));
        let eid = graph.add_edge(
            a,
            b,
            EdgeData {
                edge_type: EdgeType::Source,
            },
        );
        let mut via = HashMap::new();
        via.insert(eid, 3usize);

        let out = render_with_via(&graph, &via);
        let parsed: serde_json::Value = serde_json::from_str(&out).unwrap();
        let edges = parsed["edges"].as_array().unwrap();
        assert_eq!(edges[0]["via_hops"], 3);

        // And the default rendering omits via_hops entirely.
        let out = render_to_string(&graph);
        let parsed: serde_json::Value = serde_json::from_str(&out).unwrap();
        assert!(parsed["edges"][0].get("via_hops").is_none());
    }

    #[test]
    fn test_single_node() {
        let mut graph = LineageGraph::new();
        graph.add_node(make_node("model.orders", "orders", NodeType::Model));
        let output = render_to_string(&graph);
        let parsed: serde_json::Value = serde_json::from_str(&output).unwrap();
        let nodes = parsed["nodes"].as_array().unwrap();
        assert_eq!(nodes.len(), 1);
        assert_eq!(nodes[0]["unique_id"], "model.orders");
        assert_eq!(nodes[0]["label"], "orders");
        assert_eq!(nodes[0]["node_type"], "model");
        assert!(nodes[0].get("file_path").is_none());
        assert!(nodes[0].get("description").is_none());
    }

    #[test]
    fn test_node_with_file_path_and_description() {
        let mut graph = LineageGraph::new();
        graph.add_node(NodeData {
            unique_id: "model.orders".into(),
            label: "orders".into(),
            node_type: NodeType::Model,
            file_path: Some(PathBuf::from("models/orders.sql")),
            description: Some("Orders mart model".into()),
            materialization: None,
            tags: vec![],
            columns: vec![],
        });
        let output = render_to_string(&graph);
        let parsed: serde_json::Value = serde_json::from_str(&output).unwrap();
        let nodes = parsed["nodes"].as_array().unwrap();
        assert_eq!(nodes[0]["file_path"], "models/orders.sql");
        assert_eq!(nodes[0]["description"], "Orders mart model");
    }

    #[test]
    fn test_edges() {
        let mut graph = LineageGraph::new();
        let a = graph.add_node(make_node(
            "source.raw.orders",
            "raw.orders",
            NodeType::Source,
        ));
        let b = graph.add_node(make_node("model.stg_orders", "stg_orders", NodeType::Model));
        graph.add_edge(
            a,
            b,
            EdgeData {
                edge_type: EdgeType::Source,
            },
        );

        let output = render_to_string(&graph);
        let parsed: serde_json::Value = serde_json::from_str(&output).unwrap();
        let edges = parsed["edges"].as_array().unwrap();
        assert_eq!(edges.len(), 1);
        assert_eq!(edges[0]["source"], "source.raw.orders");
        assert_eq!(edges[0]["target"], "model.stg_orders");
        assert_eq!(edges[0]["edge_type"], "source");
    }

    #[test]
    fn test_all_edge_types() {
        assert_eq!(edge_type_label(EdgeType::Ref), "ref");
        assert_eq!(edge_type_label(EdgeType::Source), "source");
        assert_eq!(edge_type_label(EdgeType::Test), "test");
        assert_eq!(edge_type_label(EdgeType::Exposure), "exposure");
    }

    #[test]
    fn test_all_node_types() {
        let mut graph = LineageGraph::new();
        let types = [
            ("model.a", NodeType::Model, "model"),
            ("source.a.b", NodeType::Source, "source"),
            ("seed.a", NodeType::Seed, "seed"),
            ("snapshot.a", NodeType::Snapshot, "snapshot"),
            ("test.a", NodeType::Test, "test"),
            ("exposure.a", NodeType::Exposure, "exposure"),
            ("model.unknown", NodeType::Phantom, "phantom"),
        ];
        for (id, nt, _) in &types {
            graph.add_node(make_node(id, "a", *nt));
        }
        let output = render_to_string(&graph);
        let parsed: serde_json::Value = serde_json::from_str(&output).unwrap();
        let nodes = parsed["nodes"].as_array().unwrap();
        for (i, (_, _, expected_type)) in types.iter().enumerate() {
            assert_eq!(nodes[i]["node_type"], *expected_type);
        }
    }

    #[test]
    fn test_valid_json() {
        let mut graph = LineageGraph::new();
        let a = graph.add_node(make_node("model.a", "a", NodeType::Model));
        let b = graph.add_node(make_node("model.b", "b", NodeType::Model));
        graph.add_edge(
            a,
            b,
            EdgeData {
                edge_type: EdgeType::Ref,
            },
        );
        let output = render_to_string(&graph);
        // Should parse as valid JSON
        let _: serde_json::Value = serde_json::from_str(&output).unwrap();
    }

    #[test]
    fn test_node_with_materialization_tags_columns() {
        let mut graph = LineageGraph::new();
        graph.add_node(NodeData {
            unique_id: "model.orders".into(),
            label: "orders".into(),
            node_type: NodeType::Model,
            file_path: None,
            description: None,
            materialization: Some("table".into()),
            tags: vec!["daily".into(), "core".into()],
            columns: vec!["order_id".into(), "customer_id".into()],
        });
        let output = render_to_string(&graph);
        let parsed: serde_json::Value = serde_json::from_str(&output).unwrap();
        let node = &parsed["nodes"][0];
        assert_eq!(node["materialization"], "table");
        assert_eq!(node["tags"][0], "daily");
        assert_eq!(node["tags"][1], "core");
        assert_eq!(node["columns"][0], "order_id");
        assert_eq!(node["columns"][1], "customer_id");
    }
}
