use snail_ast::*;

use super::break_rewrite::rewrite_breaks_in_block;

pub(super) struct Desugarer {
    block_counter: usize,
}

impl Desugarer {
    pub(super) fn new() -> Self {
        Self { block_counter: 0 }
    }

    fn next_block_name(&mut self) -> String {
        let name = format!("__snail_block_{}", self.block_counter);
        self.block_counter += 1;
        name
    }

    fn next_expr_name(&mut self) -> String {
        let name = format!("__snail_expr_{}", self.block_counter);
        self.block_counter += 1;
        name
    }

    /// Modify the tail (last unterminated expression) of a statement block to
    /// assign its value to `tmp_name` instead of discarding it.  For compound
    /// tails (if, try, with, block) the capture is propagated recursively into
    /// each branch so that every exit path writes to the same temp variable.
    fn capture_branch_tail(stmts: &mut Vec<Stmt>, tmp_name: &str, span: &SourceSpan) {
        // Classify the last statement without borrowing mutably yet.
        enum Action {
            CaptureSimple,
            CaptureCompound,
            NoCapture,
        }

        let action = match stmts.last() {
            Some(Stmt::Expr {
                value,
                semicolon_terminated: false,
                ..
            }) => match value {
                Expr::Block { .. }
                | Expr::If { .. }
                | Expr::Try { .. }
                | Expr::With { .. }
                | Expr::While { .. }
                | Expr::For { .. } => Action::CaptureCompound,
                _ => Action::CaptureSimple,
            },
            _ => Action::NoCapture,
        };

        match action {
            Action::CaptureSimple => {
                if let Some(Stmt::Expr { value, span: s, .. }) = stmts.pop() {
                    stmts.push(Stmt::Assign {
                        targets: vec![AssignTarget::Name {
                            name: tmp_name.to_string(),
                            span: span.clone(),
                        }],
                        value,
                        span: s,
                    });
                }
            }
            Action::CaptureCompound => {
                if let Some(Stmt::Expr {
                    value,
                    semicolon_terminated,
                    span: s,
                }) = stmts.pop()
                {
                    let new_value = match value {
                        Expr::Block {
                            stmts: mut inner,
                            span: bs,
                        } => {
                            Self::capture_branch_tail(&mut inner, tmp_name, span);
                            Expr::Block {
                                stmts: inner,
                                span: bs,
                            }
                        }
                        Expr::If {
                            cond,
                            mut body,
                            mut elifs,
                            mut else_body,
                            span: is,
                        } => {
                            Self::capture_branch_tail(&mut body, tmp_name, span);
                            for (_, eb) in elifs.iter_mut() {
                                Self::capture_branch_tail(eb, tmp_name, span);
                            }
                            if let Some(eb) = &mut else_body {
                                Self::capture_branch_tail(eb, tmp_name, span);
                            }
                            Expr::If {
                                cond,
                                body,
                                elifs,
                                else_body,
                                span: is,
                            }
                        }
                        Expr::Try {
                            mut body,
                            mut handlers,
                            mut else_body,
                            finally_body,
                            span: ts,
                        } => {
                            Self::capture_branch_tail(&mut body, tmp_name, span);
                            for h in handlers.iter_mut() {
                                Self::capture_branch_tail(&mut h.body, tmp_name, span);
                            }
                            if let Some(eb) = &mut else_body {
                                Self::capture_branch_tail(eb, tmp_name, span);
                            }
                            Expr::Try {
                                body,
                                handlers,
                                else_body,
                                finally_body,
                                span: ts,
                            }
                        }
                        Expr::With {
                            items,
                            mut body,
                            span: ws,
                        } => {
                            Self::capture_branch_tail(&mut body, tmp_name, span);
                            Expr::With {
                                items,
                                body,
                                span: ws,
                            }
                        }
                        Expr::For {
                            target,
                            iter,
                            mut body,
                            mut else_body,
                            span: fs,
                        } => {
                            rewrite_breaks_in_block(&mut body, tmp_name, span);
                            Self::capture_branch_tail(&mut body, tmp_name, span);
                            if let Some(eb) = &mut else_body {
                                Self::capture_branch_tail(eb, tmp_name, span);
                            }
                            Expr::For {
                                target,
                                iter,
                                body,
                                else_body,
                                span: fs,
                            }
                        }
                        Expr::While {
                            cond,
                            mut body,
                            mut else_body,
                            span: ws,
                        } => {
                            rewrite_breaks_in_block(&mut body, tmp_name, span);
                            Self::capture_branch_tail(&mut body, tmp_name, span);
                            if let Some(eb) = &mut else_body {
                                Self::capture_branch_tail(eb, tmp_name, span);
                            }
                            Expr::While {
                                cond,
                                body,
                                else_body,
                                span: ws,
                            }
                        }
                        other => other,
                    };
                    stmts.push(Stmt::Expr {
                        value: new_value,
                        semicolon_terminated,
                        span: s,
                    });
                }
            }
            Action::NoCapture => {}
        }
    }

