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

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

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

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

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

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

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

impl std::error::Error for ParseError {}

#[derive(Debug, Clone)]
struct ArgSpec {
    long: String,
    takes_value: bool,
}

/// A small parser for long options (`--output file`), short options (`-o file`),
/// grouped flags (`-vvq`), and positional arguments.
#[derive(Debug, Default)]
pub struct ArgParser {
    specs: HashMap<String, ArgSpec>,
    shorts: HashMap<char, String>,
}

impl ArgParser {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn flag(mut self, long: &str, short: Option<char>) -> Self {
        self.add_spec(long, short, false);
        self
    }

    pub fn option(mut self, long: &str, short: Option<char>) -> Self {
        self.add_spec(long, short, true);
        self
    }

    fn add_spec(&mut self, long: &str, short: Option<char>, takes_value: bool) {
        let spec = ArgSpec {
            long: long.to_string(),
            takes_value,
        };

        self.specs.insert(long.to_string(), spec);

        if let Some(short) = short {
            self.shorts.insert(short, long.to_string());
        }
    }

    pub fn parse_env(&self) -> Result<ParsedArgs, ParseError> {
        self.parse(env::args().skip(1))
    }

    pub 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 positional_only = false;

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

            if arg == "--" {
                positional_only = true;
            } else if let Some(rest) = arg.strip_prefix("--") {
                self.parse_long(rest, &mut args, &mut parsed)?;
            } else if arg.starts_with('-') && arg.len() > 1 {
                self.parse_short_group(&arg[1..], &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 (name, inline_value) = match raw.split_once('=') {
            Some((name, value)) => (name, Some(value.to_string())),
            None => (raw, None),
        };

        let spec = self
            .specs
            .get(name)
            .ok_or_else(|| ParseError::UnknownOption(format!("--{name}")))?;

        self.store_value(spec, inline_value, args, parsed)
    }

    fn parse_short_group<I>(
        &self,
        group: &str,
        args: &mut std::iter::Peekable<I>,
        parsed: &mut ParsedArgs,
    ) -> Result<(), ParseError>
    where
        I: Iterator<Item = String>,
    {
        let mut chars = group.char_indices().peekable();

        while let Some((index, short)) = chars.next() {
            let long = self
                .shorts
                .get(&short)
                .ok_or_else(|| ParseError::UnknownOption(format!("-{short}")))?;
            let spec = self.specs.get(long).expect("short aliases point to specs");

            if spec.takes_value {
                let value_start = index + short.len_utf8();
                let rest = &group[value_start..];
                let inline_value = (!rest.is_empty()).then(|| rest.to_string());
                return self.store_value(spec, inline_value, args, parsed);
            }

            parsed.flags.insert(spec.long.clone());
        }

        Ok(())
    }

    fn store_value<I>(
        &self,
        spec: &ArgSpec,
        inline_value: Option<String>,
        args: &mut std::iter::Peekable<I>,
        parsed: &mut ParsedArgs,
    ) -> Result<(), ParseError>
    where
        I: Iterator<Item = String>,
    {
        if spec.takes_value {
            let value = match inline_value {
                Some(value) => value,
                None => args
                    .next()
                    .ok_or_else(|| ParseError::MissingValue(format!("--{}", spec.long)))?,
            };

            parsed.options.insert(spec.long.clone(), value);
        } else if inline_value.is_some() {
            return Err(ParseError::UnexpectedValue(format!("--{}", spec.long)));
        } else {
            parsed.flags.insert(spec.long.clone());
        }

        Ok(())
    }
}

fn main() {
    let parser = ArgParser::new()
        .flag("help", Some('h'))
        .flag("verbose", Some('v'))
        .option("output", Some('o'));

    match parser.parse_env() {
        Ok(args) => {
            if args.has_flag("help") {
                println!("usage: cli [-h] [-v] [-o FILE] [ARGS...]");
                return;
            }

            println!("{args:#?}");
        }
        Err(error) => {
            eprintln!("error: {error}");
            std::process::exit(2);
        }
    }
}

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

    fn parser() -> ArgParser {
        ArgParser::new()
            .flag("verbose", Some('v'))
            .flag("quiet", Some('q'))
            .option("output", Some('o'))
    }

    #[test]
    fn parses_long_and_short_options() {
        let parsed = parser()
            .parse(["-vq", "--output=log.txt", "input.txt"])
            .unwrap();

        assert!(parsed.has_flag("verbose"));
        assert!(parsed.has_flag("quiet"));
        assert_eq!(parsed.option("output"), Some("log.txt"));
        assert_eq!(parsed.positionals(), &["input.txt".to_string()]);
    }

    #[test]
    fn option_value_can_follow_short_name() {
        let parsed = parser().parse(["-o", "out.txt"]).unwrap();

        assert_eq!(parsed.option("output"), Some("out.txt"));
    }

    #[test]
    fn double_dash_stops_option_parsing() {
        let parsed = parser().parse(["--", "--output", "file"]).unwrap();

        assert_eq!(
            parsed.positionals(),
            &["--output".to_string(), "file".to_string()]
        );
    }
}