use nu_protocol::{
    Span,
    ast::{Argument, Call, Expr, Expression, Pipeline, Traverse},
};

use crate::{
    Fix, LintLevel, Replacement,
    ast::{call::CallExt, pipeline::PipelineExt},
    context::LintContext,
    rule::{DetectFix, Rule},
    violation::Detection,
};

/// Represents a cell path access like `$row.col1` that needs to be replaced
/// with `$col1`
struct CellPathReplacement {
    /// The full span of `$row.col1`
    span: Span,
    /// The field name (e.g., "col1")
    field: String,
}

struct TransposeEachPattern {
    each_call: Call,
    col1: String,
    col2: String,
    combined_span: nu_protocol::Span,
    transpose_span: nu_protocol::Span,
    each_span: nu_protocol::Span,
    closure_span: nu_protocol::Span,
    /// The closure parameter's `VarId` for AST-based span lookup
    param_var_id: nu_protocol::VarId,
    /// Cell path accesses that need replacement
    cell_path_replacements: Vec<CellPathReplacement>,
}

fn extract_transpose_column_names(call: &Call) -> Option<(String, String)> {
    if call.arguments.len() != 2 {
        return None;
    }

    let first_arg = call.arguments.first()?;
    let second_arg = call.arguments.get(1)?;

    let first_name = match first_arg {
        Argument::Positional(expr) | Argument::Unknown(expr) => {
            if let Expr::String(s) = &expr.expr {
                s.clone()
            } else {
                return None;
            }
        }
        _ => return None,
    };

    let second_name = match second_arg {
        Argument::Positional(expr) | Argument::Unknown(expr) => {
            if let Expr::String(s) = &expr.expr {
                s.clone()
            } else {
                return None;
            }
        }
        _ => return None,
    };

    Some((first_name, second_name))
}

fn closure_only_uses_fields(
    block_id: nu_protocol::BlockId,
    closure_var_id: nu_protocol::VarId,
    field1: &str,
    field2: &str,
    context: &LintContext,
) -> Option<Vec<CellPathReplacement>> {
    let block = context.working_set.get_block(block_id);

    let mut field_accesses = Vec::new();

    block.flat_map(
        context.working_set,
        &|expr| {
            if let Expr::FullCellPath(cell_path) = &expr.expr {
                if let Expr::Var(var_id) = &cell_path.head.expr
                    && *var_id == closure_var_id
                {
                    vec![expr.clone()]
                } else {
                    vec![]
                }
            } else {
                vec![]
            }
        },
        &mut field_accesses,
    );

    log::trace!("Found {} field accesses", field_accesses.len());

    if field_accesses.is_empty() {
        log::trace!("No field accesses found");
        return None;
    }

    let mut replacements = Vec::new();

    for expr in &field_accesses {
        if let Expr::FullCellPath(cell_path) = &expr.expr
            && cell_path.tail.len() == 1
        {
            let field_name = context.span_text(cell_path.tail[0].span());
            log::trace!("Field access: {field_name} (expecting {field1} or {field2})");
            if field_name == field1 || field_name == field2 {
                replacements.push(CellPathReplacement {
                    span: expr.span,
                    field: field_name.to_string(),
                });
            } else {
                log::trace!("Invalid field name");
                return None;
            }
        } else {
            log::trace!("Invalid field access pattern");
            return None;
        }
    }

    log::trace!("All {} field accesses valid", replacements.len());
    Some(replacements)
}

fn detect_pattern_in_pipeline(
    pipeline: &Pipeline,
    context: &LintContext,
) -> Vec<TransposeEachPattern> {
    pipeline
        .find_command_pairs(
            context,
            |call, ctx| call.is_call_to_command("transpose", ctx),
            |call, ctx| call.is_call_to_command("each", ctx),
        )
        .into_iter()
        .filter_map(|pair| {
            let (col1, col2) = extract_transpose_column_names(pair.first)?;

            log::trace!("Column names: {col1}, {col2}");

            let Some(Argument::Positional(closure_arg) | Argument::Unknown(closure_arg)) =
                pair.second.arguments.first()
            else {
                log::trace!("Each doesn't have positional argument");
                return None;
            };

            let Expr::Closure(block_id) = &closure_arg.expr else {
                log::trace!("Argument is not a closure");
                return None;
            };

            log::trace!("Closure found");

            let block = context.working_set.get_block(*block_id);

            if block.signature.required_positional.len() != 1 {
                log::trace!(
                    "Closure has {} parameters",
                    block.signature.required_positional.len()
                );
                return None;
            }

            log::trace!("Closure has 1 parameter");

            let param = &block.signature.required_positional[0];
            let Some(closure_var_id) = param.var_id else {
                log::trace!("Parameter doesn't have var_id");
                return None;
            };

            log::trace!("Checking field usage");

            let Some(cell_path_replacements) =
                closure_only_uses_fields(*block_id, closure_var_id, &col1, &col2, context)
            else {
                log::trace!("Closure doesn't only use the specified fields");
                return None;
            };

            log::trace!("Pattern matched!");

            Some(TransposeEachPattern {
                each_call: pair.second.clone(),
                col1,
                col2,
                combined_span: pair.span,
                transpose_span: pair.first.head,
                each_span: pair.second.head,
                closure_span: closure_arg.span,
                param_var_id: closure_var_id,
                cell_path_replacements,
            })
        })
        .collect()
}

