use petgraph::dot::{Config, Dot};
use petgraph::graph::{Graph, NodeIndex};

use crate::ast::types::{
    BlockStmt, CompileUnit, Expr, ExprValue, FuncDeclStmt, GlobalDeclStmt, Param, ReturnStmt,
    Statement, VarDeclStmt, VarDeclStmtValue,
};

pub type AstGraph = Graph<String, String>;

pub trait AstGraphExt {
    fn to_graph(&self) -> AstGraph;
    fn to_dot(&self) -> String {
        format!("{}", Dot::with_config(&self.to_graph(), &[Config::EdgeNoLabel]))
    }
}

impl AstGraphExt for CompileUnit {
    fn to_graph(&self) -> AstGraph {
        let mut builder = AstGraphBuilder::new();
        builder.add_compile_unit(self);
        builder.graph
    }
}

struct AstGraphBuilder {
    graph: AstGraph,
}

impl AstGraphBuilder {
    fn new() -> Self {
        Self { graph: Graph::new() }
    }

    fn node(&mut self, label: impl Into<String>) -> NodeIndex {
        self.graph.add_node(label.into())
    }

    fn child(&mut self, parent: NodeIndex, label: impl Into<String>) -> NodeIndex {
        let child = self.node(label);
        self.graph.add_edge(parent, child, String::new());
        child
    }

    fn add_compile_unit(&mut self, compile_unit: &CompileUnit) -> NodeIndex {
        let root = self.node(compile_unit.to_string());
        for decl in &compile_unit.global_decls {
            self.add_global_decl(root, decl);
        }
        root
    }

    fn add_global_decl(&mut self, parent: NodeIndex, decl: &GlobalDeclStmt) -> NodeIndex {
        match decl {
            GlobalDeclStmt::VarDecl(var_decl) => {
                let node = self.child(parent, decl.to_string());
                self.add_var_decl(node, var_decl);
                node
            }
            GlobalDeclStmt::FuncDecl(func_decl) => {
                let node = self.child(parent, decl.to_string());
                self.add_func_decl(node, func_decl);
                node
            }
        }
    }

    fn add_var_decl(&mut self, parent: NodeIndex, var_decl: &VarDeclStmt) -> NodeIndex {
        let node = self.child(parent, var_decl.to_string());
        for value in &var_decl.values {
            self.add_var_decl_value(node, value);
        }
        node
    }

    fn add_var_decl_value(&mut self, parent: NodeIndex, value: &VarDeclStmtValue) -> NodeIndex {
        self.child(parent, value.to_string())
    }

    fn add_func_decl(&mut self, parent: NodeIndex, func_decl: &FuncDeclStmt) -> NodeIndex {
        let node = self.child(parent, func_decl.to_string());
        let params = self.child(node, "Params");
        for param in &func_decl.params {
            self.add_param(params, param);
        }
        self.add_block_stmt(node, &func_decl.body);
        node
    }

    fn add_param(&mut self, parent: NodeIndex, param: &Param) -> NodeIndex {
        self.child(parent, param.to_string())
    }

    fn add_block_stmt(&mut self, parent: NodeIndex, block_stmt: &BlockStmt) -> NodeIndex {
        let node = self.child(parent, block_stmt.to_string());
        for stmt in &block_stmt.statements {
            self.add_statement(node, stmt);
        }
        node
    }

    fn add_statement(&mut self, parent: NodeIndex, stmt: &Statement) -> NodeIndex {
        match stmt {
            Statement::Return(return_stmt) => {
                let node = self.child(parent, stmt.to_string());
                self.add_return_stmt(node, return_stmt);
                node
            }
            Statement::Block(block_stmt) => self.add_block_stmt(parent, block_stmt),
            Statement::Expr(expr) => {
                let node = self.child(parent, stmt.to_string());
                self.add_expr(node, expr);
                node
            }
            Statement::VarDecl(var_decl) => {
                let node = self.child(parent, stmt.to_string());
                self.add_var_decl(node, var_decl);
                node
            }
        }
    }

    fn add_return_stmt(&mut self, parent: NodeIndex, return_stmt: &ReturnStmt) -> NodeIndex {
        match &return_stmt.value {
            Some(expr) => self.add_expr(parent, expr),
            None => self.child(parent, "Void"),
        }
    }

    fn add_expr(&mut self, parent: NodeIndex, expr: &Expr) -> NodeIndex {
        match &expr.value {
            ExprValue::IntLit(_) | ExprValue::Var(_) => self.child(parent, expr.value.to_string()),
            ExprValue::BinaryOp { lhs, op: _, rhs } => {
                let node = self.child(parent, expr.value.to_string());
                self.add_expr(node, lhs);
                self.add_expr(node, rhs);
                node
            }
            ExprValue::FuncCall(_, args) => {
                let node = self.child(parent, expr.value.to_string());
                let args_node = self.child(node, "Args");
                for arg in args {
                    self.add_expr(args_node, arg);
                }
                node
            }
            ExprValue::Assign { lvalue, rvalue } => {
                let node = self.child(parent, expr.value.to_string());
                self.add_expr(node, lvalue);
                self.add_expr(node, rvalue);
                node
            }
        }
    }
}
