use crate::ast::{ASTNode, BinaryOperator, LiteralValue, Span, UnaryOperator};

#[derive(Debug, Clone, PartialEq)]
pub struct StepTree {
    pub root: StepNode,
    pub features: StepTreeFeatures,
}

#[derive(Debug, Clone, PartialEq, Default)]
pub struct StepTreeFeatures {
    pub has_if: bool,
    pub has_loop: bool,
    pub has_break: bool,
    pub has_continue: bool,
    pub has_return: bool,
    pub max_if_depth: u32,
    pub max_loop_depth: u32,
}

#[derive(Debug, Clone, PartialEq)]
pub enum StepNode {
    Block(Vec<StepNode>),
    If {
        cond: AstSummary,
        then_branch: Box<StepNode>,
        else_branch: Option<Box<StepNode>>,
        span: Span,
    },
    Loop {
        cond: AstSummary,
        body: Box<StepNode>,
        span: Span,
    },
    Stmt { kind: StepStmtKind, span: Span },
}

#[derive(Debug, Clone, PartialEq)]
pub enum StepStmtKind {
    LocalDecl { vars: Vec<String> },
    Assign,
    Print,
    Return,
    Break,
    Continue,
    Other(&'static str),
}

#[derive(Debug, Clone, PartialEq)]
pub enum AstSummary {
    Variable(String),
    Literal(LiteralValue),
    Unary {
        op: UnaryOperator,
        expr: Box<AstSummary>,
    },
    Binary {
        op: BinaryOperator,
        lhs: Box<AstSummary>,
        rhs: Box<AstSummary>,
    },
    Other(&'static str),
}

impl StepTree {
    pub fn to_compact_string(&self) -> String {
        let mut out = String::new();
        self.root.write_compact(&mut out, 0);
        out
    }
}

impl StepNode {
    fn write_compact(&self, out: &mut String, indent: usize) {
        let pad = "  ".repeat(indent);
        match self {
            StepNode::Block(nodes) => {
                out.push_str(&format!("{pad}Block(len={})\n", nodes.len()));
                for n in nodes {
                    n.write_compact(out, indent + 1);
                }
            }
            StepNode::If {
                cond,
                then_branch,
                else_branch,
                ..
            } => {
                out.push_str(&format!(
                    "{pad}If(cond={})\n",
                    cond.to_compact_string()
                ));
                out.push_str(&format!("{pad}  then:\n"));
                then_branch.write_compact(out, indent + 2);
                if let Some(else_branch) = else_branch {
                    out.push_str(&format!("{pad}  else:\n"));
                    else_branch.write_compact(out, indent + 2);
                }
            }
            StepNode::Loop { cond, body, .. } => {
                out.push_str(&format!(
                    "{pad}Loop(cond={})\n",
                    cond.to_compact_string()
                ));
                body.write_compact(out, indent + 1);
            }
            StepNode::Stmt { kind, .. } => {
                out.push_str(&format!("{pad}Stmt({})\n", kind.to_compact_string()));
            }
        }
    }
}

impl StepStmtKind {
    fn to_compact_string(&self) -> String {
        match self {
            StepStmtKind::LocalDecl { vars } => format!("local({})", vars.join(",")),
            StepStmtKind::Assign => "assign".to_string(),
            StepStmtKind::Print => "print".to_string(),
            StepStmtKind::Return => "return".to_string(),
            StepStmtKind::Break => "break".to_string(),
            StepStmtKind::Continue => "continue".to_string(),
            StepStmtKind::Other(name) => format!("other:{name}"),
        }
    }
}

impl AstSummary {
    fn to_compact_string(&self) -> String {
        match self {
            AstSummary::Variable(name) => format!("var:{name}"),
            AstSummary::Literal(lit) => format!("lit:{lit:?}"),
            AstSummary::Unary { op, expr } => format!("({op:?} {})", expr.to_compact_string()),
            AstSummary::Binary { op, lhs, rhs } => format!(
                "({} {} {})",
                lhs.to_compact_string(),
                op,
                rhs.to_compact_string()
            ),
            AstSummary::Other(k) => format!("other:{k}"),
        }
    }
}

pub struct StepTreeBuilderBox;

impl StepTreeBuilderBox {
    pub fn build_from_ast(ast: &ASTNode) -> StepTree {
        match ast {
            ASTNode::Program { statements, .. } => Self::build_from_block(statements),
            ASTNode::ScopeBox { body, .. } => Self::build_from_block(body),
            _ => {
                let (node, features) = Self::build_node(ast, 0, 0);
                StepTree {
                    root: node,
                    features,
                }
            }
        }
    }