fn collect_patterns_from_block(
    block_id: nu_protocol::BlockId,
    context: &LintContext,
) -> Vec<TransposeEachPattern> {
    let block = context.working_set.get_block(block_id);
    let mut patterns = Vec::new();

    for pipeline in &block.pipelines {
        patterns.extend(detect_pattern_in_pipeline(pipeline, context));
    }

    patterns
}

fn detect_pattern(expr: &Expression, context: &LintContext) -> Vec<TransposeEachPattern> {
    match &expr.expr {
        Expr::Block(block_id) | Expr::Closure(block_id) | Expr::Subexpression(block_id) => {
            collect_patterns_from_block(*block_id, context)
        }
        _ => vec![],
    }
}

struct ItemsInsteadOfTransposeEach;

impl DetectFix for ItemsInsteadOfTransposeEach {
    type FixInput<'a> = TransposeEachPattern;

    fn id(&self) -> &'static str {
        "transpose_items"
    }

    fn short_description(&self) -> &'static str {
        "Simplify `transpose | each` to `items`"
    }

    fn source_link(&self) -> Option<&'static str> {
        Some("https://www.nushell.sh/commands/docs/items.html")
    }

    fn level(&self) -> LintLevel {
        LintLevel::Hint
    }

    fn detect<'a>(&self, context: &'a LintContext) -> Vec<(Detection, Self::FixInput<'a>)> {
        let context: &LintContext = context;
        let mut patterns = Vec::new();

        for pipeline in &context.ast.pipelines {
            patterns.extend(detect_pattern_in_pipeline(pipeline, context));
        }

        context.ast.flat_map(
            context.working_set,
            &|expr| detect_pattern(expr, context),
            &mut patterns,
        );

        patterns
            .into_iter()
            .map(|pattern| {
                let violation = Detection::from_global_span(
                    "Use 'items' instead of 'transpose | each' when iterating over record entries",
                    pattern.combined_span,
                )
                .with_primary_label("transpose | each pattern")
                .with_extra_label("converts record to table", pattern.transpose_span)
                .with_extra_label("iterates over rows", pattern.each_span)
                .with_extra_label(
                    format!("accesses ${} and ${}", pattern.col1, pattern.col2),
                    pattern.closure_span,
                );
                (violation, pattern)
            })
            .collect()
    }

    fn fix(&self, context: &LintContext, pattern: &Self::FixInput<'_>) -> Option<Fix> {
        let pattern: &TransposeEachPattern = pattern;
        let closure_arg = pattern.each_call.arguments.first()?;
        let (Argument::Positional(closure_expr) | Argument::Unknown(closure_expr)) = closure_arg
        else {
            return None;
        };

        let Expr::Closure(_block_id) = &closure_expr.expr else {
            return None;
        };

        let mut replacements = Vec::new();

        // 1. Replace `transpose col1 col2 | each` with `items`
        // The span from transpose head to just before the closure
        let transpose_each_span = Span::new(pattern.transpose_span.start, closure_expr.span.start);
        replacements.push(Replacement::new(transpose_each_span, "items ".to_string()));

        // 2. Replace the closure parameter declaration `|row|` with `|col1, col2|`
        // Use AST-based span lookup via the variable's declaration_span
        let var = context.working_set.get_variable(pattern.param_var_id);
        let param_span = var.declaration_span;
        replacements.push(Replacement::new(
            param_span,
            format!("{}, {}", pattern.col1, pattern.col2),
        ));

        // 3. Replace each cell path access `$row.col1` with `$col1`, `$row.col2` with
        //    `$col2`
        for cell_path_replacement in &pattern.cell_path_replacements {
            replacements.push(Replacement::new(
                cell_path_replacement.span,
                format!("${}", cell_path_replacement.field),
            ));
        }

        Some(Fix {
            explanation: "Replace 'transpose ... | each' with 'items' for cleaner iteration over \
                          record entries"
                .into(),
            replacements,
        })
    }
}

pub static RULE: &dyn Rule = &ItemsInsteadOfTransposeEach;

#[cfg(test)]
mod detect_bad;
#[cfg(test)]
mod generated_fix;
#[cfg(test)]
mod ignore_good;
