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

use crate::diag::{Diag, DiagCode, DiagCollector};
use crate::ir::DType;
use crate::parse::ast::{Expr, FuncDef, Spanned};

#[derive(Debug, thiserror::Error)]
pub enum RenameError {
    #[error("rename mapping references unknown parameter `{name}`")]
    UnknownParam { name: String },
    #[error("multiple source parameters map to the same target `{name}`")]
    DuplicateTarget { name: String },
    #[error("target name `{name}` conflicts with an existing symbol")]
    NameConflict { name: String },
}

#[derive(Debug, thiserror::Error)]
pub enum ComposeError {
    #[error("parameter `{name}` does not exist in the outer expression")]
    ParamNotFound { name: String },
    #[error("symbol `{symbol}` has incompatible constraints: {outer} vs {inner}")]
    ConstraintConflict {
        symbol: String,
        outer: DType,
        inner: DType,
    },
    #[error(transparent)]
    Diagnostic(#[from] Diag),
}

pub fn rename_params(
    func: &FuncDef,
    mapping: impl IntoIterator<Item = (String, String)>,
) -> Result<FuncDef, RenameError> {
    let param_names: HashSet<String> = func.params.iter().map(|(name, _)| name.clone()).collect();
    let mut rename_map = HashMap::new();
    let mut seen_targets = HashSet::new();
    for (src, dst) in mapping {
        if !param_names.contains(&src) {
            return Err(RenameError::UnknownParam { name: src });
        }
        if !seen_targets.insert(dst.clone()) {
            return Err(RenameError::DuplicateTarget { name: dst });
        }
        rename_map.insert(src, dst);
    }

    let mut reserved = collect_let_bindings(&func.body);
    for (name, _) in &func.params {
        if !rename_map.contains_key(name) {
            reserved.insert(name.clone());
        }
    }
    for target in rename_map.values() {
        if reserved.contains(target) {
            return Err(RenameError::NameConflict {
                name: target.clone(),
            });
        }
    }

    let params = func
        .params
        .iter()
        .map(|(name, dtype)| {
            (
                rename_map
                    .get(name)
                    .cloned()
                    .unwrap_or_else(|| name.clone()),
                *dtype,
            )
        })
        .collect();
    let body = rename_param_refs(&func.body, &rename_map, &mut HashSet::new());
    Ok(FuncDef { params, body })
}

pub fn compose_with_diags(
    outer: &FuncDef,
    substitute_param: &str,
    inner: &FuncDef,
    diags: &mut DiagCollector,
) -> Result<FuncDef, ComposeError> {
    if !outer
        .params
        .iter()
        .any(|(name, _)| name == substitute_param)
    {
        return Err(ComposeError::ParamNotFound {
            name: substitute_param.to_string(),
        });
    }

    let mut outer_let_names = collect_let_bindings(&outer.body);
    let mut forbidden_for_inner = outer_let_names.clone();
    for (name, _) in &outer.params {
        forbidden_for_inner.insert(name.clone());
    }
    for (name, _) in &inner.params {
        forbidden_for_inner.insert(name.clone());
    }
    let inner_body = alpha_rename_lets(&inner.body, &forbidden_for_inner);

    let replacement_free: HashSet<String> =
        inner.params.iter().map(|(name, _)| name.clone()).collect();
    outer_let_names.extend(replacement_free.iter().cloned());
    let mut used_names = outer_let_names;
    for (name, _) in &outer.params {
        used_names.insert(name.clone());
    }
    for (name, _) in &inner.params {
        used_names.insert(name.clone());
    }

    let body = substitute_param_with_expr(
        &outer.body,
        substitute_param,
        &inner_body,
        &replacement_free,
        &mut used_names,
        &mut HashMap::new(),
        false,
    );

    let mut params: Vec<(String, DType)> = Vec::new();
    let mut index: HashMap<String, usize> = HashMap::new();
    for (name, dtype) in &outer.params {
        if name == substitute_param {
            continue;
        }
        index.insert(name.clone(), params.len());
        params.push((name.clone(), *dtype));
    }
    for (name, dtype) in &inner.params {
        if let Some(existing_idx) = index.get(name).copied() {
            let existing_dtype = params[existing_idx].1;
            if existing_dtype != *dtype {
                return Err(ComposeError::ConstraintConflict {
                    symbol: name.clone(),
                    outer: existing_dtype,
                    inner: *dtype,
                });
            }
            diags.emit(Diag::new(
                DiagCode::ComposeSymbolReuse,
                format!(
                    "symbol `{name}` appears in both outer and inner expressions and is treated as the same variable"
                ),
                None,
                Some("rename one side if you intended distinct symbols".to_string()),
            ))?;
            continue;
        }
        index.insert(name.clone(), params.len());
        params.push((name.clone(), *dtype));
    }

    Ok(FuncDef { params, body })
}

pub fn compose(
    outer: &FuncDef,
    substitute_param: &str,
    inner: &FuncDef,
) -> Result<FuncDef, ComposeError> {
    let mut collector = DiagCollector::default();
    compose_with_diags(outer, substitute_param, inner, &mut collector)
}

fn rename_param_refs(
    expr: &Spanned<Expr>,
    mapping: &HashMap<String, String>,
    shadowed: &mut HashSet<String>,
) -> Spanned<Expr> {
    match &expr.node {
        Expr::LitBool(_)
        | Expr::LitInt(_, _)
        | Expr::LitFloat(_, _)
        | Expr::LitComplex(_, _, _)
        | Expr::ConstE
        | Expr::ConstPi => expr.clone(),
        Expr::Var(name) => {
            if shadowed.contains(name) {
                return expr.clone();
            }
            if let Some(new_name) = mapping.get(name) {
                Spanned::new(Expr::Var(new_name.clone()), expr.span)
            } else {
                expr.clone()
            }
        }
        Expr::Unary(op, child) => Spanned::new(
            Expr::Unary(*op, Box::new(rename_param_refs(child, mapping, shadowed))),
            expr.span,
        ),
        Expr::Binary(op, lhs, rhs) => Spanned::new(
            Expr::Binary(
                *op,
                Box::new(rename_param_refs(lhs, mapping, shadowed)),
                Box::new(rename_param_refs(rhs, mapping, shadowed)),
            ),
            expr.span,
        ),
        Expr::IfElse(cond, then_expr, else_expr) => Spanned::new(
            Expr::IfElse(
                Box::new(rename_param_refs(cond, mapping, shadowed)),
                Box::new(rename_param_refs(then_expr, mapping, shadowed)),
                Box::new(rename_param_refs(else_expr, mapping, shadowed)),
            ),
            expr.span,
        ),
        Expr::Call(name, args) => Spanned::new(
            Expr::Call(
                name.clone(),
                args.iter()
                    .map(|arg| rename_param_refs(arg, mapping, shadowed))
                    .collect(),
            ),
            expr.span,
        ),
        Expr::Cast(dtype, child) => Spanned::new(
            Expr::Cast(
                *dtype,
                Box::new(rename_param_refs(child, mapping, shadowed)),
            ),
            expr.span,
        ),
        Expr::Tuple(items) => Spanned::new(
            Expr::Tuple(
                items
                    .iter()
                    .map(|item| rename_param_refs(item, mapping, shadowed))
                    .collect(),
            ),
            expr.span,
        ),
        Expr::Let(name, value, body) => {
            let new_value = rename_param_refs(value, mapping, shadowed);
            shadowed.insert(name.clone());
            let new_body = rename_param_refs(body, mapping, shadowed);
            shadowed.remove(name);
            Spanned::new(
                Expr::Let(name.clone(), Box::new(new_value), Box::new(new_body)),
                expr.span,
            )
        }
    }
}

fn collect_let_bindings(expr: &Spanned<Expr>) -> HashSet<String> {
    let mut names = HashSet::new();
    collect_let_bindings_into(expr, &mut names);
    names
}

fn collect_let_bindings_into(expr: &Spanned<Expr>, names: &mut HashSet<String>) {
    match &expr.node {
        Expr::Unary(_, child) | Expr::Cast(_, child) => collect_let_bindings_into(child, names),
        Expr::Binary(_, lhs, rhs) => {
            collect_let_bindings_into(lhs, names);
            collect_let_bindings_into(rhs, names);
        }
        Expr::IfElse(cond, then_expr, else_expr) => {
            collect_let_bindings_into(cond, names);
            collect_let_bindings_into(then_expr, names);
            collect_let_bindings_into(else_expr, names);
        }
        Expr::Call(_, args) | Expr::Tuple(args) => {
            for arg in args {
                collect_let_bindings_into(arg, names);
            }
        }
        Expr::Let(name, value, body) => {
            names.insert(name.clone());
            collect_let_bindings_into(value, names);
            collect_let_bindings_into(body, names);
        }
        Expr::LitBool(_)
        | Expr::LitInt(_, _)
        | Expr::LitFloat(_, _)
        | Expr::LitComplex(_, _, _)
        | Expr::ConstE
        | Expr::ConstPi
        | Expr::Var(_) => {}
    }
}

fn fresh_name(base: &str, used: &mut HashSet<String>) -> String {
    let mut idx = 1usize;
    loop {
        let candidate = format!("{base}_{idx}");
        if used.insert(candidate.clone()) {
            return candidate;
        }
        idx += 1;
    }
}

fn alpha_rename_lets(expr: &Spanned<Expr>, forbidden: &HashSet<String>) -> Spanned<Expr> {
    let mut used = forbidden.clone();
    used.extend(collect_let_bindings(expr));
    alpha_rename_lets_rec(expr, forbidden, &mut used, &mut HashMap::new())
}

fn alpha_rename_lets_rec(
    expr: &Spanned<Expr>,
    forbidden: &HashSet<String>,
    used: &mut HashSet<String>,
    env: &mut HashMap<String, String>,
) -> Spanned<Expr> {
    match &expr.node {
        Expr::LitBool(_)
        | Expr::LitInt(_, _)
        | Expr::LitFloat(_, _)
        | Expr::LitComplex(_, _, _)
        | Expr::ConstE
        | Expr::ConstPi => expr.clone(),
        Expr::Var(name) => {
            if let Some(mapped) = env.get(name) {
                Spanned::new(Expr::Var(mapped.clone()), expr.span)
            } else {
                expr.clone()
            }
        }
        Expr::Unary(op, child) => Spanned::new(
            Expr::Unary(
                *op,
                Box::new(alpha_rename_lets_rec(child, forbidden, used, env)),
            ),
            expr.span,
        ),
        Expr::Binary(op, lhs, rhs) => Spanned::new(
            Expr::Binary(
                *op,
                Box::new(alpha_rename_lets_rec(lhs, forbidden, used, env)),
                Box::new(alpha_rename_lets_rec(rhs, forbidden, used, env)),
            ),
            expr.span,
        ),
        Expr::IfElse(cond, then_expr, else_expr) => Spanned::new(
            Expr::IfElse(
                Box::new(alpha_rename_lets_rec(cond, forbidden, used, env)),
                Box::new(alpha_rename_lets_rec(then_expr, forbidden, used, env)),
                Box::new(alpha_rename_lets_rec(else_expr, forbidden, used, env)),
            ),
            expr.span,
        ),
        Expr::Call(name, args) => Spanned::new(
            Expr::Call(
                name.clone(),
                args.iter()
                    .map(|arg| alpha_rename_lets_rec(arg, forbidden, used, env))
                    .collect(),
            ),
            expr.span,
        ),
        Expr::Cast(dtype, child) => Spanned::new(
            Expr::Cast(
                *dtype,
                Box::new(alpha_rename_lets_rec(child, forbidden, used, env)),
            ),
            expr.span,
        ),
        Expr::Tuple(items) => Spanned::new(
            Expr::Tuple(
                items
                    .iter()
                    .map(|item| alpha_rename_lets_rec(item, forbidden, used, env))
                    .collect(),
            ),
            expr.span,
        ),
        Expr::Let(name, value, body) => {
            let new_value = alpha_rename_lets_rec(value, forbidden, used, env);
            let needs_rename = forbidden.contains(name);
            let out_name = if needs_rename {
                fresh_name(name, used)
            } else {
                name.clone()
            };
            let previous = env.insert(name.clone(), out_name.clone());
            let new_body = alpha_rename_lets_rec(body, forbidden, used, env);
            if let Some(prev) = previous {
                env.insert(name.clone(), prev);
            } else {
                env.remove(name);
            }
            Spanned::new(
                Expr::Let(out_name, Box::new(new_value), Box::new(new_body)),
                expr.span,
            )
        }
    }
}

fn substitute_param_with_expr(
    expr: &Spanned<Expr>,
    target_param: &str,
    replacement: &Spanned<Expr>,
    replacement_free: &HashSet<String>,
    used_names: &mut HashSet<String>,
    rename_env: &mut HashMap<String, String>,
    target_shadowed: bool,
) -> Spanned<Expr> {
    match &expr.node {
        Expr::LitBool(_)
        | Expr::LitInt(_, _)
        | Expr::LitFloat(_, _)
        | Expr::LitComplex(_, _, _)
        | Expr::ConstE
        | Expr::ConstPi => expr.clone(),
        Expr::Var(name) => {
            let resolved = rename_env
                .get(name)
                .cloned()
                .unwrap_or_else(|| name.clone());
            if !target_shadowed && resolved == target_param {
                replacement.clone()
            } else if resolved == *name {
                expr.clone()
            } else {
                Spanned::new(Expr::Var(resolved), expr.span)
            }
        }
        Expr::Unary(op, child) => Spanned::new(
            Expr::Unary(
                *op,
                Box::new(substitute_param_with_expr(
                    child,
                    target_param,
                    replacement,
                    replacement_free,
                    used_names,
                    rename_env,
                    target_shadowed,
                )),
            ),
            expr.span,
        ),
        Expr::Binary(op, lhs, rhs) => Spanned::new(
            Expr::Binary(
                *op,
                Box::new(substitute_param_with_expr(
                    lhs,
                    target_param,
                    replacement,
                    replacement_free,
                    used_names,
                    rename_env,
                    target_shadowed,
                )),
                Box::new(substitute_param_with_expr(
                    rhs,
                    target_param,
                    replacement,
                    replacement_free,
                    used_names,
                    rename_env,
                    target_shadowed,
                )),
            ),
            expr.span,
        ),
        Expr::IfElse(cond, then_expr, else_expr) => Spanned::new(
            Expr::IfElse(
                Box::new(substitute_param_with_expr(
                    cond,
                    target_param,
                    replacement,
                    replacement_free,
                    used_names,
                    rename_env,
                    target_shadowed,
                )),
                Box::new(substitute_param_with_expr(
                    then_expr,
                    target_param,
                    replacement,
                    replacement_free,
                    used_names,
                    rename_env,
                    target_shadowed,
                )),
                Box::new(substitute_param_with_expr(
                    else_expr,
                    target_param,
                    replacement,
                    replacement_free,
                    used_names,
                    rename_env,
                    target_shadowed,
                )),
            ),
            expr.span,
        ),
        Expr::Call(name, args) => Spanned::new(
            Expr::Call(
                name.clone(),
                args.iter()
                    .map(|arg| {
                        substitute_param_with_expr(
                            arg,
                            target_param,
                            replacement,
                            replacement_free,
                            used_names,
                            rename_env,
                            target_shadowed,
                        )
                    })
                    .collect(),
            ),
            expr.span,
        ),
        Expr::Cast(dtype, child) => Spanned::new(
            Expr::Cast(
                *dtype,
                Box::new(substitute_param_with_expr(
                    child,
                    target_param,
                    replacement,
                    replacement_free,
                    used_names,
                    rename_env,
                    target_shadowed,
                )),
            ),
            expr.span,
        ),
        Expr::Tuple(items) => Spanned::new(
            Expr::Tuple(
                items
                    .iter()
                    .map(|item| {
                        substitute_param_with_expr(
                            item,
                            target_param,
                            replacement,
                            replacement_free,
                            used_names,
                            rename_env,
                            target_shadowed,
                        )
                    })
                    .collect(),
            ),
            expr.span,
        ),
        Expr::Let(name, value, body) => {
            let new_value = substitute_param_with_expr(
                value,
                target_param,
                replacement,
                replacement_free,
                used_names,
                rename_env,
                target_shadowed,
            );
            let mut out_name = name.clone();
            if replacement_free.contains(name) {
                out_name = fresh_name(name, used_names);
            }
            let previous = rename_env.insert(name.clone(), out_name.clone());
            let body_target_shadowed = target_shadowed || name == target_param;
            let new_body = substitute_param_with_expr(
                body,
                target_param,
                replacement,
                replacement_free,
                used_names,
                rename_env,
                body_target_shadowed,
            );
            if let Some(prev) = previous {
                rename_env.insert(name.clone(), prev);
            } else {
                rename_env.remove(name);
            }
            Spanned::new(
                Expr::Let(out_name, Box::new(new_value), Box::new(new_body)),
                expr.span,
            )
        }
    }
}
