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

use t_ree::declaration as t;
use t_ree::expression as t_expression;

fn collect_type_deps(inner_type: &t_ree::types::Type) -> HashSet<String> {
    let mut deps = HashSet::new();
    collect_type_deps_recursive(inner_type, &mut deps);
    deps
}

fn collect_type_deps_recursive(ty: &t_ree::types::Type, deps: &mut HashSet<String>) {
    use t_ree::types::Type;
    match ty {
        Type::Named(name) => {
            deps.insert(name.clone());
        }
        Type::Pointer(_, inner)
        | Type::Array(inner, _)
        | Type::Vector(inner, _)
        | Type::Slice(_, inner) => {
            collect_type_deps_recursive(inner, deps);
        }
        Type::Tuple(fields) | Type::Enum(fields) => {
            for field in fields {
                collect_type_deps_recursive(field, deps);
            }
        }
        _ => {}
    }
}

pub fn reorder_functions(module: &mut t::Module) {
    let function_indices: HashMap<String, usize> = module
        .iter()
        .enumerate()
        .filter_map(|(index, declaration)| {
            if let t::Declaration::Function(function) = declaration {
                Some((function.name.clone(), index))
            } else {
                None
            }
        })
        .collect();

    let call_deps: Vec<(usize, Vec<usize>)> = module
        .iter()
        .enumerate()
        .filter_map(|(index, declaration)| {
            if let t::Declaration::Function(function) = declaration {
                let deps = collect_called_functions(&function.body)
                    .into_iter()
                    .filter_map(|name| function_indices.get(&name).copied())
                    .filter(|&dep| dep != index)
                    .collect();
                Some((index, deps))
            } else {
                None
            }
        })
        .collect();

    let mut sorted_indices: Vec<usize> = Vec::new();
    let mut visited: HashSet<usize> = HashSet::new();
    let mut visiting: HashSet<usize> = HashSet::new();
    let dep_map: HashMap<usize, Vec<usize>> = call_deps.into_iter().collect();

    fn visit(
        index: usize,
        dep_map: &HashMap<usize, Vec<usize>>,
        visited: &mut HashSet<usize>,
        visiting: &mut HashSet<usize>,
        sorted: &mut Vec<usize>,
    ) {
        if visited.contains(&index) || visiting.contains(&index) {
            return;
        }
        visiting.insert(index);
        if let Some(deps) = dep_map.get(&index) {
            for &dep in deps {
                visit(dep, dep_map, visited, visiting, sorted);
            }
        }
        visiting.remove(&index);
        visited.insert(index);
        sorted.push(index);
    }

    for &(index, _) in &dep_map.iter().collect::<Vec<_>>() {
        visit(
            *index,
            &dep_map,
            &mut visited,
            &mut visiting,
            &mut sorted_indices,
        );
    }

    let type_indices: HashMap<String, usize> = module
        .iter()
        .enumerate()
        .filter_map(|(index, declaration)| {
            if let t::Declaration::Type(newtype) = declaration {
                Some((newtype.name.clone(), index))
            } else {
                None
            }
        })
        .collect();

    let type_deps: HashMap<usize, Vec<usize>> = module
        .iter()
        .enumerate()
        .filter_map(|(index, declaration)| {
            if let t::Declaration::Type(newtype) = declaration {
                let deps = collect_type_deps(&newtype.inner_type)
                    .into_iter()
                    .filter_map(|name| type_indices.get(&name).copied())
                    .filter(|&dep| dep != index)
                    .collect();
                Some((index, deps))
            } else {
                None
            }
        })
        .collect();

    let mut sorted_type_indices: Vec<usize> = Vec::new();
    let mut type_visited: HashSet<usize> = HashSet::new();
    let mut type_visiting: HashSet<usize> = HashSet::new();

    for &index in type_deps.keys() {
        visit(
            index,
            &type_deps,
            &mut type_visited,
            &mut type_visiting,
            &mut sorted_type_indices,
        );
    }

    let non_functions: Vec<_> = module
        .iter()
        .enumerate()
        .filter(|(index, declaration)| {
            !matches!(declaration, t::Declaration::Function(_)) && !type_deps.contains_key(index)
        })
        .map(|(index, _)| index)
        .collect();

    let mut result: Vec<t::Declaration> = Vec::with_capacity(module.len());
    let placeholder = t::Declaration::Import(String::new());
    let mut originals: Vec<t::Declaration> = std::mem::take(module);

    for &index in &non_functions {
        result.push(std::mem::replace(
            &mut originals[index],
            placeholder.clone(),
        ));
    }
    for &index in &sorted_type_indices {
        result.push(std::mem::replace(
            &mut originals[index],
            placeholder.clone(),
        ));
    }
    for &index in &sorted_indices {
        result.push(std::mem::replace(
            &mut originals[index],
            placeholder.clone(),
        ));
    }
    for original in originals {
        if !matches!(&original, t::Declaration::Import(path) if path.is_empty()) {
            result.push(original);
        }
    }

    *module = result;
}

