use std::collections::HashMap;

use nom::{
    IResult, Parser,
    branch::alt,
    bytes::complete::take_while1,
    character::complete::{char, multispace0},
    combinator::map,
    multi::fold,
    sequence::{delimited, preceded},
};
use roaring::RoaringBitmap;

#[derive(Debug, Clone)]
enum Expr<'a> {
    Term(&'a str),
    And(Box<Expr<'a>>, Box<Expr<'a>>),
    Or(Box<Expr<'a>>, Box<Expr<'a>>),
    Not(Box<Expr<'a>>),
}

fn identifier(i: &str) -> IResult<&str, &str> {
    take_while1(|c: char| !matches!(c, '&' | '|' | '!' | '(' | ')') && !c.is_whitespace()).parse(i)
}

fn parens(i: &str) -> IResult<&str, Expr<'_>> {
    delimited(
        (multispace0, char('('), multispace0),
        or_expr,
        (multispace0, char(')')),
    )
    .parse(i)
}

fn atom(i: &str) -> IResult<&str, Expr<'_>> {
    delimited(
        multispace0,
        alt((parens, map(identifier, |s: &str| Expr::Term(s)))),
        multispace0,
    )
    .parse(i)
}

fn not_expr(i: &str) -> IResult<&str, Expr<'_>> {
    alt((
        map(
            preceded((multispace0, char('!'), multispace0), not_expr),
            |e| Expr::Not(Box::new(e)),
        ),
        atom,
    ))
    .parse(i)
}

fn and_expr(i: &str) -> IResult<&str, Expr<'_>> {
    let (i, first) = not_expr(i)?;
    fold(
        0..,
        preceded((multispace0, char('&'), multispace0), not_expr),
        move || first.clone(),
        |acc, rhs| Expr::And(Box::new(acc), Box::new(rhs)),
    )
    .parse(i)
}

fn or_expr(i: &str) -> IResult<&str, Expr<'_>> {
    let (i, first) = and_expr(i)?;
    fold(
        0..,
        preceded((multispace0, char('|'), multispace0), and_expr),
        move || first.clone(),
        |acc, rhs| Expr::Or(Box::new(acc), Box::new(rhs)),
    )
    .parse(i)
}

fn parse(input: &str) -> eyre::Result<Expr<'_>> {
    let (rest, expr) = or_expr(input).map_err(|e| eyre::eyre!("{e:?}"))?;
    let rest = rest.trim();
    if !rest.is_empty() {
        eyre::bail!("trailing input after expression: {rest:?}")
    }
    Ok(expr)
}

struct InventoryIndex<'a> {
    term_to_hosts: HashMap<&'a str, RoaringBitmap>,
    all_hosts: RoaringBitmap,
}
impl<'a> InventoryIndex<'a> {
    fn new<H, FN, FT>(hosts: &'a [H], name_of: FN, tags_of: FT) -> Self
    where
        FN: Fn(&H) -> &str,
        FT: Fn(&H) -> &[String],
    {
        let mut term_to_hosts = HashMap::new();
        let mut all_hosts = RoaringBitmap::new();

        for (idx, host) in hosts.iter().enumerate() {
            let u_idx = idx as u32;
            all_hosts.insert(u_idx);

            term_to_hosts
                .entry(name_of(host))
                .or_insert_with(RoaringBitmap::new)
                .insert(u_idx);

            for tag in tags_of(host) {
                term_to_hosts
                    .entry(tag.as_str())
                    .or_insert_with(RoaringBitmap::new)
                    .insert(u_idx);
            }
        }

        Self {
            term_to_hosts,
            all_hosts,
        }
    }

    fn eval(&self, expr: &Expr) -> RoaringBitmap {
        match expr {
            Expr::Term(t) => self.term_to_hosts.get(t).cloned().unwrap_or_default(),
            Expr::And(a, b) => self.eval(a) & self.eval(b),
            Expr::Or(a, b) => self.eval(a) | self.eval(b),
            Expr::Not(e) => &self.all_hosts - self.eval(e),
        }
    }
}

pub fn select_indices<H, FN, FT>(
    hosts: &[H],
    limit: Option<&str>,
    name_of: FN,
    tags_of: FT,
) -> eyre::Result<Vec<usize>>
where
    FN: Fn(&H) -> &str,
    FT: Fn(&H) -> &[String],
{
    let Some(expr_src) = limit.map(str::trim).filter(|s| !s.is_empty()) else {
        return Ok((0..hosts.len()).collect());
    };

    let expr = parse(expr_src)?;

    let inventory_index = InventoryIndex::new(hosts, name_of, tags_of);
    Ok(inventory_index
        .eval(&expr)
        .into_iter()
        .map(|idx| idx as usize)
        .collect())
}

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