    /// Replace a compound expression with an inline expansion: emit the
    /// compound's statements into the prelude and capture the tail value in a
    /// temp variable.  No new scope is created, so variables assigned inside
    /// the compound remain visible in the enclosing scope.
    fn inline_compound_expr(
        &mut self,
        expr: Expr,
        prelude: &mut Vec<Stmt>,
        span: &SourceSpan,
    ) -> Expr {
        let tmp_name = self.next_expr_name();

        // Push `tmp = None` as default value.
        prelude.push(Stmt::Assign {
            targets: vec![AssignTarget::Name {
                name: tmp_name.clone(),
                span: span.clone(),
            }],
            value: Expr::None { span: span.clone() },
            span: span.clone(),
        });

        match expr {
            Expr::Block { mut stmts, .. } => {
                Self::capture_branch_tail(&mut stmts, &tmp_name, span);
                prelude.extend(stmts);
            }
            Expr::If {
                cond,
                mut body,
                mut elifs,
                mut else_body,
                span: is,
            } => {
                Self::capture_branch_tail(&mut body, &tmp_name, span);
                for (_, eb) in elifs.iter_mut() {
                    Self::capture_branch_tail(eb, &tmp_name, span);
                }
                if let Some(eb) = &mut else_body {
                    Self::capture_branch_tail(eb, &tmp_name, span);
                }
                prelude.push(Stmt::Expr {
                    value: Expr::If {
                        cond,
                        body,
                        elifs,
                        else_body,
                        span: is,
                    },
                    semicolon_terminated: false,
                    span: span.clone(),
                });
            }
            Expr::Try {
                mut body,
                mut handlers,
                mut else_body,
                finally_body,
                span: ts,
            } => {
                Self::capture_branch_tail(&mut body, &tmp_name, span);
                for h in handlers.iter_mut() {
                    Self::capture_branch_tail(&mut h.body, &tmp_name, span);
                }
                if let Some(eb) = &mut else_body {
                    Self::capture_branch_tail(eb, &tmp_name, span);
                }
                prelude.push(Stmt::Expr {
                    value: Expr::Try {
                        body,
                        handlers,
                        else_body,
                        finally_body,
                        span: ts,
                    },
                    semicolon_terminated: false,
                    span: span.clone(),
                });
            }
            Expr::With {
                items,
                mut body,
                span: ws,
            } => {
                Self::capture_branch_tail(&mut body, &tmp_name, span);
                prelude.push(Stmt::Expr {
                    value: Expr::With {
                        items,
                        body,
                        span: ws,
                    },
                    semicolon_terminated: false,
                    span: span.clone(),
                });
            }
            Expr::For {
                target,
                iter,
                mut body,
                mut else_body,
                span: fs,
            } => {
                rewrite_breaks_in_block(&mut body, &tmp_name, span);
                Self::capture_branch_tail(&mut body, &tmp_name, span);
                if let Some(eb) = &mut else_body {
                    Self::capture_branch_tail(eb, &tmp_name, span);
                }
                prelude.push(Stmt::Expr {
                    value: Expr::For {
                        target,
                        iter,
                        body,
                        else_body,
                        span: fs,
                    },
                    semicolon_terminated: false,
                    span: span.clone(),
                });
            }
            Expr::While {
                cond,
                mut body,
                mut else_body,
                span: ws,
            } => {
                rewrite_breaks_in_block(&mut body, &tmp_name, span);
                Self::capture_branch_tail(&mut body, &tmp_name, span);
                if let Some(eb) = &mut else_body {
                    Self::capture_branch_tail(eb, &tmp_name, span);
                }
                prelude.push(Stmt::Expr {
                    value: Expr::While {
                        cond,
                        body,
                        else_body,
                        span: ws,
                    },
                    semicolon_terminated: false,
                    span: span.clone(),
                });
            }
            // Other compounds don't produce values; push as-is.
            other => {
                prelude.push(Stmt::Expr {
                    value: other,
                    semicolon_terminated: false,
                    span: span.clone(),
                });
            }
        }

        Expr::Name {
            name: tmp_name,
            span: span.clone(),
        }
    }

