use std::collections::{HashMap, HashSet};

use crate::frontend::token::Span;
use crate::middle::mir::*;

#[derive(Debug, PartialEq, Eq)]
pub struct MirWarning {
    pub message: String,
    pub span: Span,
}

/// Eliminates unreachable basic blocks from the Control Flow Graph.
pub fn eliminate_dead_code(module: &mut MirModule) -> Vec<MirWarning> {
    let mut warnings = Vec::new();
    for func in &mut module.functions {
        optimize_function(func, &mut warnings);
    }
    warnings.sort_by_key(|w| w.span.start);
    warnings.dedup_by_key(|w| w.span.start);
    warnings
}

fn optimize_function(func: &mut MirFunction, warnings: &mut Vec<MirWarning>) {
    if func.blocks.is_empty() {
        return;
    }

    let mut visited = HashSet::new();
    // The first block generated is always the entry block.
    let mut worklist = vec![func.blocks[0].id];

    let block_map: HashMap<BlockId, &BasicBlock> = func.blocks.iter().map(|b| (b.id, b)).collect();

    while let Some(id) = worklist.pop() {
        if !visited.insert(id) {
            continue; // Already visited
        }

        if let Some(block) = block_map.get(&id) {
            match &block.terminator.kind {
                TerminatorKind::Goto { target } => worklist.push(*target),
                TerminatorKind::CondBranch {
                    cond,
                    target_true,
                    target_false,
                } => match cond {
                    Operand::Constant(ConstantValue::Boolean(true)) => {
                        worklist.push(*target_true);
                    }
                    Operand::Constant(ConstantValue::Boolean(false)) => {
                        worklist.push(*target_false);
                    }
                    _ => {
                        worklist.push(*target_true);
                        worklist.push(*target_false);
                    }
                },
                TerminatorKind::Return { .. } | TerminatorKind::Unreachable => {}
            }
        }
    }

    // Emit warnings for the dead code we are about to remove
    for block in &func.blocks {
        if !visited.contains(&block.id) {
            let is_implicit_unreachable = block.statements.is_empty()
                && matches!(block.terminator.kind, TerminatorKind::Unreachable);

            // Do not warn about implicit unreachables (e.g. missing returns handled by builder)
            if !is_implicit_unreachable {
                let span = block
                    .statements
                    .first()
                    .map(|s| s.span)
                    .unwrap_or(block.terminator.span);
                warnings.push(MirWarning {
                    message: "unreachable code".to_string(),
                    span,
                });
            }
        }
    }

    // Retain only the blocks that were successfully reached during our traversal
    func.blocks.retain(|b| visited.contains(&b.id));
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::frontend::sema::Ty;
    use crate::frontend::token::Span;

    #[test]
    fn test_eliminate_dead_blocks() {
        let mut module = MirModule {
            structs: HashMap::new(),
            functions: vec![MirFunction {
                name: "test_func".to_string(),
                params: vec![],
                return_type: Ty::Unit,
                locals: vec![],
                blocks: vec![
                    BasicBlock {
                        id: BlockId(0), // Entry block
                        statements: vec![],
                        terminator: Terminator {
                            kind: TerminatorKind::Goto { target: BlockId(1) },
                            span: Span::new(0, 0),
                        },
                    },
                    BasicBlock {
                        id: BlockId(1), // Reachable
                        statements: vec![],
                        terminator: Terminator {
                            kind: TerminatorKind::Return { value: None },
                            span: Span::new(0, 0),
                        },
                    },
                    BasicBlock {
                        id: BlockId(2), // Unreachable
                        statements: vec![],
                        terminator: Terminator {
                            kind: TerminatorKind::Goto { target: BlockId(1) },
                            span: Span::new(0, 0),
                        },
                    },
                ],
            }],
            extern_functions: vec![],
        };

        let warnings = eliminate_dead_code(&mut module);
        assert_eq!(warnings.len(), 1);

        let blocks = &module.functions[0].blocks;
        assert_eq!(blocks.len(), 2, "Expected exactly 2 reachable blocks");
        assert!(blocks.iter().any(|b| b.id == BlockId(0)));
        assert!(blocks.iter().any(|b| b.id == BlockId(1)));
        assert!(
            !blocks.iter().any(|b| b.id == BlockId(2)),
            "Block 2 should have been eliminated"
        );
    }

    #[test]
    fn test_eliminate_dead_cond_branch() {
        let mut module = MirModule {
            structs: HashMap::new(),
            functions: vec![MirFunction {
                name: "test_cond_func".to_string(),
                params: vec![],
                return_type: Ty::Unit,
                locals: vec![],
                blocks: vec![
                    BasicBlock {
                        id: BlockId(0), // Entry block
                        statements: vec![],
                        terminator: Terminator {
                            kind: TerminatorKind::CondBranch {
                                cond: Operand::Constant(ConstantValue::Boolean(true)),
                                target_true: BlockId(1),
                                target_false: BlockId(2),
                            },
                            span: Span::new(0, 0),
                        },
                    },
                    BasicBlock {
                        id: BlockId(1), // Reachable (true branch)
                        statements: vec![],
                        terminator: Terminator {
                            kind: TerminatorKind::Return { value: None },
                            span: Span::new(0, 0),
                        },
                    },
                    BasicBlock {
                        id: BlockId(2), // Unreachable (false branch)
                        statements: vec![],
                        terminator: Terminator {
                            kind: TerminatorKind::Return { value: None },
                            span: Span::new(0, 0),
                        },
                    },
                ],
            }],
            extern_functions: vec![],
        };

        let warnings = eliminate_dead_code(&mut module);
        assert_eq!(warnings.len(), 1);

        let blocks = &module.functions[0].blocks;
        assert_eq!(blocks.len(), 2, "Expected exactly 2 reachable blocks");
        assert!(blocks.iter().any(|b| b.id == BlockId(0)));
        assert!(blocks.iter().any(|b| b.id == BlockId(1)));
        assert!(
            !blocks.iter().any(|b| b.id == BlockId(2)),
            "Block 2 should have been eliminated"
        );
    }
}
