copper/src/commands/argument.rs

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

#[derive(Debug, Clone)]
pub enum ArgumentType {
    String,
    Integer,
    Float,
    Boolean,
    Enum(&'static [&'static str]),
    GreedyString,
}

impl fmt::Display for ArgumentType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ArgumentType::String => write!(f, "String"),
            ArgumentType::Integer => write!(f, "Integer"),
            ArgumentType::Enum(variants) => write!(f, "Enum({:?})", variants),
            ArgumentType::Float => write!(f, "Float"),
            ArgumentType::Boolean => write!(f, "Boolean"),
            ArgumentType::GreedyString => write!(f, "GreedyString"),
        }
    }
}

#[derive(Debug, Clone)]
pub struct ArgumentSpec {
    pub name: &'static str,
    pub arg_type: ArgumentType,
    pub required: bool,
    pub default: Option<&'static str>,
    pub children: &'static [ArgumentSpec],
}

#[derive(Debug, Clone)]
pub enum ParsedArgument {
    String(String),
    Integer(i64),
    Float(f64),
    Boolean(bool),
    Enum(String),
    GreedyString(String),
}

impl fmt::Display for ParsedArgument {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ParsedArgument::String(s) => write!(f, "{}", s),
            ParsedArgument::Integer(i) => write!(f, "{}", i),
            ParsedArgument::Float(fl) => write!(f, "{}", fl),
            ParsedArgument::Boolean(b) => write!(f, "{}", b),
            ParsedArgument::Enum(e) => write!(f, "{}", e),
            ParsedArgument::GreedyString(s) => write!(f, "{}", s),
        }
    }
}

pub fn parse_argument(arg: &str, spec: &ArgumentSpec) -> Result<ParsedArgument, String> {
    match &spec.arg_type {
        ArgumentType::String => Ok(ParsedArgument::String(arg.to_string())),
        ArgumentType::Integer => arg
            .parse::<i64>()
            .map(ParsedArgument::Integer)
            .map_err(|_| format!("Argument '{}' must be an integer, got '{}'", spec.name, arg)),
        ArgumentType::Float => arg
            .parse::<f64>()
            .map(ParsedArgument::Float)
            .map_err(|_| format!("Argument '{}' must be a float, got '{}'", spec.name, arg)),
        ArgumentType::Boolean => match arg.to_lowercase().as_str() {
            "true" | "yes" | "1" => Ok(ParsedArgument::Boolean(true)),
            "false" | "no" | "0" => Ok(ParsedArgument::Boolean(false)),
            _ => Err(format!(
                "Argument '{}' must be a boolean, got '{}'",
                spec.name, arg
            )),
        },
        ArgumentType::Enum(variants) => {
            if variants.contains(&arg) {
                Ok(ParsedArgument::Enum(arg.to_string()))
            } else {
                Err(format!(
                    "Argument '{}' must be one of {:?}, got '{}'",
                    spec.name, variants, arg
                ))
            }
        }
        ArgumentType::GreedyString => Ok(ParsedArgument::GreedyString(arg.to_string())),
    }
}

#[derive(Debug, Clone)]
pub struct ParsedArguments {
    pub raw: Vec<String>,
    pub parsed: HashMap<&'static str, ParsedArgument>,
}

impl ParsedArguments {
    pub fn get(&self, name: &str) -> Option<&ParsedArgument> {
        self.parsed.get(name)
    }
}

fn token_matches_spec(token: &str, spec: &ArgumentSpec) -> bool {
    match &spec.arg_type {
        ArgumentType::String => true,
        ArgumentType::Integer => token.parse::<i64>().is_ok(),
        ArgumentType::Float => token.parse::<f64>().is_ok(),
        ArgumentType::Boolean => {
            matches!(
                token.to_lowercase().as_str(),
                "true" | "yes" | "1" | "false" | "no" | "0"
            )
        }
        ArgumentType::Enum(variants) => variants.contains(&token),
        ArgumentType::GreedyString => true,
    }
}

fn parse_spec_chain(
    specs: &[ArgumentSpec],
    raw_args: &[String],
    parsed: &mut HashMap<&'static str, ParsedArgument>,
    depth: usize,
) -> Result<usize, String> {
    if depth > 10 {
        return Err("Too many nested arguments (depth > 10)".into());
    }

    let mut consumed: usize = 0;

    for spec in specs {
        if let ArgumentType::GreedyString = spec.arg_type {
            if consumed >= raw_args.len() {
                if let Some(def) = spec.default {
                    let parsed_arg = parse_argument(def, spec)?;
                    parsed.insert(spec.name, parsed_arg);
                } else if spec.required {
                    return Err(format!("Missing required argument '{}'", spec.name));
                }
            } else {
                let rest = raw_args[consumed..].join(" ");
                let parsed_arg = parse_argument(rest.as_str(), spec)?;
                parsed.insert(spec.name, parsed_arg);
                consumed = raw_args.len();
            }
            break;
        }

        if consumed < raw_args.len() {
            let next = &raw_args[consumed];
            if token_matches_spec(next, spec) {
                let parsed_arg = parse_argument(next.as_str(), spec)?;
                parsed.insert(spec.name, parsed_arg);
                consumed += 1;

                if !spec.children.is_empty() {
                    let used =
                        parse_spec_chain(spec.children, &raw_args[consumed..], parsed, depth + 1)?;
                    consumed += used;
                }
            } else if let Some(def) = spec.default {
                let parsed_arg = parse_argument(def, spec)?;
                parsed.insert(spec.name, parsed_arg);

                if !spec.children.is_empty() {
                    let used =
                        parse_spec_chain(spec.children, &raw_args[consumed..], parsed, depth + 1)?;
                    consumed += used;
                }
            } else if spec.required {
                return Err(format!("Missing required argument '{}'", spec.name));
            }
        } else if let Some(def) = spec.default {
            let parsed_arg = parse_argument(def, spec)?;
            parsed.insert(spec.name, parsed_arg);

            if !spec.children.is_empty() {
                let used =
                    parse_spec_chain(spec.children, &raw_args[consumed..], parsed, depth + 1)?;
                consumed += used;
            }
        } else if spec.required {
            return Err(format!("Missing required argument '{}'", spec.name));
        }
    }

    Ok(consumed)
}

pub fn parse_arguments(
    specs: &[ArgumentSpec],
    raw_args: &[String],
) -> Result<ParsedArguments, String> {
    let mut parsed = HashMap::new();
    let _ = parse_spec_chain(specs, raw_args, &mut parsed, 0)?;

    Ok(ParsedArguments {
        raw: raw_args.to_vec(),
        parsed,
    })
}