use std::collections::HashMap;
use std::env;
use std::fmt;

#[derive(Debug, Clone)]
struct OptionSpec {
    name: &'static str,
    short: Option<char>,
    long: Option<&'static str>,
    takes_value: bool,
    help: &'static str,
}

#[derive(Debug, Default, PartialEq, Eq)]
struct ParsedArgs {
    flags: HashMap<String, bool>,
    values: HashMap<String, Vec<String>>,
    positionals: Vec<String>,
}

impl ParsedArgs {
    fn flag(&self, name: &str) -> bool {
        self.flags.get(name).copied().unwrap_or(false)
    }

    fn value(&self, name: &str) -> Option<&str> {
        self.values.get(name).and_then(|v| v.last()).map(String::as_str)
    }

    fn values(&self, name: &str) -> &[String] {
        self.values.get(name).map(Vec::as_slice).unwrap_or(&[])
    }
}

#[derive(Debug, PartialEq, Eq)]
enum ParseError {
    UnknownOption(String),
    MissingValue(String),
    UnexpectedValue(String),
    InvalidShortOption(String),
}

impl fmt::Display for ParseError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ParseError::UnknownOption(opt) => write!(f, "unknown option: {opt}"),
            ParseError::MissingValue(opt) => write!(f, "missing value for option: {opt}"),
            ParseError::UnexpectedValue(opt) => write!(f, "option does not take a value: {opt}"),
            ParseError::InvalidShortOption(opt) => write!(f, "invalid short option: {opt}"),
        }
    }
}

struct Parser<'a> {
    specs: &'a [OptionSpec],
}

impl<'a> Parser<'a> {
    fn new(specs: &'a [OptionSpec]) -> Self {
        Self { specs }
    }

    fn parse<I, S>(&self, args: I) -> Result<ParsedArgs, ParseError>
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        let mut parsed = ParsedArgs::default();
        let mut args = args.into_iter().map(Into::into).peekable();
        let mut positionals_only = false;

        while let Some(arg) = args.next() {
            if positionals_only {
                parsed.positionals.push(arg);
                continue;
            }

            if arg == "--" {
                positionals_only = true;
            } else if let Some(rest) = arg.strip_prefix("--") {
                self.parse_long(rest, &mut args, &mut parsed)?;
            } else if arg.starts_with('-') && arg != "-" {
                self.parse_short_cluster(&arg, &mut args, &mut parsed)?;
            } else {
                parsed.positionals.push(arg);
            }
        }

        Ok(parsed)
    }

    fn parse_long<I>(
        &self,
        raw: &str,
        args: &mut std::iter::Peekable<I>,
        parsed: &mut ParsedArgs,
    ) -> Result<(), ParseError>
    where
        I: Iterator<Item = String>,
    {
        let (long, inline_value) = match raw.split_once('=') {
            Some((name, value)) => (name, Some(value.to_string())),
            None => (raw, None),
        };

        let spec = self
            .specs
            .iter()
            .find(|spec| spec.long == Some(long))
            .ok_or_else(|| ParseError::UnknownOption(format!("--{long}")))?;

        if spec.takes_value {
            let value = match inline_value {
                Some(value) => value,
                None => args
                    .next()
                    .ok_or_else(|| ParseError::MissingValue(format!("--{long}")))?,
            };
            parsed.values.entry(spec.name.to_string()).or_default().push(value);
        } else {
            if inline_value.is_some() {
                return Err(ParseError::UnexpectedValue(format!("--{long}")));
            }
            parsed.flags.insert(spec.name.to_string(), true);
        }

        Ok(())
    }

    fn parse_short_cluster<I>(
        &self,
        arg: &str,
        args: &mut std::iter::Peekable<I>,
        parsed: &mut ParsedArgs,
    ) -> Result<(), ParseError>
    where
        I: Iterator<Item = String>,
    {
        let mut chars = arg
            .strip_prefix('-')
            .ok_or_else(|| ParseError::InvalidShortOption(arg.to_string()))?
            .chars()
            .peekable();

        while let Some(short) = chars.next() {
            let spec = self
                .specs
                .iter()
                .find(|spec| spec.short == Some(short))
                .ok_or_else(|| ParseError::UnknownOption(format!("-{short}")))?;

            if spec.takes_value {
                // A value-taking short option consumes the rest of the cluster, or the next arg.
                let rest: String = chars.collect();
                let value = if rest.is_empty() {
                    args.next()
                        .ok_or_else(|| ParseError::MissingValue(format!("-{short}")))?
                } else {
                    rest
                };
                parsed.values.entry(spec.name.to_string()).or_default().push(value);
                break;
            } else {
                parsed.flags.insert(spec.name.to_string(), true);
            }
        }

        Ok(())
    }

    fn usage(&self, program: &str) -> String {
        let mut out = format!("Usage: {program} [options] [--] [args...]\n\nOptions:\n");

        for spec in self.specs {
            let short = spec.short.map(|c| format!("-{c}")).unwrap_or_default();
            let long = spec.long.map(|s| format!("--{s}")).unwrap_or_default();
            let value = if spec.takes_value { " <value>" } else { "" };

            let names = match (short.is_empty(), long.is_empty()) {
                (false, false) => format!("{short}, {long}{value}"),
                (false, true) => format!("{short}{value}"),
                (true, false) => format!("{long}{value}"),
                (true, true) => spec.name.to_string(),
            };

            out.push_str(&format!("  {:24} {}\n", names, spec.help));
        }

        out
    }
}

