only-engine 0.3.2

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

use only_semantic::{DirectiveAst, DocumentAst, ShellKind, TaskAst, TaskStepAst};

use crate::planner::PlanError;
use crate::probe::probe_matches;
use crate::{ExecutionNode, ExecutionStep, PlanParam};

pub(crate) type TaskIndex<'a> = HashMap<String, Vec<&'a TaskAst>>;

pub(crate) fn build_task_index(document: &DocumentAst) -> TaskIndex<'_> {
    let mut tasks = HashMap::<String, Vec<&TaskAst>>::new();
    for task in &document.tasks {
        tasks
            .entry(task.qualified_name().to_string())
            .or_default()
            .push(task);
    }
    tasks
}

pub(crate) fn resolve_root_task<'a>(
    tasks: &'a TaskIndex<'a>,
    target: &str,
) -> Result<&'a TaskAst, PlanError> {
    let Some(root_variants) = tasks.get(target) else {
        return Err(PlanError::UnknownTask(target.to_string()));
    };
    select_root_task_from_variants(root_variants, target)
}

pub(crate) fn resolve_root_task_in_document<'a>(
    document: &'a DocumentAst,
    target: &str,
) -> Result<&'a TaskAst, PlanError> {
    let root_variants = document
        .tasks
        .iter()
        .filter(|task| task.qualified_name() == target)
        .collect::<Vec<_>>();

    if root_variants.is_empty() {
        return Err(PlanError::UnknownTask(target.to_string()));
    }

    select_root_task_from_variants(&root_variants, target)
}

fn select_root_task_from_variants<'a>(
    root_variants: &[&'a TaskAst],
    target: &str,
) -> Result<&'a TaskAst, PlanError> {
    let Some(root) = select_task_variant(root_variants) else {
        return Err(PlanError::TaskUnavailable(target.to_string()));
    };
    if root.is_helper() {
        return Err(PlanError::HelperTask(target.to_string()));
    }
    Ok(root)
}

pub(crate) fn build_execution_nodes(
    ordered: Vec<(usize, &TaskAst, HashMap<String, String>)>,
) -> Vec<ExecutionNode> {
    ordered
        .into_iter()
        .map(|(stage, task, bindings)| ExecutionNode {
            stage,
            name: task.qualified_name().to_string(),
            steps: task
                .steps
                .iter()
                .map(|step| match step {
                    TaskStepAst::Command(command) => {
                        ExecutionStep::Command(command.text.to_string())
                    }
                    TaskStepAst::CommandBlock(block) => ExecutionStep::CommandBlock {
                        source: block.source.to_string(),
                        line_count: block.line_ranges.len(),
                    },
                })
                .collect(),
            params: build_plan_params(task, &bindings),
            shell: task.shell.as_ref().map(|shell| shell.selection.clone()),
        })
        .collect()
}

fn build_plan_params(task: &TaskAst, bindings: &HashMap<String, String>) -> Vec<PlanParam> {
    let mut params = task
        .params
        .iter()
        .map(|param| PlanParam {
            name: param.name.to_string(),
            default_value: param.default_value.as_ref().map(ToString::to_string),
            value: bindings.get(param.name.as_str()).cloned(),
        })
        .collect::<Vec<_>>();
    let mut globals = bindings
        .iter()
        .filter(|(name, _)| {
            !task
                .params
                .iter()
                .any(|param| param.name.as_str() == name.as_str())
        })
        .map(|(name, value)| PlanParam {
            name: name.clone(),
            default_value: Some(value.clone()),
            value: Some(value.clone()),
        })
        .collect::<Vec<_>>();
    globals.sort_by(|left, right| left.name.cmp(&right.name));
    params.extend(globals);
    params
}

pub(crate) fn select_task_variant<'a>(variants: &[&'a TaskAst]) -> Option<&'a TaskAst> {
    let mut fallback = None;

    for task in variants {
        if task.guards.is_empty() {
            fallback = Some(*task);
        } else if task
            .guards
            .iter()
            .all(|guard| probe_matches(&guard.kind, guard.argument.as_str()))
        {
            return Some(task);
        }
    }

    fallback
}

pub(crate) fn document_shell(document: &DocumentAst) -> Option<ShellKind> {
    document
        .directives
        .iter()
        .find_map(|directive| match directive {
            DirectiveAst::Version { .. } => None,
            DirectiveAst::Shell { shell, .. } => Some(shell.clone()),
            DirectiveAst::Variable { .. } => None,
        })
}

pub(crate) fn bind_parameters(
    task: &TaskAst,
    inputs: Option<&HashMap<String, String>>,
    globals: &HashMap<String, String>,
) -> Result<HashMap<String, String>, PlanError> {
    let input_map = inputs.cloned().unwrap_or_default();
    let mut parameters = globals.clone();

    for param in &task.params {
        if let Some(value) = input_map.get(param.name.as_str()) {
            parameters.insert(param.name.to_string(), value.clone());
            continue;
        }

        if param.is_slice {
            parameters.insert(param.name.to_string(), String::new());
            continue;
        }

        if let Some(default) = &param.default_value {
            parameters.insert(param.name.to_string(), default.to_string());
            continue;
        }

        return Err(PlanError::MissingRequiredParameter(param.name.to_string()));
    }

    Ok(parameters)
}

pub(crate) fn merge_parameter_inputs(
    positional_args: Vec<&str>,
    named_overrides: Vec<(&str, &str)>,
    task: &TaskAst,
    globals: &HashMap<String, String>,
) -> Result<HashMap<String, String>, PlanError> {
    let slice_index = task
        .params
        .iter()
        .enumerate()
        .find(|(index, param)| {
            param.is_slice && param.default_value.is_none() && index + 1 == task.params.len()
        })
        .map(|(index, _)| index);

    if slice_index.is_none() && positional_args.len() > task.params.len() {
        return Err(PlanError::TooManyArguments {
            task: task.qualified_name().to_string(),
            expected: task.params.len(),
            got: positional_args.len(),
        });
    }

    let allowed = task
        .params
        .iter()
        .map(|param| param.name.as_str())
        .collect::<HashSet<_>>();

    let mut merged = HashMap::new();
    for (index, value) in positional_args.iter().copied().enumerate() {
        if Some(index) == slice_index {
            let joined = positional_args[index..].join(" ");
            merged.insert(task.params[index].name.to_string(), joined);
            break;
        }

        let parameter = &task.params[index];
        merged.insert(parameter.name.to_string(), value.to_string());
    }

    for (name, value) in globals {
        if !allowed.contains(name.as_str()) {
            merged.insert(name.clone(), value.clone());
        }
    }

    let mut seen_overrides = HashSet::new();
    for (name, value) in named_overrides {
        if !allowed.contains(name) && !globals.contains_key(name) {
            return Err(PlanError::UnknownParameter {
                task: task.qualified_name().to_string(),
                name: name.to_string(),
            });
        }

        if !seen_overrides.insert(name) {
            return Err(PlanError::DuplicateOverride(name.to_string()));
        }
        merged.insert(name.to_string(), value.to_string());
    }

    Ok(merged)
}