use std::borrow::Cow;
use std::collections::HashMap;
use std::error;
use std::fmt;
use std::slice::from_ref;

pub trait ArgHandler {
    type Error: error::Error + 'static;

    fn on_literal(&mut self, name: &str) -> Result<(), Self::Error>;

    fn on_short(&mut self, name: char, value: Option<&str>) -> Result<(), Self::Error>;

    fn on_long(&mut self, name: &str, value: Option<&str>) -> Result<(), Self::Error>;
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Error<E: error::Error + 'static> {
    Handler { pos: usize, val: E },
    EmptyName { pos: usize },
}

impl<E: error::Error + 'static> fmt::Display for Error<E> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Handler { pos, val } => write!(f, "{val} (at #{pos})"),
            Self::EmptyName { pos } => write!(f, "malformed argument (at #{pos})"),
        }
    }
}

impl<E: error::Error + 'static> error::Error for Error<E> {
    fn source(&self) -> Option<&(dyn error::Error + 'static)> {
        match self {
            // forward past the inner error because we decorate it in our Display impl
            Self::Handler { val, .. } => val.source(),
            _ => None,
        }
    }
}

pub fn parse<I: Iterator, H: ArgHandler>(
    args: &mut I,
    handler: &mut H,
    arg_off: usize,
) -> Result<bool, Error<H::Error>>
where
    I::Item: AsRef<str>,
{
    for (pos, arg) in args.enumerate() {
        let arg = arg.as_ref();
        if !arg.is_empty() {
            if arg.as_bytes()[0] == b'-' {
                if arg.len() >= 2 && arg.as_bytes()[1] == b'-' {
                    if arg == "--" {
                        return Ok(false);
                    }
                    let (name, value) = match arg.bytes().enumerate().find(|(_, b)| *b == b'=') {
                        None => (&arg[2..], None),
                        Some((i, _)) => (&arg[2..i], Some(&arg[i + 1..])),
                    };
                    if name.is_empty() {
                        return Err(Error::EmptyName { pos: arg_off + pos });
                    }
                    if let Err(val) = handler.on_long(name, value) {
                        return Err(Error::Handler {
                            pos: arg_off + pos,
                            val,
                        });
                    }
                } else {
                    let (value, end) = match arg.bytes().enumerate().find(|(_, b)| *b == b'=') {
                        None => (None, arg.len()),
                        Some((i, _)) => (Some(&arg[i + 1..]), i),
                    };
                    if end > 1 {
                        for c in arg[1..end].chars() {
                            if let Err(val) = handler.on_short(c, value) {
                                return Err(Error::Handler {
                                    pos: arg_off + pos,
                                    val,
                                });
                            }
                        }
                    } else {
                        return Err(Error::EmptyName { pos: arg_off + pos });
                    }
                }
            } else {
                if let Err(val) = handler.on_literal(arg) {
                    return Err(Error::Handler {
                        pos: arg_off + pos,
                        val,
                    });
                }
            }
        }
    }
    Ok(true)
}

