// Pruebas públicas de los pasos 4 y 5. Se corre una por proceso: huffman_tests nombre.
#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <fstream>
#include <functional>
#include <iostream>
#include <queue>
#include <random>
#include <stdexcept>
#include <string>
#include <vector>

#include "byte_file.h"
#include "huffman.h"

namespace {

Bytes bytes_of(const std::string& text) { return Bytes(text.begin(), text.end()); }

Frequencies frequencies_of(const std::vector<std::pair<int, std::uint64_t>>& pairs) {
    Frequencies frequencies{};
    for (const auto& [symbol, count] : pairs) {
        frequencies[symbol] = count;
    }
    return frequencies;
}

std::uint64_t total_bits(const Frequencies& frequencies, const Codes& codes) {
    std::uint64_t bits = 0;
    for (int symbol = 0; symbol < 256; symbol++) {
        bits += frequencies[symbol] * codes[symbol].size();
    }
    return bits;
}

// El costo óptimo: la suma de los pesos de cada unión, con la cola de la biblioteca.
std::uint64_t optimal_bits(const Frequencies& frequencies) {
    std::priority_queue<std::uint64_t, std::vector<std::uint64_t>, std::greater<std::uint64_t>> queue;
    for (std::uint64_t count : frequencies) {
        if (count > 0) {
            queue.push(count);
        }
    }
    std::uint64_t total = 0;
    while (queue.size() > 1) {
        std::uint64_t first = queue.top();
        queue.pop();
        std::uint64_t second = queue.top();
        queue.pop();
        total += first + second;
        queue.push(first + second);
    }
    return total;
}

void assert_prefix_free(const Frequencies& frequencies, const Codes& codes) {
    for (int a = 0; a < 256; a++) {
        for (int b = 0; b < 256; b++) {
            if (a != b && frequencies[a] > 0 && frequencies[b] > 0) {
                assert(codes[b].rfind(codes[a], 0) != 0);
            }
        }
    }
}

void assert_round_trip(const Bytes& data) {
    Bytes compressed = compress_bytes(data);
    assert(decompress_bytes(compressed) == data);
}

// Comprueba que descomprimir lance std::runtime_error, y no otra cosa.
void assert_rejected(const Bytes& compressed) {
    bool rejected = false;
    try {
        decompress_bytes(compressed);
    } catch (const std::runtime_error&) {
        rejected = true;
    }
    assert(rejected);
}

void test_count_frequencies() {
    Bytes data = {'a', 'b', 'a', 0x00, 0xFF, 0xFF, 0xC3, 0xFF};
    Frequencies frequencies = count_frequencies(data);
    assert(frequencies['a'] == 2);
    assert(frequencies['b'] == 1);
    assert(frequencies[0x00] == 1);
    assert(frequencies[0xFF] == 3);
    assert(frequencies[0xC3] == 1);
    assert(frequencies['c'] == 0);
    assert(count_frequencies(Bytes{})[0] == 0);
}

void test_tree_guide_example() {
    Frequencies frequencies = frequencies_of({{'a', 5}, {'b', 2}, {'c', 1}, {'d', 1}});
    HuffmanTree tree = build_tree(frequencies);
    assert(tree.nodes.size() == 7);
    assert(tree.root == 6);
    assert(tree.nodes[tree.root].weight == 9);
    Codes codes = make_codes(tree);
    assert(codes['a'] == "1");
    assert(codes['b'] == "00");
    assert(codes['c'] == "010");
    assert(codes['d'] == "011");
    assert(total_bits(frequencies, codes) == 15);
}

void test_tree_ties_follow_creation_order() {
    Frequencies frequencies = frequencies_of({{'a', 1}, {'b', 1}, {'c', 1}, {'d', 1}});
    Codes codes = make_codes(build_tree(frequencies));
    assert(codes['a'] == "00");
    assert(codes['b'] == "01");
    assert(codes['c'] == "10");
    assert(codes['d'] == "11");
}

void test_tree_weights_above_32_bits() {
    Frequencies frequencies = frequencies_of({{'a', 5000000000ULL}, {'b', 1}, {'c', 1}, {'d', 3000000000ULL}});
    Codes codes = make_codes(build_tree(frequencies));
    assert(codes['a'].size() == 1);
    assert(codes['d'].size() == 2);
    assert(codes['b'].size() == 3);
}

void test_tree_empty() {
    HuffmanTree tree = build_tree(Frequencies{});
    assert(tree.root == -1);
    assert(tree.nodes.empty());
    assert(make_codes(tree)[0].empty());
}

void test_tree_one_symbol() {
    Frequencies frequencies = frequencies_of({{0xFF, 42}});
    HuffmanTree tree = build_tree(frequencies);
    assert(tree.nodes.size() == 1);
    assert(tree.root == 0);
    assert(tree.nodes[0].is_leaf());
    assert(tree.nodes[0].symbol == 0xFF);
    assert(make_codes(tree)[0xFF].empty());
}

void test_tree_random_is_optimal() {
    std::mt19937 generator(13);
    for (int round = 0; round < 300; round++) {
        Frequencies frequencies{};
        int symbols = 2 + static_cast<int>(generator() % 255);
        for (int i = 0; i < symbols; i++) {
            frequencies[generator() % 256] = 1 + generator() % 1000;
        }
        HuffmanTree tree = build_tree(frequencies);
        Codes codes = make_codes(tree);
        assert_prefix_free(frequencies, codes);
        assert(total_bits(frequencies, codes) == optimal_bits(frequencies));
    }
}

void test_codes_longer_than_a_word() {
    // Frecuencias de Fibonacci: cada unión absorbe al árbol anterior, y los códigos se alargan.
    Frequencies frequencies{};
    std::uint64_t a = 1;
    std::uint64_t b = 1;
    for (int symbol = 0; symbol < 80; symbol++) {
        frequencies[symbol] = a;
        std::uint64_t next = a + b;
        a = b;
        b = next;
    }
    HuffmanTree tree = build_tree(frequencies);
    Codes codes = make_codes(tree);
    std::size_t longest = 0;
    for (const std::string& code : codes) {
        longest = std::max(longest, code.size());
    }
    assert(longest == 79);
    assert_prefix_free(frequencies, codes);
    // El árbol también viaja por el archivo.
    Bytes out;
    BitWriter writer(out);
    write_tree(tree, writer);
    writer.finish();
    BitReader reader(out, 0);
    assert(make_codes(read_tree(reader)) == codes);
}

void test_tree_write_guide_example() {
    Frequencies frequencies = frequencies_of({{'a', 2}, {'b', 1}});
    HuffmanTree tree = build_tree(frequencies);
    Bytes out;
    BitWriter writer(out);
    write_tree(tree, writer);
    writer.finish();
    // 0 | 1 01100010 | 1 01100001, con relleno
    assert((out == Bytes{0x58, 0xAC, 0x20}));
}

void test_tree_write_and_read_all_symbols() {
    Frequencies frequencies{};
    for (int symbol = 0; symbol < 256; symbol++) {
        frequencies[symbol] = 1 + symbol % 7;
    }
    HuffmanTree tree = build_tree(frequencies);
    Bytes out;
    BitWriter writer(out);
    write_tree(tree, writer);
    writer.finish();
    assert(out.size() == 320);  // 255 + 256 * 9 = 2559 bits
    BitReader reader(out, 0);
    HuffmanTree copy = read_tree(reader);
    assert(make_codes(copy) == make_codes(tree));
}

void test_tree_read_one_leaf() {
    Bytes bits = {0xFF, 0x80};  // 1, y el símbolo 11111111
    BitReader reader(bits, 0);
    HuffmanTree tree = read_tree(reader);
    assert(tree.nodes.size() == 1);
    assert(tree.root == 0);
    assert(tree.nodes[0].symbol == 0xFF);
}

void test_exact_bytes() {
    Bytes expected = {'H', 'U', 'F', '1', 3, 0, 0, 0, 0, 0, 0, 0, 0x58, 0xAC, 0x38};
    assert(compress_bytes(bytes_of("aab")) == expected);
}

void test_exact_size_in_header() {
    Bytes data(70000, 'x');
    Bytes compressed = compress_bytes(data);
    // 70000 = 0x11170: el byte menos significativo va primero.
    assert(compressed[4] == 0x70);
    assert(compressed[5] == 0x11);
    assert(compressed[6] == 0x01);
    assert(compressed[7] == 0x00);
}

void test_empty_input() {
    Bytes compressed = compress_bytes(Bytes{});
    assert((compressed == Bytes{'H', 'U', 'F', '1', 0, 0, 0, 0, 0, 0, 0, 0}));
    assert(decompress_bytes(compressed).empty());
}

void test_one_byte_input() {
    Bytes data = {0xFF};
    Bytes compressed = compress_bytes(data);
    assert(compressed.size() == header_size + 2);
    assert(decompress_bytes(compressed) == data);
}

void test_one_symbol_input() {
    Bytes data(100000, 'a');
    Bytes compressed = compress_bytes(data);
    // Una hoja son 9 bits, y sus códigos no ocupan nada.
    assert(compressed.size() == header_size + 2);
    assert(decompress_bytes(compressed) == data);
}

void test_two_symbols() {
    assert_round_trip(bytes_of("ab"));
    assert_round_trip(bytes_of("aaaaaaab"));
    Bytes data = compress_bytes(Bytes(16, 'a'));
    assert(decompress_bytes(data).size() == 16);
}

void test_padding_is_not_data() {
    // Cualquier relleno, de 0 a 7 bits, debe ignorarse: el tamaño manda.
    for (std::size_t n = 1; n <= 40; n++) {
        Bytes data(n, 'b');
        data[0] = 'a';
        assert_round_trip(data);
    }
}

void test_all_byte_values() {
    Bytes data;
    for (int round = 0; round < 5; round++) {
        for (int value = 0; value < 256; value++) {
            data.push_back(static_cast<std::uint8_t>(value));
        }
    }
    assert_round_trip(data);
}

void test_special_bytes() {
    assert_round_trip(Bytes{0x00, 0x00, 0xFF, 0x0A, 0x0D, 0x0A, 0x1A, 0x20, 0x00});
}

void test_random_round_trips() {
    std::mt19937 generator(13);
    for (int round = 0; round < 200; round++) {
        Bytes data(generator() % 2000);
        int alphabet = 1 + static_cast<int>(generator() % 256);
        for (std::uint8_t& byte : data) {
            byte = static_cast<std::uint8_t>(generator() % alphabet);
        }
        assert_round_trip(data);
    }
}

void test_reference_file() {
    // Lo comprimió la solución del curso. Si lee esto, su formato coincide.
    Bytes compressed = read_file("lab/tests/reference.huf");
    Bytes original = read_file("lab/tests/fixture.txt");
    assert(decompress_bytes(compressed) == original);
    assert(compress_bytes(original) == compressed);
}

void test_rejects_bad_header() {
    assert_rejected(Bytes{});
    assert_rejected(bytes_of("HUF"));
    assert_rejected(bytes_of("HUF1abc"));
    assert_rejected(bytes_of("esto no es un archivo comprimido"));
    // Un archivo correcto, con un solo byte cambiado en la marca.
    Bytes good = compress_bytes(bytes_of("aab"));
    for (std::size_t i = 0; i < 4; i++) {
        Bytes bad = good;
        bad[i] = 'x';
        assert_rejected(bad);
    }
}

void test_rejects_truncated_file() {
    Bytes compressed = compress_bytes(bytes_of("el ñandú llegó después, y la zorra no llegó"));
    for (std::size_t size = 0; size < compressed.size(); size++) {
        Bytes cut(compressed.begin(), compressed.begin() + static_cast<std::ptrdiff_t>(size));
        assert_rejected(cut);
    }
}

void test_rejects_endless_tree() {
    // Todo ceros: un nodo interno tras otro. Sin un límite, la recursión agota la pila.
    Bytes compressed = {'H', 'U', 'F', '1', 1, 0, 0, 0, 0, 0, 0, 0};
    compressed.resize(compressed.size() + 1000000, 0);
    assert_rejected(compressed);
}

// Escribe en preorden un árbol perfecto con 2^depth hojas.
void write_perfect_tree(BitWriter& writer, int depth) {
    if (depth == 0) {
        writer.put_bit(true);
        writer.put_byte('x');
        return;
    }
    writer.put_bit(false);
    write_perfect_tree(writer, depth - 1);
    write_perfect_tree(writer, depth - 1);
}

void test_rejects_too_many_leaves() {
    // 512 hojas, sin pasar del límite de profundidad. Le siguen bits de sobra para decodificar.
    Bytes out = {'H', 'U', 'F', '1', 1, 0, 0, 0, 0, 0, 0, 0};
    BitWriter writer(out);
    write_perfect_tree(writer, 9);
    for (int i = 0; i < 16; i++) {
        writer.put_bit(false);
    }
    writer.finish();
    assert_rejected(out);
}

struct Test {
    const char* name;
    void (*run)();
};

const Test tree_tests[] = {
    {"count_frequencies", test_count_frequencies},
    {"tree_guide_example", test_tree_guide_example},
    {"tree_ties_follow_creation_order", test_tree_ties_follow_creation_order},
    {"tree_weights_above_32_bits", test_tree_weights_above_32_bits},
    {"tree_empty", test_tree_empty},
    {"tree_one_symbol", test_tree_one_symbol},
    {"tree_random_is_optimal", test_tree_random_is_optimal},
    {"codes_longer_than_a_word", test_codes_longer_than_a_word},
    {"tree_write_guide_example", test_tree_write_guide_example},
    {"tree_write_and_read_all_symbols", test_tree_write_and_read_all_symbols},
    {"tree_read_one_leaf", test_tree_read_one_leaf},
    {"exact_bytes", test_exact_bytes},
    {"exact_size_in_header", test_exact_size_in_header},
    {"empty_input", test_empty_input},
    {"one_byte_input", test_one_byte_input},
    {"one_symbol_input", test_one_symbol_input},
    {"two_symbols", test_two_symbols},
    {"padding_is_not_data", test_padding_is_not_data},
    {"all_byte_values", test_all_byte_values},
    {"special_bytes", test_special_bytes},
    {"random_round_trips", test_random_round_trips},
    {"reference_file", test_reference_file},
    {"rejects_bad_header", test_rejects_bad_header},
    {"rejects_truncated_file", test_rejects_truncated_file},
    {"rejects_endless_tree", test_rejects_endless_tree},
    {"rejects_too_many_leaves", test_rejects_too_many_leaves},
};

}  // namespace

int main(int argc, char** argv) {
    if (argc != 2) {
        std::cerr << "uso: huffman_tests nombre\n";
        return 2;
    }
    for (const Test& test : tree_tests) {
        if (std::strcmp(test.name, argv[1]) == 0) {
            test.run();
            return 0;
        }
    }
    std::cerr << "no existe la prueba " << argv[1] << "\n";
    return 2;
}