fn main() {
    let specs = [
        OptionSpec {
            name: "help",
            short: Some('h'),
            long: Some("help"),
            takes_value: false,
            help: "show this help text",
        },
        OptionSpec {
            name: "verbose",
            short: Some('v'),
            long: Some("verbose"),
            takes_value: false,
            help: "enable verbose output",
        },
        OptionSpec {
            name: "output",
            short: Some('o'),
            long: Some("output"),
            takes_value: true,
            help: "write output to a file",
        },
        OptionSpec {
            name: "define",
            short: Some('D'),
            long: Some("define"),
            takes_value: true,
            help: "set a key=value definition; may be repeated",
        },
    ];

    let parser = Parser::new(&specs);
    let program = env::args().next().unwrap_or_else(|| "app".to_string());

    match parser.parse(env::args().skip(1)) {
        Ok(args) if args.flag("help") => print!("{}", parser.usage(&program)),
        Ok(args) => {
            println!("verbose: {}", args.flag("verbose"));
            println!("output: {:?}", args.value("output"));
            println!("defines: {:?}", args.values("define"));
            println!("positionals: {:?}", args.positionals);
        }
        Err(err) => {
            eprintln!("{err}\n");
            eprintln!("{}", parser.usage(&program));
            std::process::exit(2);
        }
    }
}

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

    fn specs() -> Vec<OptionSpec> {
        vec![
            OptionSpec {
                name: "verbose",
                short: Some('v'),
                long: Some("verbose"),
                takes_value: false,
                help: "verbose mode",
            },
            OptionSpec {
                name: "output",
                short: Some('o'),
                long: Some("output"),
                takes_value: true,
                help: "output path",
            },
        ]
    }

    #[test]
    fn parses_long_short_and_positionals() {
        let specs = specs();
        let parser = Parser::new(&specs);

        let args = parser
            .parse(["-v", "--output=target.txt", "input.txt"])
            .unwrap();

        assert!(args.flag("verbose"));
        assert_eq!(args.value("output"), Some("target.txt"));
        assert_eq!(args.positionals, vec!["input.txt"]);
    }

    #[test]
    fn parses_short_value_and_double_dash() {
        let specs = specs();
        let parser = Parser::new(&specs);

        let args = parser.parse(["-oout.txt", "--", "--not-an-option"]).unwrap();

        assert_eq!(args.value("output"), Some("out.txt"));
        assert_eq!(args.positionals, vec!["--not-an-option"]);
    }

    #[test]
    fn reports_missing_value() {
        let specs = specs();
        let parser = Parser::new(&specs);

        assert_eq!(
            parser.parse(["--output"]),
            Err(ParseError::MissingValue("--output".to_string()))
        );
    }
}