pub fn parse_args<H: ArgHandler>(handler: &mut H) -> Result<(), Error<H::Error>> {
    parse(&mut std::env::args(), handler, 0)?;
    Ok(())
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ArgCount {
    Forbidden,
    Optional(usize),
    Required(usize),
}

impl ArgCount {
    pub const fn has_value(&self) -> bool {
        match self {
            Self::Optional(..) | ArgCount::Required(..) => true,
            _ => false,
        }
    }

    pub const fn is_required(&self) -> bool {
        match self {
            Self::Required(..) => true,
            _ => false,
        }
    }

    pub const fn get_max_count(&self) -> Option<usize> {
        match self {
            Self::Optional(max) | ArgCount::Required(max) => Some(*max),
            _ => None,
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ArgOption {
    short: Option<char>,
    long: Option<Cow<'static, str>>,
    count: ArgCount,
}

impl ArgOption {
    pub const fn new(
        short: Option<char>,
        long: Option<Cow<'static, str>>,
        count: ArgCount,
    ) -> Self {
        if short.is_none() && long.is_none() {
            panic!("option must have at least a short or long name");
        }
        if let Some(max) = count.get_max_count() {
            if max == 0 {
                panic!("argument must be allowed to appear at least once");
            }
        }
        Self { short, long, count }
    }

    pub fn get_short(&self) -> Option<char> {
        self.short
    }

    pub fn get_long(&self) -> Option<&str> {
        match self.long {
            None => None,
            Some(Cow::Borrowed(r)) => Some(r),
            Some(Cow::Owned(ref s)) => Some(s.as_str()),
        }
    }

    pub const fn get_count(&self) -> &ArgCount {
        &self.count
    }
}

impl fmt::Display for ArgOption {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match (self.get_short(), self.get_long()) {
            (None, None) => unreachable!("unnamed ArgOption"),
            (None, Some(long)) => write!(f, "\"--{long}\""),
            (Some(short), None) => write!(f, "\"-{short}\""),
            (Some(short), Some(long)) => write!(f, "\"--{long}\" / \"-{short}\""),
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum OptionValue {
    Absent,
    Present,
    Value(String),
    Values(Vec<String>),
}

impl OptionValue {
    pub const fn is_absent(&self) -> bool {
        match self {
            Self::Absent => true,
            _ => false,
        }
    }

    pub const fn is_present(&self) -> bool {
        match self {
            Self::Present | Self::Value(..) | Self::Values(..) => true,
            _ => false,
        }
    }

    pub const fn has_value(&self) -> bool {
        match self {
            Self::Value(..) => true,
            Self::Values(..) => true,
            _ => false,
        }
    }

    pub const fn get_value(&self) -> Option<&String> {
        match self {
            Self::Value(v) => Some(v),
            _ => None,
        }
    }

    pub fn get_values(&self) -> Option<&[String]> {
        match self {
            Self::Value(v) => Some(from_ref(v)),
            Self::Values(v) => Some(v.as_ref()),
            _ => None,
        }
    }
}

#[derive(Clone, Debug)]
pub struct OptionHandler {
    options: Vec<(ArgOption, OptionValue)>,
    short_map: HashMap<char, usize>,
    long_map: HashMap<String, usize>,
    literals: Vec<String>,
}

impl OptionHandler {
    pub fn new() -> Self {
        Self {
            options: Vec::new(),
            short_map: HashMap::new(),
            long_map: HashMap::new(),
            literals: Vec::new(),
        }
    }

    pub fn add(&mut self, opt: ArgOption) -> Result<OptionRef, AddArgError> {
        match opt.short {
            Some(c) => match self.short_map.get(&c) {
                Some(&i) => {
                    return Err(AddArgError {
                        to_add: opt,
                        existing: &self.options[i].0,
                    })
                }
                _ => (),
            },
            _ => (),
        }
        match opt.long {
            Some(ref s) => match self.long_map.get(&**s) {
                Some(&i) => {
                    return Err(AddArgError {
                        to_add: opt,
                        existing: &self.options[i].0,
                    })
                }
                _ => (),
            },
            _ => (),
        }

        let idx = self.options.len();
        self.options.push((opt, OptionValue::Absent));
        let opt = &self.options[idx].0;
        if let Some(c) = opt.short {
            self.short_map.insert(c, idx);
        }
        if let Some(ref s) = opt.long {
            let k = &**s;
            self.long_map.insert(k.to_owned(), idx);
        }
        Ok(OptionRef(idx))
    }

    pub fn options(&self) -> &Vec<(ArgOption, OptionValue)> {
        &self.options
    }

    pub fn get(&self, opt_ref: OptionRef) -> (&ArgOption, &OptionValue) {
        let opt = &self.options[opt_ref.0];
        (&opt.0, &opt.1)
    }

    pub fn get_option(&self, opt_ref: OptionRef) -> &ArgOption {
        &self.options[opt_ref.0].0
    }

    pub fn get_value(&self, opt_ref: OptionRef) -> &OptionValue {
        &self.options[opt_ref.0].1
    }

    pub fn get_short(&self, name: char) -> Option<&(ArgOption, OptionValue)> {
        self.short_map.get(&name).map(|&i| &self.options[i])
    }

    pub fn get_long(&self, name: &str) -> Option<&(ArgOption, OptionValue)> {
        self.long_map.get(name).map(|&i| &self.options[i])
    }

    pub fn get_literals(&self) -> &Vec<String> {
        &self.literals
    }

    fn set_arg(&mut self, idx: usize, value: Option<&str>) -> Result<(), OptionError> {
        let (ref o, ref mut curr) = self.options[idx];
        match o.count {
            ArgCount::Forbidden => {
                if let None = value {
                    if curr.is_absent() {
                        *curr = OptionValue::Present;
                    }
                    Ok(())
                } else {
                    Err(OptionError::ValueForbidden(o.clone()))
                }
            }
            ArgCount::Optional(max) => match curr {
                OptionValue::Absent | OptionValue::Present => {
                    if let Some(v) = value {
                        if max == 1 {
                            *curr = OptionValue::Value(v.to_owned());
                        } else {
                            *curr = OptionValue::Values(vec![v.to_owned()]);
                        }
                    } else {
                        *curr = OptionValue::Present;
                    }
                    Ok(())
                }
                OptionValue::Value(..) => Err(OptionError::TooMany(o.clone())),
                OptionValue::Values(vec) => {
                    if vec.len() <= max {
                        if let Some(v) = value {
                            vec.push(v.to_owned());
                        }
                        Ok(())
                    } else {
                        Err(OptionError::TooMany(o.clone()))
                    }
                }
            },
            ArgCount::Required(max) => {
                if let Some(v) = value {
                    match curr {
                        OptionValue::Absent => {
                            if max == 1 {
                                *curr = OptionValue::Value(v.to_owned());
                            } else {
                                *curr = OptionValue::Values(vec![v.to_owned()]);
                            }
                            Ok(())
                        }
                        OptionValue::Present => unreachable!("argument missing required value"),
                        OptionValue::Value(..) => Err(OptionError::TooMany(o.clone())),
                        OptionValue::Values(vec) => {
                            if vec.len() <= max {
                                vec.push(v.to_owned());
                                Ok(())
                            } else {
                                Err(OptionError::TooMany(o.clone()))
                            }
                        }
                    }
                } else {
                    Err(OptionError::ValueRequired(o.clone()))
                }
            }
        }
    }

    pub fn clear(&mut self) {
        self.options
            .iter_mut()
            .for_each(|(_, v)| *v = OptionValue::Absent);
    }
}

#[derive(Clone, Copy, Debug)]
pub struct OptionRef(usize);

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct AddArgError<'l> {
    pub to_add: ArgOption,
    pub existing: &'l ArgOption,
}

impl<'l> fmt::Display for AddArgError<'l> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "duplicate argument {} (already have {})",
            self.to_add, self.existing
        )
    }
}

impl<'l> error::Error for AddArgError<'l> {}

impl ArgHandler for OptionHandler {
    type Error = OptionError;

    fn on_literal(&mut self, name: &str) -> Result<(), Self::Error> {
        self.literals.push(name.to_owned());
        Ok(())
    }

    fn on_short(&mut self, name: char, value: Option<&str>) -> Result<(), Self::Error> {
        match self.short_map.get(&name) {
            None => Err(OptionError::NoSuchShort(name)),
            Some(&i) => self.set_arg(i, value),
        }
    }

    fn on_long(&mut self, name: &str, value: Option<&str>) -> Result<(), Self::Error> {
        match self.long_map.get(name) {
            None => Err(OptionError::NoSuchLong(name.to_owned())),
            Some(&i) => self.set_arg(i, value),
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum OptionError {
    NoSuchShort(char),
    NoSuchLong(String),
    ValueForbidden(ArgOption),
    ValueRequired(ArgOption),
    TooMany(ArgOption),
}

impl fmt::Display for OptionError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::NoSuchShort(short) => write!(f, "invalid argument \"-{short}\""),
            Self::NoSuchLong(long) => write!(f, "invalid argument \"--{long}\""),
            Self::ValueForbidden(opt) => write!(f, "argument {opt} has no value"),
            Self::ValueRequired(opt) => write!(f, "argument {opt} requires a value"),
            Self::TooMany(opt) => {
                if let Some(max) = opt.count.get_max_count() {
                    write!(f, "too many {opt} (max {max})")
                } else {
                    write!(f, "duplicate argument {opt}")
                }
            }
        }
    }
}

impl error::Error for OptionError {}