use std::collections::HashMap;
use std::collections::HashSet;

#[path = "../lib.rs"] mod lib;

#[derive(Clone)]
struct Valve {
    rate: u64,
    tunnels: HashMap<String, u64>,
}

struct QueueEntry {
    pos1: String,
    pos2: String,
    time1: u64,
    time2: u64,
    open: HashSet<String>,
    pressure: u64,
    potential: u64,
}

fn get_input() -> HashMap<String, Valve> {
    let mut map = HashMap::new();

    for line in lib::iter_input() {
        let (line1, line2) = lib::split_once(line.as_str(), ';').unwrap();
        let (line11, line12) = lib::split_once(line1, '=').unwrap();
        let name = line11[6..8].to_string();
        let tunnels_s = match line2.strip_prefix(" tunnels lead to valves ") {
            Some(s) => s,
            _ => line2.strip_prefix(" tunnel leads to valve ").unwrap(),
        };
        let mut tunnels = HashMap::new();
        for s in tunnels_s.split(", ") {
            tunnels.insert(s.to_string(), 1);
        }
        map.insert(name, Valve {
            rate: line12.parse().unwrap(),
            tunnels: tunnels,
        });
    }

    return map;
}

fn make_entry(
    valves: &HashMap<String, Valve>,
    pos1: String,
    pos2: String,
    time1: u64,
    time2: u64,
    open: HashSet<String>,
    pressure: u64,
) -> QueueEntry {
    let min_distance = valves.values().map(|valve|
        valve.tunnels.values().min().unwrap()
    ).min().unwrap();

    let mut rates = valves.iter()
        .filter(|(key, _valve)| !open.contains(*key))
        .map(|(_key, valve)| valve.rate)
        .collect::<Vec<u64>>();
    rates.sort();
    rates.reverse();

    let mut potential = pressure;
    let mut t = (time1, time2);
    for i in 0..rates.len() {
        let t1 = t.0.max(t.1);
        let t2 = t.0.min(t.1);
        potential += rates[i] * (t1 - 1);
        if t1 > min_distance + 1 {
            t = (t1 - min_distance - 1, t2);
        } else {
            break;
        }
    }

    // swap so that the pos1 always refers to the one behind on time
    let swap = time1 < time2;
    let (time1, time2) = if swap { (time2, time1) } else { (time1, time2) };
    let (pos1, pos2) = if swap { (pos2, pos1) } else { (pos1, pos2) };

    return QueueEntry { pos1, pos2, time1, time2, open, pressure, potential };
}

fn get_next(queue: &mut Vec<QueueEntry>) -> Option<QueueEntry> {
    // from experiments this seems to be a good way to pick the next entry
    // let (i, _) = queue.iter()
    //     .enumerate()
    //     .max_by_key(|(_, entry)| entry.potential + entry.pressure)?;
    // return Some(queue.remove(i));

    return queue.pop();
}

fn part1(valves: &HashMap<String, Valve>, time: u64) -> u64 {
    let mut max_pressure = 0;
    let mut queue = vec![make_entry(
        valves,
        "AA".to_string(),
        "AA".to_string(),
        time,
        0,
        HashSet::new(),
        0,
    )];

    while let Some(entry) = get_next(&mut queue) {
        // PERF: stop early if there is no way to reach max_pressure
        if entry.potential <= max_pressure {
            continue;
        }

        if !entry.open.contains(&entry.pos1) {
            let pressure = entry.pressure + valves[&entry.pos1].rate * (entry.time1 - 1);
            if pressure > max_pressure {
                max_pressure = pressure;
            }

            let mut open = entry.open.clone();
            open.insert(entry.pos1.clone());
            if entry.time1 > 1 {
                queue.push(make_entry(
                    valves,
                    entry.pos1.clone(),
                    entry.pos2.clone(),
                    entry.time1 - 1,
                    entry.time2,
                    open,
                    pressure,
                ));
            }
        }

        for (name, len) in valves[&entry.pos1].tunnels.iter() {
            if entry.time1 > *len {
                queue.push(make_entry(
                    valves,
                    name.clone(),
                    entry.pos2.clone(),
                    entry.time1 - len,
                    entry.time2,
                    entry.open.clone(),
                    entry.pressure,
                ));
            }
        }
    }

    return max_pressure;
}

fn part2(valves: &HashMap<String, Valve>, time: u64) -> u64 {
    let mut max_pressure = 0;
    let mut queue = vec![make_entry(
        valves,
        "AA".to_string(),
        "AA".to_string(),
        time,
        time,
        HashSet::new(),
        0,
    )];

    while let Some(entry) = get_next(&mut queue) {
        // PERF: stop early if there is no way to reach max_pressure
        if entry.potential <= max_pressure {
            continue;
        }

        if !entry.open.contains(&entry.pos1) {
            let pressure = entry.pressure + valves[&entry.pos1].rate * (entry.time1 - 1);
            if pressure > max_pressure {
                max_pressure = pressure;
                println!("{} {} {}", entry.time1, max_pressure, queue.len());
            }

            let mut open = entry.open.clone();
            open.insert(entry.pos1.clone());
            if entry.time1 > 1 {
                queue.push(make_entry(
                    valves,
                    entry.pos1.clone(),
                    entry.pos2.clone(),
                    entry.time1 - 1,
                    entry.time2,
                    open,
                    pressure,
                ));
            }
        }

        for (name, len) in valves[&entry.pos1].tunnels.iter() {
            if entry.time1 > *len {
                queue.push(make_entry(
                    valves,
                    name.clone(),
                    entry.pos2.clone(),
                    entry.time1 - len,
                    entry.time2,
                    entry.open.clone(),
                    entry.pressure,
                ));
            }
        }
    }

    return max_pressure;
}

fn optimize_tunnels(valves: &HashMap<String, Valve>, name: &String, path: &Vec<String>) -> HashMap<String, u64> {
    let mut tunnels = HashMap::new();
    let mut p = path.clone();
    p.push(name.clone());

    for (tunnel_name, d) in valves[name].tunnels.iter() {
        if valves[tunnel_name].rate > 0 || tunnel_name == "AA" {
            if *tunnels.get(tunnel_name).unwrap_or(&u64::MAX) > *d {
                tunnels.insert(tunnel_name.clone(), *d);
            }
        } else if !p.contains(tunnel_name) {
            for (tunnel_name2, d2) in optimize_tunnels(valves, tunnel_name, &p).iter() {
                if !p.contains(tunnel_name2) {
                    if *tunnels.get(tunnel_name2).unwrap_or(&u64::MAX) > *d + *d2 {
                        tunnels.insert(tunnel_name2.clone(), *d + *d2);
                    }
                }
            }
        }
    }

    return tunnels;
}

fn optimize_graph(valves: &HashMap<String, Valve>) -> HashMap<String, Valve> {
    let mut optimized = HashMap::new();
    for (name, valve) in valves.iter() {
        if valve.rate > 0 || name == "AA" {
            optimized.insert(name.clone(), Valve {
                rate: valve.rate,
                tunnels: optimize_tunnels(&valves, name, &vec![]),
            });
        }
    }
    return optimized;
}

fn main() {
    let valves = get_input();
    let optimized = optimize_graph(&valves);
    println!("part1: {}", part1(&optimized, 30));
    println!("part2: {}", part2(&optimized, 26));
}