use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use flowscope_core::{AnalyzeResult, EdgeType, NodeType};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MermaidView {
    All,
    Script,
    Table,
    Column,
    Hybrid,
}

pub fn export_mermaid(result: &AnalyzeResult, view: MermaidView) -> String {
    match view {
        MermaidView::All => generate_all_views(result),
        MermaidView::Script => generate_script_view(result),
        MermaidView::Table => generate_table_view(result),
        MermaidView::Column => generate_column_view(result),
        MermaidView::Hybrid => generate_hybrid_view(result),
    }
}

fn generate_all_views(result: &AnalyzeResult) -> String {
    let sections = vec![
        "# Lineage Diagrams".to_string(),
        String::new(),
        "## Script View".to_string(),
        "```mermaid".to_string(),
        generate_script_view(result),
        "```".to_string(),
        String::new(),
        "## Hybrid View (Scripts + Tables)".to_string(),
        "```mermaid".to_string(),
        generate_hybrid_view(result),
        "```".to_string(),
        String::new(),
        "## Table View".to_string(),
        "```mermaid".to_string(),
        generate_table_view(result),
        "```".to_string(),
        String::new(),
        "## Column View".to_string(),
        "```mermaid".to_string(),
        generate_column_view(result),
        "```".to_string(),
    ];

    sections.join("\n")
}

fn sanitize_id(id: &str) -> String {
    id.chars()
        .map(|c| {
            if c.is_alphanumeric() || c == '_' {
                c
            } else {
                '_'
            }
        })
        .collect()
}

fn escape_label(label: &str) -> String {
    label.replace('"', "\\\"").replace('\n', " ")
}

#[derive(Debug)]
struct ScriptInfo {
    source_name: String,
    tables_read: HashSet<Arc<str>>,
    tables_written: HashSet<Arc<str>>,
}

fn extract_script_info(result: &AnalyzeResult) -> Vec<ScriptInfo> {
    let mut script_map: HashMap<String, ScriptInfo> = HashMap::new();

    for stmt in &result.statements {
        let source_name = stmt
            .source_name
            .clone()
            .unwrap_or_else(|| "default".to_string());

        let entry = script_map
            .entry(source_name.clone())
            .or_insert_with(|| ScriptInfo {
                source_name: source_name.clone(),
                tables_read: HashSet::new(),
                tables_written: HashSet::new(),
            });

        let stmt_edges: Vec<_> = result.edges_in_statement(stmt.statement_index).collect();

        for node in result.nodes_in_statement(stmt.statement_index) {
            if matches!(node.node_type, NodeType::Table | NodeType::View) {
                let is_written = stmt_edges
                    .iter()
                    .any(|edge| edge.to == node.id && edge.edge_type == EdgeType::DataFlow);
                let is_read = stmt_edges
                    .iter()
                    .any(|edge| edge.from == node.id && edge.edge_type == EdgeType::DataFlow);

                let table_name = node
                    .qualified_name
                    .clone()
                    .unwrap_or_else(|| node.label.clone());

                if is_written {
                    entry.tables_written.insert(table_name.clone());
                }
                if is_read || !is_written {
                    entry.tables_read.insert(table_name);
                }
            }
        }
    }

    script_map.into_values().collect()
}

fn generate_script_view(result: &AnalyzeResult) -> String {
    let scripts = extract_script_info(result);
    let mut lines = vec!["flowchart LR".to_string()];

    for script in &scripts {
        let id = sanitize_id(&script.source_name);
        let label = escape_label(&script.source_name);
        lines.push(format!("    {id}[\"{label}\"]"));
    }

    for producer in &scripts {
        for consumer in &scripts {
            if producer.source_name == consumer.source_name {
                continue;
            }

            let shared_tables: Vec<_> = producer
                .tables_written
                .iter()
                .filter(|table| consumer.tables_read.contains(*table))
                .collect();

            if !shared_tables.is_empty() {
                let producer_id = sanitize_id(&producer.source_name);
                let consumer_id = sanitize_id(&consumer.source_name);
                let label = if shared_tables.len() > 3 {
                    let first_three: Vec<_> = shared_tables.iter().take(3).collect();
                    format!(
                        "{}...",
                        first_three
                            .iter()
                            .map(|value| value.as_ref())
                            .collect::<Vec<_>>()
                            .join(", ")
                    )
                } else {
                    shared_tables
                        .iter()
                        .map(|value| value.as_ref())
                        .collect::<Vec<_>>()
                        .join(", ")
                };
                lines.push(format!(
                    "    {producer_id} -->|\"{}\"| {consumer_id}",
                    escape_label(&label)
                ));
            }
        }
    }

    lines.join("\n")
}