    fn h(name: &'static str, tags: &[&'static str]) -> (String, Vec<String>) {
        (
            name.to_string(),
            tags.iter().map(|s| s.to_string()).collect(),
        )
    }

    fn run(hosts: &[(String, Vec<String>)], limit: Option<&str>) -> Vec<String> {
        let idx = select_indices(hosts, limit, |(n, _)| n.as_str(), |(_, t)| t.as_slice())
            .expect("expected limit to parse");
        idx.into_iter().map(|i| hosts[i].0.clone()).collect()
    }

    #[test]
    fn no_limit_returns_all() {
        let hs = vec![h("a", &[]), h("b", &[])];
        assert_eq!(run(&hs, None), vec!["a", "b"]);
        assert_eq!(run(&hs, Some("")), vec!["a", "b"]);
        assert_eq!(run(&hs, Some("   ")), vec!["a", "b"]);
    }

    #[test]
    fn term_matches_name_or_tag() {
        let hs = vec![h("vyos-01", &["router"]), h("ap-01", &["openwrt"])];
        assert_eq!(run(&hs, Some("vyos-01")), vec!["vyos-01"]);
        assert_eq!(run(&hs, Some("router")), vec!["vyos-01"]);
        assert_eq!(run(&hs, Some("openwrt")), vec!["ap-01"]);
    }

    #[test]
    fn intersection() {
        let hs = vec![
            h("a", &["eu", "openwrt"]),
            h("b", &["eu", "router"]),
            h("c", &["us", "openwrt"]),
        ];
        assert_eq!(run(&hs, Some("eu & openwrt")), vec!["a"]);
        assert_eq!(run(&hs, Some("eu & !openwrt")), vec!["b"]);
    }

    #[test]
    fn union() {
        let hs = vec![h("a", &["router"]), h("b", &["ap"]), h("c", &[])];
        assert_eq!(run(&hs, Some("router | ap")), vec!["a", "b"]);
    }

    #[test]
    fn negation() {
        let hs = vec![h("a", &["router"]), h("b", &["router"]), h("c", &["ap"])];
        assert_eq!(run(&hs, Some("!ap")), vec!["a", "b"]);
        assert_eq!(run(&hs, Some("router & !b")), vec!["a"]);
    }

    #[test]
    fn double_negation() {
        let hs = vec![h("a", &["x"]), h("b", &[])];
        assert_eq!(run(&hs, Some("!!x")), vec!["a"]);
    }

    #[test]
    fn precedence_and_binds_tighter_than_or() {
        let hs = vec![h("a", &["x"]), h("b", &["y", "z"]), h("c", &["y"])];
        assert_eq!(run(&hs, Some("x | y & z")), vec!["a", "b"]);
    }

    #[test]
    fn parens_override_precedence() {
        let hs = vec![
            h("a", &["x", "y"]),
            h("b", &["x", "z"]),
            h("c", &["y", "z"]),
        ];
        assert_eq!(run(&hs, Some("(x | y) & !z")), vec!["a"]);
        assert_eq!(run(&hs, Some("(x | z) & !y")), vec!["b"]);
    }

    #[test]
    fn example_1() {
        let hs = vec![
            h("ap-par", &["eu", "par", "openwrt"]),
            h("ap-prg", &["eu", "prg", "openwrt"]),
            h("ap-lon", &["eu", "lon", "openwrt"]),
            h("ap-nyc", &["us", "nyc", "openwrt"]),
            h("rt-par", &["eu", "par", "router"]),
        ];
        assert_eq!(
            run(&hs, Some("eu & openwrt & !prg")),
            vec!["ap-par", "ap-lon"]
        );
    }

    #[test]
    fn whitespace_optional() {
        let hs = vec![h("a", &["x", "y"]), h("b", &["x"])];
        assert_eq!(run(&hs, Some("x&y")), vec!["a"]);
        assert_eq!(run(&hs, Some(" x  &  y ")), vec!["a"]);
        assert_eq!(run(&hs, Some("!!x|y")), vec!["a", "b"]);
    }

    #[test]
    fn identifier_specials() {
        let hs = vec![h("ap.east.01", &["openwrt"]), h("vyos-01-eu", &["router"])];
        assert_eq!(run(&hs, Some("ap.east.01")), vec!["ap.east.01"]);
        assert_eq!(run(&hs, Some("vyos-01-eu")), vec!["vyos-01-eu"]);
    }
}
