use std::{
    collections::HashMap,
    env,
    fs::File,
    io::{Read, Write},
    path::Path,
};

fn main() {
    let args: Vec<String> = env::args().collect();
    let file_name = Path::new(&args[1]);

    let content = read_file(file_name);
    let stats = get_stats(&content);
    let tree = create_tree(stats);

    let result = compress(content, tree);

    std::io::stdout()
        .write_all(&result)
        .expect("Couldn't write the result");
}

fn read_file(path: &Path) -> String {
    let mut file = File::open(path).expect("Can't open file");

    let mut contents = String::new();
    file.read_to_string(&mut contents)
        .expect("Couldn't read the file");

    contents
}

fn get_stats(text: &String) -> HashMap<String, u32> {
    let mut letters = HashMap::new();

    for letter in text.chars() {
        *letters.entry(letter.to_string()).or_insert(0) += 1;
    }

    letters
}

fn create_tree(letters: HashMap<String, u32>) -> HashMap<String, String> {
    let mut tree = HashMap::new();
    let mut sorted: Vec<String> = letters.keys().cloned().collect();

    sorted.sort_by_key(|key| std::cmp::Reverse(letters[key]));

    let mut a: u32 = 0;

    for letter in sorted {
        let mut zeroes = "".to_string();

        for _ in 0..a {
            zeroes.push('0');
        }

        tree.insert(letter, format!("{zeroes}1"));
        a += 1;
    }

    tree
}

fn compress(text: String, tree: HashMap<String, String>) -> Vec<u8> {
    let mut compressed = Vec::new();

    let mut current_byte: u8 = 0;
    let mut bit_count: u8 = 0;

    for letter in text.chars() {
        let code = tree.get(&letter.to_string()).unwrap();

        for bit in code.chars() {
            current_byte <<= 1;

            if bit == '1' {
                current_byte |= 1;
            }

            bit_count += 1;

            if bit_count == 8 {
                compressed.push(current_byte);

                current_byte = 0;
                bit_count = 0;
            }
        }
    }

    if bit_count > 0 {
        current_byte <<= 8 - bit_count;
        compressed.push(current_byte);
    }

    compressed
}