fn generate_table_view(result: &AnalyzeResult) -> String {
    let mut lines = vec!["flowchart LR".to_string()];
    let mut table_ids: HashMap<String, String> = HashMap::new();
    let mut edges = HashSet::new();

    let table_nodes: Vec<_> = result
        .nodes
        .iter()
        .filter(|node| node.node_type.is_table_like())
        .collect();

    for node in &table_nodes {
        let key = node
            .qualified_name
            .as_deref()
            .unwrap_or(&node.label)
            .to_string();
        if !table_ids.contains_key(&key) {
            let id = sanitize_id(&key);
            table_ids.insert(key.clone(), id.clone());
            let escaped_label = escape_label(&node.label);
            let shape = match node.node_type {
                NodeType::Cte => format!("([\"{escaped_label}\"])"),
                NodeType::View => format!("[/\"{escaped_label}\"/]"),
                _ => format!("[\"{escaped_label}\"]"),
            };
            lines.push(format!("    {id}{shape}"));
        }
    }

    for edge in &result.edges {
        if edge.edge_type == EdgeType::DataFlow || edge.edge_type == EdgeType::Derivation {
            let source_node = table_nodes.iter().find(|node| node.id == edge.from);
            let target_node = table_nodes.iter().find(|node| node.id == edge.to);

            if let (Some(source), Some(target)) = (source_node, target_node) {
                let source_key = source
                    .qualified_name
                    .as_deref()
                    .unwrap_or(&source.label)
                    .to_string();
                let target_key = target
                    .qualified_name
                    .as_deref()
                    .unwrap_or(&target.label)
                    .to_string();
                let edge_key = format!("{source_key}->{target_key}");

                if source_key != target_key && edges.insert(edge_key) {
                    let source_id = table_ids.get(&source_key).cloned().unwrap_or_else(|| {
                        let id = sanitize_id(&source_key);
                        table_ids.insert(source_key.clone(), id.clone());
                        id
                    });
                    let target_id = table_ids.get(&target_key).cloned().unwrap_or_else(|| {
                        let id = sanitize_id(&target_key);
                        table_ids.insert(target_key.clone(), id.clone());
                        id
                    });
                    lines.push(format!("    {source_id} --> {target_id}"));
                }
            }
        }
    }

    lines.join("\n")
}

#[derive(Debug)]
struct ColumnMapping {
    source_table: String,
    source_column: String,
    target_table: String,
    target_column: String,
    edge_type: EdgeType,
}

