use std::fmt::Display;

use crate::misc::indent;

use super::Expr;

#[derive(Debug, Clone)]
pub enum Stmt {
    Print {
        expr: Box<Expr>,
    },
    IfStmt {
        condition: Box<Expr>,
        then_branch: Box<Stmt>,
        else_branch: Option<Box<Stmt>>,
    },
    While {
        condition: Box<Expr>,
        body: Box<Stmt>,
    },
    VarDecl {
        name: String,
        initializer: Box<Expr>,
    },
    Block {
        statements: Vec<Stmt>,
    },
    ExprStmt {
        expr: Box<Expr>,
    },
    Break,
    Return {
        expr: Box<Expr>,
    },
}

impl Stmt {
    pub fn print_stmt(expr: Expr) -> Self {
        let expr = Box::new(expr);
        Stmt::Print { expr }
    }

    pub fn if_stmt(
        condition: impl Into<Box<Expr>>,
        then_branch: impl Into<Box<Stmt>>,
        else_branch: Option<impl Into<Box<Stmt>>>,
    ) -> Self {
        let condition = condition.into();
        let then_branch = then_branch.into();
        let else_branch = else_branch.map(|stmt| stmt.into());

        Stmt::IfStmt {
            condition,
            then_branch,
            else_branch,
        }
    }

    pub fn while_stmt(condition: impl Into<Box<Expr>>, body: impl Into<Box<Stmt>>) -> Self {
        let condition = condition.into();
        let body = body.into();

        Stmt::While { condition, body }
    }

    pub fn var_decl(name: impl Into<String>, initializer: impl Into<Box<Expr>>) -> Self {
        let name = name.into();
        let initializer = initializer.into();
        Stmt::VarDecl { name, initializer }
    }

    pub fn expr_stmt(expr: impl Into<Box<Expr>>) -> Self {
        let expr = expr.into();
        Stmt::ExprStmt { expr }
    }

    pub fn return_stmt(expr: impl Into<Box<Expr>>) -> Self {
        let expr = expr.into();
        Stmt::Return { expr }
    }
}

impl Display for Stmt {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Stmt::Print { expr } => write!(f, "print {expr};"),
            Stmt::IfStmt {
                condition,
                then_branch,
                else_branch,
            } => {
                writeln!(f, "if {condition}")?;
                match then_branch.as_ref() {
                    Stmt::Block { .. } => write!(f, "{then_branch}")?,
                    _ => write!(f, "{}", indent(format!("{then_branch}")))?,
                }
                if let Some(else_branch) = else_branch {
                    writeln!(f, "\nelse")?;
                    match else_branch.as_ref() {
                        Stmt::Block { .. } => write!(f, "{else_branch}")?,
                        _ => write!(f, "{}", indent(format!("{else_branch}")))?,
                    }
                }
                Ok(())
            }
            Stmt::While { condition, body } => {
                writeln!(f, "{condition}")?;
                match body.as_ref() {
                    Stmt::Block { .. } => write!(f, "{body}")?,
                    _ => write!(f, "{}", indent(format!("{body}")))?,
                }
                Ok(())
            }
            Stmt::VarDecl { name, initializer } => write!(f, "var {name} = {initializer};"),
            Stmt::Block { statements } => {
                writeln!(f, "{{")?;
                for statement in statements {
                    write!(f, "{}", indent(format!("{statement}")))?;
                }
                write!(f, "}}")
            }
            Stmt::ExprStmt { expr } => write!(f, "{expr};"),
            Stmt::Break => write!(f, "break;"),
            Stmt::Return { expr } => {
                write!(f, "return {expr};")
            }
        }
    }
}