    pub(super) fn desugar_program(&mut self, program: &Program) -> Program {
        Program {
            stmts: self.desugar_block(&program.stmts),
            span: program.span.clone(),
        }
    }

    pub(super) fn desugar_block(&mut self, block: &[Stmt]) -> Vec<Stmt> {
        let mut out = Vec::new();
        for stmt in block {
            let mut prelude = Vec::new();
            let stmt = self.desugar_stmt(stmt, &mut prelude);
            out.extend(prelude);
            out.push(stmt);
        }
        out
    }

    fn desugar_stmt(&mut self, stmt: &Stmt, prelude: &mut Vec<Stmt>) -> Stmt {
        match stmt {
            Stmt::Return { value, span } => Stmt::Return {
                value: value
                    .as_ref()
                    .map(|value| self.desugar_expr(value, prelude)),
                span: span.clone(),
            },
            Stmt::Raise { value, from, span } => Stmt::Raise {
                value: value
                    .as_ref()
                    .map(|value| self.desugar_expr(value, prelude)),
                from: from.as_ref().map(|from| self.desugar_expr(from, prelude)),
                span: span.clone(),
            },
            Stmt::Assert {
                test,
                message,
                span,
            } => Stmt::Assert {
                test: self.desugar_expr(test, prelude),
                message: message
                    .as_ref()
                    .map(|message| self.desugar_expr(message, prelude)),
                span: span.clone(),
            },
            Stmt::Delete { targets, span } => Stmt::Delete {
                targets: targets
                    .iter()
                    .map(|target| self.desugar_assign_target(target, prelude))
                    .collect(),
                span: span.clone(),
            },
            Stmt::Break { value, span } => Stmt::Break {
                value: value
                    .as_ref()
                    .map(|value| self.desugar_expr(value, prelude)),
                span: span.clone(),
            },
            Stmt::Continue { span } => Stmt::Continue { span: span.clone() },
            Stmt::Pass { span } => Stmt::Pass { span: span.clone() },
            Stmt::Import { items, span } => Stmt::Import {
                items: items.clone(),
                span: span.clone(),
            },
            Stmt::ImportFrom {
                level,
                module,
                items,
                span,
            } => Stmt::ImportFrom {
                level: *level,
                module: module.clone(),
                items: items.clone(),
                span: span.clone(),
            },
            Stmt::Assign {
                targets,
                value,
                span,
            } => Stmt::Assign {
                targets: targets
                    .iter()
                    .map(|target| self.desugar_assign_target(target, prelude))
                    .collect(),
                value: self.desugar_expr(value, prelude),
                span: span.clone(),
            },
            Stmt::Expr {
                value,
                semicolon_terminated,
                span,
            } => Stmt::Expr {
                value: self.desugar_expr_no_hoist(value, prelude),
                semicolon_terminated: *semicolon_terminated,
                span: span.clone(),
            },
            Stmt::PatternAction {
                pattern,
                action,
                span,
            } => Stmt::PatternAction {
                pattern: pattern.as_ref().map(|p| self.desugar_expr(p, prelude)),
                action: action.as_ref().map(|a| self.desugar_block(a)),
                span: span.clone(),
            },
            Stmt::SegmentBreak { span } => Stmt::SegmentBreak { span: span.clone() },
        }
    }

    fn desugar_condition(&mut self, cond: &Condition, prelude: &mut Vec<Stmt>) -> Condition {
        match cond {
            Condition::Expr(expr) => Condition::Expr(Box::new(self.desugar_expr(expr, prelude))),
            Condition::Let {
                target,
                value,
                guard,
                span,
            } => Condition::Let {
                target: Box::new(self.desugar_assign_target(target, prelude)),
                value: Box::new(self.desugar_expr(value, prelude)),
                guard: guard
                    .as_ref()
                    .map(|guard| Box::new(self.desugar_expr(guard, prelude))),
                span: span.clone(),
            },
        }
    }

