use std::convert::Infallible;
use std::ops::Add;
use std::str::FromStr;

#[derive(Clone, PartialEq, Eq)]
struct Segment
{
    data: String,
}

struct Instruction
{
    inputs:  Vec<Segment>,
    outputs: Vec<Segment>,
}

fn part1(instructions: &[Instruction])
{
    let mut num_immediately_identifyable = 0;
    for i in instructions {
        for o in i.outputs() {
            match o.data().len() {
                2 | 4 | 3 | 7 => num_immediately_identifyable += 1,
                _ => {},
            }
        }
    }

    println!("Solution to part 1: {}", num_immediately_identifyable);
}

fn part2(instructions: Vec<Instruction>)
{
    let sum: u64 = instructions.into_iter().map(|i| i.decode()).sum();
    println!("Solution to part 2: {}", sum);
}

pub fn run(input: Vec<String>)
{
    let instructions: Vec<Instruction> = input
        .into_iter()
        .map(|s| Instruction::from_str(&s).unwrap())
        .collect();

    part1(&instructions);
    part2(instructions);
}

impl Segment
{
    pub fn data(&self) -> &String { &self.data }

    pub fn is_superset(&self, of: &Segment) -> bool
    {
        of.data().chars().all(|c| self.data.contains(c))
    }
}

impl Add for Segment
{
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output
    {
        let data = self.data + &rhs.data;
        let mut bytes: Vec<u8> = data.bytes().collect();

        bytes.sort();
        bytes.dedup();

        Self {
            data: unsafe { String::from_utf8_unchecked(bytes) },
        }
    }
}

impl Instruction
{
    pub fn outputs(&self) -> &Vec<Segment> { &self.outputs }

    pub fn decode(mut self) -> u64
    {
        let mut codes: [Option<Segment>; 10] =
            [None, None, None, None, None, None, None, None, None, None];
        // Decode the easy ones
        self.inputs.drain_filter(|s| match s.data.len() {
            2 => {
                codes[1] = Some(s.clone());
                true
            },
            3 => {
                codes[7] = Some(s.clone());
                true
            },
            4 => {
                codes[4] = Some(s.clone());
                true
            },
            7 => {
                codes[8] = Some(s.clone());
                true
            },
            _ => false,
        });
        assert_eq!(self.inputs.len(), 6);

        // Only 3 has length 5 and contains all 1 segments
        self.inputs.drain_filter(|s| {
            if s.data().len() == 5 && s.is_superset(&codes[1].as_ref().unwrap()) {
                codes[3] = Some(s.clone());
                true
            }
            else {
                false
            }
        });
        assert_eq!(self.inputs.len(), 5);

        // Combine 3 and 4 to get 9
        let nine_seg = codes[3].as_ref().unwrap().clone() + codes[4].as_ref().unwrap().clone();
        self.inputs.drain_filter(|s| s == &nine_seg);
        codes[9] = Some(nine_seg);
        assert_eq!(self.inputs.len(), 4);

        // 6 and 0 are the only remaining numbers with 6 segments on, only 0 is a
        // superset of 1s segments
        self.inputs.drain_filter(|s| {
            if s.data().len() == 6 {
                if s.is_superset(codes[1].as_ref().unwrap()) {
                    codes[0] = Some(s.clone());
                }
                else {
                    codes[6] = Some(s.clone());
                }
                true
            }
            else {
                false
            }
        });

        assert_eq!(self.inputs.len(), 2);
        // 6 is a superset of 5, but not 2
        self.inputs.drain_filter(|s| {
            assert_eq!(s.data().len(), 5);
            if codes[6].as_ref().unwrap().is_superset(s) {
                codes[5] = Some(s.clone());
            }
            else {
                codes[2] = Some(s.clone());
            }

            true
        });

        let codes: Vec<Segment> = codes.into_iter().map(|c| c.unwrap()).collect();

        let mut acc = 0;
        for (i, o) in self.outputs.iter().rev().enumerate() {
            let num = codes.iter().enumerate().find(|(_, s)| s == &o).unwrap().0;
            acc += num as u64 * 10_u64.pow(i as u32);
        }

        acc
    }
}

impl FromStr for Instruction
{
    type Err = Infallible;

    fn from_str(s: &str) -> Result<Self, Self::Err>
    {
        let parts: Vec<&str> = s.split('|').collect();
        assert_eq!(parts.len(), 2);

        Ok(Self {
            inputs:  parts[0]
                .split_whitespace()
                .map(|s| Segment::from_str(s).unwrap())
                .collect(),
            outputs: parts[1]
                .split_whitespace()
                .map(|s| Segment::from_str(s).unwrap())
                .collect(),
        })
    }
}

impl FromStr for Segment
{
    type Err = Infallible;

    fn from_str(s: &str) -> Result<Self, Self::Err>
    {
        let mut s = s.to_owned();

        // Sort the letters in the string to make later operations easier.
        unsafe { s.as_bytes_mut().sort() };

        Ok(Self { data: s })
    }
}