    pub fn build_from_block(stmts: &[ASTNode]) -> StepTree {
        let mut nodes = Vec::with_capacity(stmts.len());
        let mut features = StepTreeFeatures::default();
        for stmt in stmts {
            let (node, node_features) = Self::build_node(stmt, 0, 0);
            nodes.push(node);
            features = merge_features(features, node_features);
        }
        StepTree {
            root: StepNode::Block(nodes),
            features,
        }
    }

    fn build_node(ast: &ASTNode, if_depth: u32, loop_depth: u32) -> (StepNode, StepTreeFeatures) {
        match ast {
            ASTNode::If {
                condition,
                then_body,
                else_body,
                span,
            } => {
                let cond = summarize_ast(condition);
                let (then_node, then_features) =
                    Self::build_block_node(then_body, if_depth + 1, loop_depth);
                let (else_node, else_features) = match else_body {
                    Some(else_body) => {
                        let (node, f) =
                            Self::build_block_node(else_body, if_depth + 1, loop_depth);
                        (Some(Box::new(node)), f)
                    }
                    None => (None, StepTreeFeatures::default()),
                };
                let mut features = StepTreeFeatures {
                    has_if: true,
                    max_if_depth: (if_depth + 1).max(then_features.max_if_depth),
                    ..StepTreeFeatures::default()
                };
                features = merge_features(features, then_features);
                features = merge_features(features, else_features);

                (
                    StepNode::If {
                        cond,
                        then_branch: Box::new(then_node),
                        else_branch: else_node,
                        span: span.clone(),
                    },
                    features,
                )
            }
            ASTNode::Loop {
                condition, body, span, ..
            } => {
                let cond = summarize_ast(condition);
                let (body_node, body_features) =
                    Self::build_block_node(body, if_depth, loop_depth + 1);
                let mut features = StepTreeFeatures {
                    has_loop: true,
                    max_loop_depth: (loop_depth + 1).max(body_features.max_loop_depth),
                    ..StepTreeFeatures::default()
                };
                features = merge_features(features, body_features);
                (
                    StepNode::Loop {
                        cond,
                        body: Box::new(body_node),
                        span: span.clone(),
                    },
                    features,
                )
            }
            ASTNode::ScopeBox { body, span } => {
                let (node, features) = Self::build_block_node(body, if_depth, loop_depth);
                (node.with_span(span.clone()), features)
            }
            ASTNode::Return { span, .. } => (
                StepNode::Stmt {
                    kind: StepStmtKind::Return,
                    span: span.clone(),
                },
                StepTreeFeatures {
                    has_return: true,
                    ..StepTreeFeatures::default()
                },
            ),
            ASTNode::Break { span } => (
                StepNode::Stmt {
                    kind: StepStmtKind::Break,
                    span: span.clone(),
                },
                StepTreeFeatures {
                    has_break: true,
                    ..StepTreeFeatures::default()
                },
            ),
            ASTNode::Continue { span } => (
                StepNode::Stmt {
                    kind: StepStmtKind::Continue,
                    span: span.clone(),
                },
                StepTreeFeatures {
                    has_continue: true,
                    ..StepTreeFeatures::default()
                },
            ),
            ASTNode::Local {
                variables, span, ..
            } => (
                StepNode::Stmt {
                    kind: StepStmtKind::LocalDecl {
                        vars: variables.clone(),
                    },
                    span: span.clone(),
                },
                StepTreeFeatures::default(),
            ),
            ASTNode::Assignment { span, .. } => (
                StepNode::Stmt {
                    kind: StepStmtKind::Assign,
                    span: span.clone(),
                },
                StepTreeFeatures::default(),
            ),
            ASTNode::Print { span, .. } => (
                StepNode::Stmt {
                    kind: StepStmtKind::Print,
                    span: span.clone(),
                },
                StepTreeFeatures::default(),
            ),
            other => (
                StepNode::Stmt {
                    kind: StepStmtKind::Other(ast_kind_name(other)),
                    span: other.span(),
                },
                StepTreeFeatures::default(),
            ),
        }
    }