    fn desugar_with_item(&mut self, item: &WithItem, prelude: &mut Vec<Stmt>) -> WithItem {
        WithItem {
            context: self.desugar_expr(&item.context, prelude),
            target: item
                .target
                .as_ref()
                .map(|target| self.desugar_assign_target(target, prelude)),
            span: item.span.clone(),
        }
    }

    fn desugar_except_handler(
        &mut self,
        handler: &ExceptHandler,
        prelude: &mut Vec<Stmt>,
    ) -> ExceptHandler {
        ExceptHandler {
            type_name: handler
                .type_name
                .as_ref()
                .map(|expr| self.desugar_expr(expr, prelude)),
            name: handler.name.clone(),
            body: self.desugar_block(&handler.body),
            span: handler.span.clone(),
        }
    }

    fn desugar_params(&mut self, params: &[Parameter], prelude: &mut Vec<Stmt>) -> Vec<Parameter> {
        params
            .iter()
            .map(|param| match param {
                Parameter::Regular {
                    name,
                    default,
                    span,
                } => Parameter::Regular {
                    name: name.clone(),
                    default: default
                        .as_ref()
                        .map(|default| self.desugar_expr(default, prelude)),
                    span: span.clone(),
                },
                Parameter::VarArgs { name, span } => Parameter::VarArgs {
                    name: name.clone(),
                    span: span.clone(),
                },
                Parameter::KwArgs { name, span } => Parameter::KwArgs {
                    name: name.clone(),
                    span: span.clone(),
                },
            })
            .collect()
    }

    fn desugar_argument(&mut self, arg: &Argument, prelude: &mut Vec<Stmt>) -> Argument {
        match arg {
            Argument::Positional { value, span } => Argument::Positional {
                value: self.desugar_expr(value, prelude),
                span: span.clone(),
            },
            Argument::Keyword { name, value, span } => Argument::Keyword {
                name: name.clone(),
                value: self.desugar_expr(value, prelude),
                span: span.clone(),
            },
            Argument::Star { value, span } => Argument::Star {
                value: self.desugar_expr(value, prelude),
                span: span.clone(),
            },
            Argument::KwStar { value, span } => Argument::KwStar {
                value: self.desugar_expr(value, prelude),
                span: span.clone(),
            },
        }
    }

    fn desugar_assign_target(
        &mut self,
        target: &AssignTarget,
        prelude: &mut Vec<Stmt>,
    ) -> AssignTarget {
        match target {
            AssignTarget::Name { name, span } => AssignTarget::Name {
                name: name.clone(),
                span: span.clone(),
            },
            AssignTarget::Attribute { value, attr, span } => AssignTarget::Attribute {
                value: Box::new(self.desugar_expr(value, prelude)),
                attr: attr.clone(),
                span: span.clone(),
            },
            AssignTarget::Index { value, index, span } => AssignTarget::Index {
                value: Box::new(self.desugar_expr(value, prelude)),
                index: Box::new(self.desugar_expr(index, prelude)),
                span: span.clone(),
            },
            AssignTarget::Starred { target, span } => AssignTarget::Starred {
                target: Box::new(self.desugar_assign_target(target, prelude)),
                span: span.clone(),
            },
            AssignTarget::Tuple { elements, span } => AssignTarget::Tuple {
                elements: elements
                    .iter()
                    .map(|element| self.desugar_assign_target(element, prelude))
                    .collect(),
                span: span.clone(),
            },
            AssignTarget::List { elements, span } => AssignTarget::List {
                elements: elements
                    .iter()
                    .map(|element| self.desugar_assign_target(element, prelude))
                    .collect(),
                span: span.clone(),
            },
        }
    }

    fn desugar_regex_pattern(
        &mut self,
        pattern: &RegexPattern,
        prelude: &mut Vec<Stmt>,
    ) -> RegexPattern {
        match pattern {
            RegexPattern::Literal(text) => RegexPattern::Literal(text.clone()),
            RegexPattern::Interpolated(parts) => RegexPattern::Interpolated(
                parts
                    .iter()
                    .map(|part| self.desugar_fstring_part(part, prelude))
                    .collect(),
            ),
        }
    }

    fn desugar_fstring_part(&mut self, part: &FStringPart, prelude: &mut Vec<Stmt>) -> FStringPart {
        match part {
            FStringPart::Text(text) => FStringPart::Text(text.clone()),
            FStringPart::Expr(expr) => FStringPart::Expr(self.desugar_fstring_expr(expr, prelude)),
        }
    }