fn collect_called_functions(block: &t_expression::Block) -> HashSet<String> {
    let mut names = HashSet::new();
    collect_calls_in_block(block, &mut names);
    names
}

fn collect_calls_in_block(block: &t_expression::Block, names: &mut HashSet<String>) {
    for statement in &block.statements {
        collect_calls_in_statement(statement, names);
    }
    if let Some(result) = &block.result {
        collect_calls_in_expression(result, names);
    }
}

fn collect_calls_in_statement(statement: &t_expression::Statement, names: &mut HashSet<String>) {
    use t_expression::Statement as S;
    match statement {
        S::Expression(expression) => collect_calls_in_expression(expression, names),
        S::Let { value, .. } => collect_calls_in_expression(value, names),
        S::Assign(target, value) => {
            collect_calls_in_expression(target, names);
            collect_calls_in_expression(value, names);
        }
        S::Return(Some(expression)) => collect_calls_in_expression(expression, names),
        S::Return(None) => {}
        S::Label {
            initial_arguments, ..
        } => {
            for argument in initial_arguments {
                collect_calls_in_expression(argument, names);
            }
        }
        S::Jump { arguments, .. } => {
            for argument in arguments {
                collect_calls_in_expression(argument, names);
            }
        }
        S::Defer(inner) => collect_calls_in_statement(inner, names),
        S::MultiReplace {
            targets, values, ..
        } => {
            for target in targets {
                collect_calls_in_expression(target, names);
            }
            for value in values {
                collect_calls_in_expression(value, names);
            }
        }
    }
}

fn collect_calls_in_expression(expression: &t_expression::Expression, names: &mut HashSet<String>) {
    use t_expression::ExpressionKind as E;
    match &expression.kind {
        E::Call(callee, arguments) => {
            if let E::Variable(name) = &callee.kind {
                names.insert(name.clone());
            }
            collect_calls_in_expression(callee, names);
            for argument in arguments {
                collect_calls_in_expression(argument, names);
            }
        }
        E::BinaryOperation(_, left, right)
        | E::Replace(left, right)
        | E::OpAssign(_, left, right) => {
            collect_calls_in_expression(left, names);
            collect_calls_in_expression(right, names);
        }
        E::UnaryOperation(_, operand)
        | E::Dereference(operand)
        | E::Convert(operand, _)
        | E::Transmute(operand, _)
        | E::Field(operand, _) => {
            collect_calls_in_expression(operand, names);
        }
        E::Index(object, index) => {
            collect_calls_in_expression(object, names);
            collect_calls_in_expression(index, names);
        }
        E::TypeConstruction(_, fields) => {
            for (_, expression) in fields {
                collect_calls_in_expression(expression, names);
            }
        }
        E::ArrayLiteral(elements) | E::TupleLiteral(elements) | E::Print(elements) => {
            for element in elements {
                collect_calls_in_expression(element, names);
            }
        }
        E::Block(block) => collect_calls_in_block(block, names),
        E::If {
            condition,
            then_branch,
            else_branch,
        } => {
            collect_calls_in_expression(condition, names);
            collect_calls_in_block(then_branch, names);
            if let Some(else_branch) = else_branch {
                collect_calls_in_block(else_branch, names);
            }
        }
        E::Match { value, arms } => {
            collect_calls_in_expression(value, names);
            for arm in arms {
                collect_calls_in_block(&arm.body, names);
            }
        }
        E::Slice(array, start, end) => {
            collect_calls_in_expression(array, names);
            if let Some(start) = start {
                collect_calls_in_expression(start, names);
            }
            if let Some(end) = end {
                collect_calls_in_expression(end, names);
            }
        }
        E::Literal(_) | E::Variable(_) | E::SizeOf(_) => {}
    }
}