    fn build_block_node(
        stmts: &[ASTNode],
        if_depth: u32,
        loop_depth: u32,
    ) -> (StepNode, StepTreeFeatures) {
        let mut nodes = Vec::with_capacity(stmts.len());
        let mut features = StepTreeFeatures::default();
        for stmt in stmts {
            let (node, node_features) = Self::build_node(stmt, if_depth, loop_depth);
            nodes.push(node);
            features = merge_features(features, node_features);
        }
        (StepNode::Block(nodes), features)
    }
}

fn merge_features(mut a: StepTreeFeatures, b: StepTreeFeatures) -> StepTreeFeatures {
    a.has_if |= b.has_if;
    a.has_loop |= b.has_loop;
    a.has_break |= b.has_break;
    a.has_continue |= b.has_continue;
    a.has_return |= b.has_return;
    a.max_if_depth = a.max_if_depth.max(b.max_if_depth);
    a.max_loop_depth = a.max_loop_depth.max(b.max_loop_depth);
    a
}

fn summarize_ast(ast: &ASTNode) -> AstSummary {
    match ast {
        ASTNode::Variable { name, .. } => AstSummary::Variable(name.clone()),
        ASTNode::Literal { value, .. } => AstSummary::Literal(value.clone()),
        ASTNode::UnaryOp {
            operator, operand, ..
        } => AstSummary::Unary {
            op: operator.clone(),
            expr: Box::new(summarize_ast(operand)),
        },
        ASTNode::BinaryOp {
            operator,
            left,
            right,
            ..
        } => AstSummary::Binary {
            op: operator.clone(),
            lhs: Box::new(summarize_ast(left)),
            rhs: Box::new(summarize_ast(right)),
        },
        other => AstSummary::Other(ast_kind_name(other)),
    }
}

fn ast_kind_name(ast: &ASTNode) -> &'static str {
    match ast {
        ASTNode::Program { .. } => "Program",
        ASTNode::Assignment { .. } => "Assignment",
        ASTNode::Print { .. } => "Print",
        ASTNode::If { .. } => "If",
        ASTNode::Loop { .. } => "Loop",
        ASTNode::While { .. } => "While",
        ASTNode::ForRange { .. } => "ForRange",
        ASTNode::Return { .. } => "Return",
        ASTNode::Break { .. } => "Break",
        ASTNode::Continue { .. } => "Continue",
        ASTNode::UsingStatement { .. } => "UsingStatement",
        ASTNode::ImportStatement { .. } => "ImportStatement",
        ASTNode::Nowait { .. } => "Nowait",
        ASTNode::AwaitExpression { .. } => "AwaitExpression",
        ASTNode::QMarkPropagate { .. } => "QMarkPropagate",
        ASTNode::MatchExpr { .. } => "MatchExpr",
        ASTNode::ArrayLiteral { .. } => "ArrayLiteral",
        ASTNode::MapLiteral { .. } => "MapLiteral",
        ASTNode::Lambda { .. } => "Lambda",
        ASTNode::Arrow { .. } => "Arrow",
        ASTNode::TryCatch { .. } => "TryCatch",
        ASTNode::Throw { .. } => "Throw",
        ASTNode::BoxDeclaration { .. } => "BoxDeclaration",
        ASTNode::FunctionDeclaration { .. } => "FunctionDeclaration",
        ASTNode::GlobalVar { .. } => "GlobalVar",
        ASTNode::Literal { .. } => "Literal",
        ASTNode::Variable { .. } => "Variable",
        ASTNode::UnaryOp { .. } => "UnaryOp",
        ASTNode::BinaryOp { .. } => "BinaryOp",
        ASTNode::GroupedAssignmentExpr { .. } => "GroupedAssignmentExpr",
        ASTNode::MethodCall { .. } => "MethodCall",
        ASTNode::Call { .. } => "Call",
        ASTNode::FunctionCall { .. } => "FunctionCall",
        ASTNode::FieldAccess { .. } => "FieldAccess",
        ASTNode::Index { .. } => "Index",
        ASTNode::New { .. } => "New",
        ASTNode::This { .. } => "This",
        ASTNode::Me { .. } => "Me",
        ASTNode::FromCall { .. } => "FromCall",
        ASTNode::ThisField { .. } => "ThisField",
        ASTNode::MeField { .. } => "MeField",
        ASTNode::Local { .. } => "Local",
        ASTNode::ScopeBox { .. } => "ScopeBox",
        ASTNode::Outbox { .. } => "Outbox",
    }
}

impl StepNode {
    fn with_span(self, span: Span) -> StepNode {
        match self {
            StepNode::Block(nodes) => StepNode::Block(nodes),
            StepNode::If {
                cond,
                then_branch,
                else_branch,
                ..
            } => StepNode::If {
                cond,
                then_branch,
                else_branch,
                span,
            },
            StepNode::Loop { cond, body, .. } => StepNode::Loop { cond, body, span },
            StepNode::Stmt { kind, .. } => StepNode::Stmt { kind, span },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::{Span, ASTNode, BinaryOperator, LiteralValue};

    fn str_lit(s: &str) -> ASTNode {
        ASTNode::Literal {
            value: LiteralValue::String(s.to_string()),
            span: Span::unknown(),
        }
    }

    fn eq(a: ASTNode, b: ASTNode) -> ASTNode {
        ASTNode::BinaryOp {
            operator: BinaryOperator::Equal,
            left: Box::new(a),
            right: Box::new(b),
            span: Span::unknown(),
        }
    }

    fn assign_x(num: i64) -> ASTNode {
        ASTNode::Assignment {
            target: Box::new(ASTNode::Variable {
                name: "x".to_string(),
                span: Span::unknown(),
            }),
            value: Box::new(ASTNode::Literal {
                value: LiteralValue::Integer(num),
                span: Span::unknown(),
            }),
            span: Span::unknown(),
        }
    }

    #[test]
    fn build_step_tree_if_only_nested_if_is_structural() {
        // Equivalent shape to Phase103 "if-only merge" fixture:
        //
        // local x = 0
        // if "x" == "x" { if "y" == "z" { x=1 } else { x=2 } } else { x=3 }
        // print(x)
        let ast = vec![
            ASTNode::Local {
                variables: vec!["x".to_string()],
                initial_values: vec![Some(Box::new(ASTNode::Literal {
                    value: LiteralValue::Integer(0),
                    span: Span::unknown(),
                }))],
                span: Span::unknown(),
            },
            ASTNode::If {
                condition: Box::new(eq(str_lit("x"), str_lit("x"))),
                then_body: vec![ASTNode::If {
                    condition: Box::new(eq(str_lit("y"), str_lit("z"))),
                    then_body: vec![assign_x(1)],
                    else_body: Some(vec![assign_x(2)]),
                    span: Span::unknown(),
                }],
                else_body: Some(vec![assign_x(3)]),
                span: Span::unknown(),
            },
            ASTNode::Print {
                expression: Box::new(ASTNode::Variable {
                    name: "x".to_string(),
                    span: Span::unknown(),
                }),
                span: Span::unknown(),
            },
        ];

        let tree = StepTreeBuilderBox::build_from_block(&ast);
        assert!(tree.features.has_if);
        assert!(!tree.features.has_loop);
        assert_eq!(tree.features.max_if_depth, 2);

        match tree.root {
            StepNode::Block(nodes) => {
                assert_eq!(nodes.len(), 3);
                match &nodes[1] {
                    StepNode::If { then_branch, .. } => match &**then_branch {
                        StepNode::Block(inner_nodes) => match &inner_nodes[0] {
                            StepNode::If { .. } => {}
                            other => panic!("expected nested If, got {other:?}"),
                        },
                        other => panic!("expected Block in then_branch, got {other:?}"),
                    },
                    other => panic!("expected If at index 1, got {other:?}"),
                }
            }
            other => panic!("expected root Block, got {other:?}"),
        }
    }
}