    fn desugar_fstring_expr(&mut self, expr: &FStringExpr, prelude: &mut Vec<Stmt>) -> FStringExpr {
        let format_spec = expr.format_spec.as_ref().map(|parts| {
            parts
                .iter()
                .map(|part| self.desugar_fstring_part(part, prelude))
                .collect()
        });
        FStringExpr {
            expr: Box::new(self.desugar_expr(&expr.expr, prelude)),
            conversion: expr.conversion,
            format_spec,
        }
    }

    fn is_hoistable_compound(expr: &Expr) -> bool {
        matches!(
            expr,
            Expr::Block { .. }
                | Expr::If { .. }
                | Expr::While { .. }
                | Expr::For { .. }
                | Expr::Try { .. }
                | Expr::With { .. }
        )
    }

    fn desugar_expr(&mut self, expr: &Expr, prelude: &mut Vec<Stmt>) -> Expr {
        let result = self.desugar_expr_no_hoist(expr, prelude);
        // Named def in expression context: hoist and return name reference.
        // (Anonymous defs were already hoisted in desugar_expr_no_hoist.)
        if let Expr::Def {
            name: Some(ref name),
            ref span,
            ..
        } = result
        {
            let name = name.clone();
            let span = span.clone();
            prelude.push(Stmt::Expr {
                value: result,
                semicolon_terminated: false,
                span: span.clone(),
            });
            return Expr::Name { name, span };
        }
        if Self::is_hoistable_compound(&result) {
            let span = Self::expr_span(&result);
            self.inline_compound_expr(result, prelude, &span)
        } else {
            result
        }
    }

    fn expr_span(expr: &Expr) -> SourceSpan {
        match expr {
            Expr::Block { span, .. }
            | Expr::If { span, .. }
            | Expr::While { span, .. }
            | Expr::For { span, .. }
            | Expr::Try { span, .. }
            | Expr::With { span, .. } => span.clone(),
            _ => unreachable!(),
        }
    }

