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

use crate::ast::types::{
    ArrayDimension, BlockStmt, BreakStmt, CompileUnit, ContinueStmt, Expr, ExprValue, ForInit, ForStmt, FuncDeclStmt, GlobalDeclStmt, IfStmt, Param, ReturnStmt, Statement, VarDeclStmt, VarDeclStmtValue, WhileStmt
};

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 {
        let node = self.child(parent, value.to_string());
        for dimension in &value.dimensions {
            self.add_array_dimension(node, dimension);
        }
        if let Some(expr) = &value.value {
            self.add_expr(node, expr);
        }
        node
    }

    fn add_array_dimension(&mut self, parent: NodeIndex, dimension: &ArrayDimension) -> NodeIndex {
        let node = self.child(parent, dimension.to_string());
        if let Some(value) = &dimension.value {
            self.add_expr(node, value);
        }
        node
    }

    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 {
        let node = self.child(parent, param.to_string());
        for dimension in &param.dimensions {
            self.add_array_dimension(node, dimension);
        }
        node
    }

    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
            },
            Statement::If(if_stmt) => self.add_if_stmt(parent, if_stmt),
            Statement::While(while_stmt) => self.add_while_stmt(parent, while_stmt),
            Statement::For(for_stmt) => self.add_for_stmt(parent, for_stmt),
            Statement::Break(break_stmt) => self.add_break_stmt(parent, break_stmt),
            Statement::Continue(continue_stmt) => self.add_continue_stmt(parent, continue_stmt),
        }
    }

    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_if_stmt(&mut self, parent: NodeIndex, if_stmt: &IfStmt) -> NodeIndex {
        let node = self.child(parent, if_stmt.to_string());
        self.add_expr(node, &if_stmt.condition);
        self.add_block_stmt(node, &if_stmt.then_branch);
        if let Some(else_branch) = &if_stmt.else_branch {
            self.add_block_stmt(node, else_branch);
        }
        node
    }
    fn add_while_stmt(&mut self, parent: NodeIndex, while_stmt: &WhileStmt) -> NodeIndex {
        let node = self.child(parent, while_stmt.to_string());
        self.add_expr(node, &while_stmt.condition);
        self.add_block_stmt(node, &while_stmt.body);
        node
    }
    fn add_for_stmt(&mut self, parent: NodeIndex, for_stmt: &ForStmt) -> NodeIndex {
        let node = self.child(parent, for_stmt.to_string());
        match &for_stmt.init {
            Some(ForInit::Expr(expr)) => {
                let init = self.child(node, "Init");
                self.add_expr(init, expr);
            }
            Some(ForInit::VarDecl(var_decl)) => {
                let init = self.child(node, "Init");
                self.add_var_decl(init, var_decl);
            }
            None => {
                self.child(node, "NoInit");
            }
        }
        match &for_stmt.condition {
            Some(expr) => {
                let condition = self.child(node, "Condition");
                self.add_expr(condition, expr);
            }
            None => {
                self.child(node, "NoCondition");
            }
        }
        match &for_stmt.update {
            Some(expr) => {
                let update = self.child(node, "Update");
                self.add_expr(update, expr);
            }
            None => {
                self.child(node, "NoUpdate");
            }
        }
        self.add_block_stmt(node, &for_stmt.body);
        node
    }
    fn add_break_stmt(&mut self, parent: NodeIndex, break_stmt: &BreakStmt) -> NodeIndex {
        self.child(parent, break_stmt.to_string())
    }
    fn add_continue_stmt(&mut self, parent: NodeIndex, continue_stmt: &ContinueStmt) -> NodeIndex {
        self.child(parent, continue_stmt.to_string())
    }
    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::ArrayAccess { array, index } => {
                let node = self.child(parent, expr.value.to_string());
                self.add_expr(node, array);
                self.add_expr(node, index);
                node
            }
            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
            },
            ExprValue::CompoundAssign { lvalue, rvalue, .. } => {
                let node = self.child(parent, expr.value.to_string());
                self.add_expr(node, lvalue);
                self.add_expr(node, rvalue);
                node
            },
            ExprValue::UnaryOp { op: _, operand } => {
                let node = self.child(parent, expr.value.to_string());
                self.add_expr(node, operand);
                node
            }
            ExprValue::IncDec { operand, .. } => {
                let node = self.child(parent, expr.value.to_string());
                self.add_expr(node, operand);
                node
            }
        }
    }
}
