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

/// Parsed command-line arguments: named options plus positional values.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Args {
    flags: HashSet<String>,
    options: HashMap<String, String>,
    positionals: Vec<String>,
}

impl Args {
    pub fn has(&self, name: &str) -> bool {
        self.flags.contains(name)
    }

    pub fn get(&self, name: &str) -> Option<&str> {
        self.options.get(name).map(String::as_str)
    }

    pub fn positionals(&self) -> &[String] {
        &self.positionals
    }
}

/// Declares which options require a value.
#[derive(Debug, Clone)]
pub struct Parser {
    value_options: HashSet<String>,
}

impl Parser {
    pub fn new() -> Self {
        Self {
            value_options: HashSet::new(),
        }
    }

    pub fn value_option(mut self, name: impl Into<String>) -> Self {
        self.value_options.insert(name.into());
        self
    }

    pub fn parse<I, S>(&self, input: I) -> Result<Args, ParseError>
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        let mut flags = HashSet::new();
        let mut options = HashMap::new();
        let mut positionals = Vec::new();

        let mut args = input.into_iter().map(Into::into).peekable();

        while let Some(arg) = args.next() {
            if arg == "--" {
                positionals.extend(args);
                break;
            }

            if let Some(rest) = arg.strip_prefix("--") {
                if rest.is_empty() {
                    positionals.push(arg);
                    continue;
                }

                if let Some((name, value)) = rest.split_once('=') {
                    if name.is_empty() {
                        return Err(ParseError::InvalidOption(arg));
                    }
                    options.insert(name.to_string(), value.to_string());
                } else if self.value_options.contains(rest) {
                    let value = args
                        .next()
                        .ok_or_else(|| ParseError::MissingValue(rest.to_string()))?;
                    options.insert(rest.to_string(), value);
                } else {
                    flags.insert(rest.to_string());
                }

                continue;
            }

            if let Some(rest) = arg.strip_prefix('-') {
                if rest.is_empty() {
                    positionals.push(arg);
                    continue;
                }

                // A single short option may take a value if registered, e.g. `-o out.txt`.
                if rest.chars().count() == 1 && self.value_options.contains(rest) {
                    let value = args
                        .next()
                        .ok_or_else(|| ParseError::MissingValue(rest.to_string()))?;
                    options.insert(rest.to_string(), value);
                } else {
                    // Grouped short flags are supported, e.g. `-abc`.
                    for ch in rest.chars() {
                        flags.insert(ch.to_string());
                    }
                }

                continue;
            }

            positionals.push(arg);
        }

        Ok(Args {
            flags,
            options,
            positionals,
        })
    }
}

impl Default for Parser {
    fn default() -> Self {
        Self::new()
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ParseError {
    MissingValue(String),
    InvalidOption(String),
}

fn main() {
    let parser = Parser::new()
        .value_option("output")
        .value_option("o")
        .value_option("config");

    match parser.parse(env::args().skip(1)) {
        Ok(args) => {
            println!("{args:#?}");
        }
        Err(err) => {
            eprintln!("error: {err:?}");
            std::process::exit(2);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parses_flags_options_and_positionals() {
        let parser = Parser::new().value_option("output").value_option("o");

        let args = parser
            .parse(["--verbose", "--output", "app.log", "-xz", "-o", "short.log", "run"])
            .unwrap();

        assert!(args.has("verbose"));
        assert!(args.has("x"));
        assert!(args.has("z"));
        assert_eq!(args.get("output"), Some("app.log"));
        assert_eq!(args.get("o"), Some("short.log"));
        assert_eq!(args.positionals(), &["run"]);
    }

    #[test]
    fn respects_double_dash_separator() {
        let parser = Parser::new();

        let args = parser.parse(["--dry-run", "--", "--not-a-flag", "file"]).unwrap();

        assert!(args.has("dry-run"));
        assert_eq!(args.positionals(), &["--not-a-flag", "file"]);
    }

    #[test]
    fn reports_missing_value() {
        let parser = Parser::new().value_option("config");

        let err = parser.parse(["--config"]).unwrap_err();

        assert_eq!(err, ParseError::MissingValue("config".to_string()));
    }
}