fn extract_column_mappings(result: &AnalyzeResult) -> Vec<ColumnMapping> {
    let mut mappings = Vec::new();

    let table_nodes: Vec<_> = result
        .nodes
        .iter()
        .filter(|node| node.node_type.is_table_like())
        .collect();
    let column_nodes: Vec<_> = result
        .nodes
        .iter()
        .filter(|node| node.node_type == NodeType::Column)
        .collect();

    let mut column_to_table: HashMap<&str, &str> = HashMap::new();
    for edge in &result.edges {
        if edge.edge_type == EdgeType::Ownership {
            if let Some(table_node) = table_nodes.iter().find(|node| node.id == edge.from) {
                let table_name = table_node
                    .qualified_name
                    .as_deref()
                    .unwrap_or(&table_node.label);
                column_to_table.insert(edge.to.as_ref(), table_name);
            }
        }
    }

    for edge in &result.edges {
        if edge.edge_type == EdgeType::Derivation || edge.edge_type == EdgeType::DataFlow {
            let source_col = column_nodes.iter().find(|col| col.id == edge.from);
            let target_col = column_nodes.iter().find(|col| col.id == edge.to);

            if let (Some(source), Some(target)) = (source_col, target_col) {
                let source_table = column_to_table
                    .get(edge.from.as_ref())
                    .copied()
                    .unwrap_or("Output");
                let target_table = column_to_table
                    .get(edge.to.as_ref())
                    .copied()
                    .unwrap_or("Output");

                mappings.push(ColumnMapping {
                    source_table: source_table.to_string(),
                    source_column: source.label.to_string(),
                    target_table: target_table.to_string(),
                    target_column: target.label.to_string(),
                    edge_type: edge.edge_type,
                });
            }
        }
    }

    mappings
}

fn generate_column_view(result: &AnalyzeResult) -> String {
    let mut lines = vec!["flowchart LR".to_string()];
    let mappings = extract_column_mappings(result);
    let mut nodes = HashSet::new();
    let mut edges = HashSet::new();

    for mapping in mappings {
        let source_id = sanitize_id(&format!(
            "{}_{}",
            mapping.source_table, mapping.source_column
        ));
        let target_id = sanitize_id(&format!(
            "{}_{}",
            mapping.target_table, mapping.target_column
        ));
        let source_label = format!("{}.{}", mapping.source_table, mapping.source_column);
        let target_label = format!("{}.{}", mapping.target_table, mapping.target_column);

        if nodes.insert(source_id.clone()) {
            lines.push(format!(
                "    {source_id}[\"{}\"]",
                escape_label(&source_label)
            ));
        }
        if nodes.insert(target_id.clone()) {
            lines.push(format!(
                "    {target_id}[\"{}\"]",
                escape_label(&target_label)
            ));
        }

        let edge_key = format!("{source_id}->{target_id}");
        if edges.insert(edge_key) {
            let edge_label = match mapping.edge_type {
                EdgeType::Derivation => "derived",
                _ => "flows",
            };
            lines.push(format!("    {source_id} -->|{edge_label}| {target_id}"));
        }
    }

    lines.join("\n")
}

fn generate_hybrid_view(result: &AnalyzeResult) -> String {
    let mut lines = vec!["flowchart LR".to_string()];
    let scripts = extract_script_info(result);

    let mut script_ids = HashMap::new();
    for script in &scripts {
        let id = sanitize_id(&format!("script_{}", script.source_name));
        script_ids.insert(script.source_name.clone(), id.clone());
        lines.push(format!(
            "    {id}{{\"{}\"}}",
            escape_label(&script.source_name)
        ));
    }

    let mut table_ids = HashMap::new();
    for node in &result.nodes {
        if node.node_type.is_table_like() {
            let key = node
                .qualified_name
                .as_deref()
                .unwrap_or(&node.label)
                .to_string();
            if !table_ids.contains_key(&key) {
                let id = sanitize_id(&format!("table_{}", key));
                table_ids.insert(key.clone(), id.clone());
                lines.push(format!("    {id}[\"{}\"]", escape_label(&node.label)));
            }
        }
    }

    for script in scripts {
        let script_id = script_ids
            .get(&script.source_name)
            .cloned()
            .unwrap_or_else(|| {
                let id = sanitize_id(&format!("script_{}", script.source_name));
                script_ids.insert(script.source_name.clone(), id.clone());
                id
            });

        for table in &script.tables_read {
            if let Some(table_id) = table_ids.get(table.as_ref()) {
                lines.push(format!("    {script_id} --> {table_id}"));
            }
        }
        for table in &script.tables_written {
            if let Some(table_id) = table_ids.get(table.as_ref()) {
                lines.push(format!("    {table_id} --> {script_id}"));
            }
        }
    }

    lines.join("\n")
}