    fn desugar_expr_no_hoist(&mut self, expr: &Expr, prelude: &mut Vec<Stmt>) -> Expr {
        match expr {
            Expr::Name { name, span } => Expr::Name {
                name: name.clone(),
                span: span.clone(),
            },
            Expr::Placeholder { span } => Expr::Placeholder { span: span.clone() },
            Expr::Number { value, span } => Expr::Number {
                value: value.clone(),
                span: span.clone(),
            },
            Expr::String {
                value,
                raw,
                bytes,
                delimiter,
                span,
            } => Expr::String {
                value: value.clone(),
                raw: *raw,
                bytes: *bytes,
                delimiter: *delimiter,
                span: span.clone(),
            },
            Expr::FString { parts, bytes, span } => Expr::FString {
                parts: parts
                    .iter()
                    .map(|part| self.desugar_fstring_part(part, prelude))
                    .collect(),
                bytes: *bytes,
                span: span.clone(),
            },
            Expr::Bool { value, span } => Expr::Bool {
                value: *value,
                span: span.clone(),
            },
            Expr::None { span } => Expr::None { span: span.clone() },
            Expr::Unary { op, expr, span } => Expr::Unary {
                op: *op,
                expr: Box::new(self.desugar_expr(expr, prelude)),
                span: span.clone(),
            },
            Expr::Binary {
                left,
                op,
                right,
                span,
            } => Expr::Binary {
                left: Box::new(self.desugar_expr(left, prelude)),
                op: *op,
                right: Box::new(self.desugar_expr(right, prelude)),
                span: span.clone(),
            },
            Expr::AugAssign {
                target,
                op,
                value,
                span,
            } => Expr::AugAssign {
                target: Box::new(self.desugar_assign_target(target, prelude)),
                op: *op,
                value: Box::new(self.desugar_expr(value, prelude)),
                span: span.clone(),
            },
            Expr::PrefixIncr { op, target, span } => Expr::PrefixIncr {
                op: *op,
                target: Box::new(self.desugar_assign_target(target, prelude)),
                span: span.clone(),
            },
            Expr::PostfixIncr { op, target, span } => Expr::PostfixIncr {
                op: *op,
                target: Box::new(self.desugar_assign_target(target, prelude)),
                span: span.clone(),
            },
            Expr::Compare {
                left,
                ops,
                comparators,
                span,
            } => Expr::Compare {
                left: Box::new(self.desugar_expr(left, prelude)),
                ops: ops.clone(),
                comparators: comparators
                    .iter()
                    .map(|expr| self.desugar_expr(expr, prelude))
                    .collect(),
                span: span.clone(),
            },
            Expr::Yield { value, span } => Expr::Yield {
                value: value
                    .as_ref()
                    .map(|expr| Box::new(self.desugar_expr(expr, prelude))),
                span: span.clone(),
            },
            Expr::YieldFrom { expr, span } => Expr::YieldFrom {
                expr: Box::new(self.desugar_expr(expr, prelude)),
                span: span.clone(),
            },
            Expr::Regex { pattern, span } => Expr::Regex {
                pattern: self.desugar_regex_pattern(pattern, prelude),
                span: span.clone(),
            },
            Expr::RegexMatch {
                value,
                pattern,
                span,
            } => Expr::RegexMatch {
                value: Box::new(self.desugar_expr(value, prelude)),
                pattern: self.desugar_regex_pattern(pattern, prelude),
                span: span.clone(),
            },
            Expr::Subprocess { kind, parts, span } => Expr::Subprocess {
                kind: *kind,
                parts: parts
                    .iter()
                    .map(|part| self.desugar_fstring_part(part, prelude))
                    .collect(),
                span: span.clone(),
            },
            Expr::StructuredAccessor { query, span } => Expr::StructuredAccessor {
                query: query.clone(),
                span: span.clone(),
            },
            Expr::Call { func, args, span } => Expr::Call {
                func: Box::new(self.desugar_expr(func, prelude)),
                args: args
                    .iter()
                    .map(|arg| self.desugar_argument(arg, prelude))
                    .collect(),
                span: span.clone(),
            },
            Expr::Attribute { value, attr, span } => Expr::Attribute {
                value: Box::new(self.desugar_expr(value, prelude)),
                attr: attr.clone(),
                span: span.clone(),
            },
            Expr::Index { value, index, span } => Expr::Index {
                value: Box::new(self.desugar_expr(value, prelude)),
                index: Box::new(self.desugar_expr(index, prelude)),
                span: span.clone(),
            },
            Expr::Paren { expr, span } => Expr::Paren {
                expr: Box::new(self.desugar_expr(expr, prelude)),
                span: span.clone(),
            },
            Expr::FieldIndex { index, span } => Expr::FieldIndex {
                index: index.clone(),
                span: span.clone(),
            },
            Expr::List { elements, span } => Expr::List {
                elements: elements
                    .iter()
                    .map(|expr| self.desugar_expr(expr, prelude))
                    .collect(),
                span: span.clone(),
            },
            Expr::Tuple { elements, span } => Expr::Tuple {
                elements: elements
                    .iter()
                    .map(|expr| self.desugar_expr(expr, prelude))
                    .collect(),
                span: span.clone(),
            },
            Expr::Set { elements, span } => Expr::Set {
                elements: elements
                    .iter()
                    .map(|expr| self.desugar_expr(expr, prelude))
                    .collect(),
                span: span.clone(),
            },
            Expr::Dict { entries, span } => Expr::Dict {
                entries: entries
                    .iter()
                    .map(|(key, value)| {
                        (
                            self.desugar_expr(key, prelude),
                            self.desugar_expr(value, prelude),
                        )
                    })
                    .collect(),
                span: span.clone(),
            },
            Expr::Slice { start, end, span } => Expr::Slice {
                start: start
                    .as_ref()
                    .map(|expr| Box::new(self.desugar_expr(expr, prelude))),
                end: end
                    .as_ref()
                    .map(|expr| Box::new(self.desugar_expr(expr, prelude))),
                span: span.clone(),
            },
            Expr::ListComp {
                element,
                target,
                iter,
                ifs,
                span,
            } => Expr::ListComp {
                element: Box::new(self.desugar_expr(element, prelude)),
                target: target.clone(),
                iter: Box::new(self.desugar_expr(iter, prelude)),
                ifs: ifs
                    .iter()
                    .map(|expr| self.desugar_expr(expr, prelude))
                    .collect(),
                span: span.clone(),
            },
            Expr::DictComp {
                key,
                value,
                target,
                iter,
                ifs,
                span,
            } => Expr::DictComp {
                key: Box::new(self.desugar_expr(key, prelude)),
                value: Box::new(self.desugar_expr(value, prelude)),
                target: target.clone(),
                iter: Box::new(self.desugar_expr(iter, prelude)),
                ifs: ifs
                    .iter()
                    .map(|expr| self.desugar_expr(expr, prelude))
                    .collect(),
                span: span.clone(),
            },
            Expr::Block { stmts, span } => Expr::Block {
                stmts: self.desugar_block(stmts),
                span: span.clone(),
            },
            Expr::If {
                cond,
                body,
                elifs,
                else_body,
                span,
            } => {
                let cond = self.desugar_condition(cond, prelude);
                let body = self.desugar_block(body);
                let elifs: Vec<_> = elifs
                    .iter()
                    .map(|(elif_cond, elif_body)| {
                        let mut elif_prelude = Vec::new();
                        let c = self.desugar_condition(elif_cond, &mut elif_prelude);
                        let b = self.desugar_block(elif_body);
                        (c, b)
                    })
                    .collect();
                let else_body = else_body.as_ref().map(|body| self.desugar_block(body));
                Expr::If {
                    cond,
                    body,
                    elifs,
                    else_body,
                    span: span.clone(),
                }
            }
            Expr::While {
                cond,
                body,
                else_body,
                span,
            } => {
                let cond = self.desugar_condition(cond, prelude);
                let body = self.desugar_block(body);
                let else_body = else_body.as_ref().map(|body| self.desugar_block(body));
                Expr::While {
                    cond,
                    body,
                    else_body,
                    span: span.clone(),
                }
            }
            Expr::For {
                target,
                iter,
                body,
                else_body,
                span,
            } => {
                let target = self.desugar_assign_target(target, prelude);
                let iter = Box::new(self.desugar_expr(iter, prelude));
                let body = self.desugar_block(body);
                let else_body = else_body.as_ref().map(|body| self.desugar_block(body));
                Expr::For {
                    target,
                    iter,
                    body,
                    else_body,
                    span: span.clone(),
                }
            }
            Expr::Def {
                name,
                params,
                body,
                span,
            } => {
                let params = self.desugar_params(params, prelude);
                let body = self.desugar_block(body);
                if name.is_some() {
                    // Named def: pass through. desugar_expr will hoist if in expression context.
                    Expr::Def {
                        name: name.clone(),
                        params,
                        body,
                        span: span.clone(),
                    }
                } else {
                    // Anonymous def: always hoist (needs a generated name to exist).
                    let gen_name = self.next_block_name();
                    prelude.push(Stmt::Expr {
                        value: Expr::Def {
                            name: Some(gen_name.clone()),
                            params,
                            body,
                            span: span.clone(),
                        },
                        semicolon_terminated: false,
                        span: span.clone(),
                    });
                    Expr::Name {
                        name: gen_name,
                        span: span.clone(),
                    }
                }
            }
            Expr::Class { name, body, span } => {
                let body = self.desugar_block(body);
                Expr::Class {
                    name: name.clone(),
                    body,
                    span: span.clone(),
                }
            }
            Expr::Try {
                body,
                handlers,
                else_body,
                finally_body,
                span,
            } => {
                let body = self.desugar_block(body);
                let handlers = handlers
                    .iter()
                    .map(|handler| self.desugar_except_handler(handler, prelude))
                    .collect();
                let else_body = else_body.as_ref().map(|body| self.desugar_block(body));
                let finally_body = finally_body.as_ref().map(|body| self.desugar_block(body));
                Expr::Try {
                    body,
                    handlers,
                    else_body,
                    finally_body,
                    span: span.clone(),
                }
            }
            Expr::With { items, body, span } => {
                let items = items
                    .iter()
                    .map(|item| self.desugar_with_item(item, prelude))
                    .collect();
                let body = self.desugar_block(body);
                Expr::With {
                    items,
                    body,
                    span: span.clone(),
                }
            }
            Expr::Awk {
                sources,
                body,
                span,
            } => Expr::Awk {
                sources: sources
                    .iter()
                    .map(|src| self.desugar_argument(src, prelude))
                    .collect(),
                body: self.desugar_block(body),
                span: span.clone(),
            },
            Expr::Xargs {
                sources,
                body,
                span,
            } => Expr::Xargs {
                sources: sources
                    .iter()
                    .map(|src| self.desugar_argument(src, prelude))
                    .collect(),
                body: self.desugar_block(body),
                span: span.clone(),
            },
        }
    }
}
