#include "compressor.h"
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>

#define MAGIC "HUF1"
#define SYMBOLS 256

typedef struct Node {
    uint8_t symbol;
    uint64_t freq;
    struct Node *left, *right;
} Node;

typedef struct {
    uint32_t bits;
    uint8_t length;
} Code;

static uint64_t frequencies[SYMBOLS];
static Code codes[SYMBOLS];

/* ---------- Huffman strom ---------- */

static Node* create_node(uint8_t symbol, uint64_t freq, Node* l, Node* r) {
    Node* n = malloc(sizeof(Node));
    n->symbol = symbol;
    n->freq = freq;
    n->left = l;
    n->right = r;
    return n;
}

static Node* build_tree() {
    Node* nodes[SYMBOLS];
    int count = 0;

    for (int i = 0; i < SYMBOLS; i++)
        if (frequencies[i])
            nodes[count++] = create_node(i, frequencies[i], NULL, NULL);

    if (count == 1)
        return create_node(0, nodes[0]->freq, nodes[0], NULL);

    while (count > 1) {
        int a = 0, b = 1;
        if (nodes[b]->freq < nodes[a]->freq) { int t=a;a=b;b=t; }

        for (int i = 2; i < count; i++) {
            if (nodes[i]->freq < nodes[a]->freq) {
                b = a;
                a = i;
            } else if (nodes[i]->freq < nodes[b]->freq) {
                b = i;
            }
        }

        Node* merged = create_node(
            0,
            nodes[a]->freq + nodes[b]->freq,
            nodes[a], nodes[b]
        );

        if (a > b) { int t=a;a=b;b=t; }
        nodes[a] = merged;
        nodes[b] = nodes[count - 1];
        count--;
    }
    return nodes[0];
}

static void build_codes(Node* n, uint32_t bits, uint8_t len) {
    if (!n->left && !n->right) {
        codes[n->symbol].bits = bits;
        codes[n->symbol].length = len;
        return;
    }
    if (n->left)  build_codes(n->left,  bits << 1,     len + 1);
    if (n->right) build_codes(n->right, (bits << 1)|1, len + 1);
}

/* ---------- Bitový výstup ---------- */

typedef struct {
    FILE* f;
    uint8_t buf;
    uint8_t count;
} BitWriter;

static void bw_init(BitWriter* w, FILE* f) {
    w->f = f;
    w->buf = 0;
    w->count = 0;
}

static void bw_write(BitWriter* w, uint32_t bits, uint8_t len) {
    for (int i = len - 1; i >= 0; i--) {
        w->buf = (w->buf << 1) | ((bits >> i) & 1);
        if (++w->count == 8) {
            fwrite(&w->buf, 1, 1, w->f);
            w->count = 0;
        }
    }
}

static void bw_flush(BitWriter* w) {
    if (w->count) {
        w->buf <<= (8 - w->count);
        fwrite(&w->buf, 1, 1, w->f);
    }
}

/* ---------- Bitový vstup ---------- */

typedef struct {
    FILE* f;
    uint8_t buf;
    uint8_t count;
} BitReader;

static void br_init(BitReader* r, FILE* f) {
    r->f = f;
    r->count = 0;
}

static int br_read(BitReader* r) {
    if (!r->count) {
        if (fread(&r->buf, 1, 1, r->f) != 1)
            return -1;
        r->count = 8;
    }
    int bit = (r->buf >> 7) & 1;
    r->buf <<= 1;
    r->count--;
    return bit;
}

/* ---------- Kompresia ---------- */

void compress_file(const char* input, const char* output) {
    FILE* fi = fopen(input, "rb");
    FILE* fo = fopen(output, "wb");
    if (!fi || !fo) exit(1);

    memset(frequencies, 0, sizeof(frequencies));

    int c;
    uint64_t size = 0;
    while ((c = fgetc(fi)) != EOF) {
        frequencies[c]++;
        size++;
    }
    rewind(fi);

    Node* root = build_tree();
    build_codes(root, 0, 0);

    fwrite(MAGIC, 1, 4, fo);
    fwrite(&size, sizeof(uint64_t), 1, fo);

    for (int i = 0; i < SYMBOLS; i++) {
        uint32_t f = frequencies[i];
        fwrite(&f, sizeof(uint32_t), 1, fo);
    }

    BitWriter bw;
    bw_init(&bw, fo);

    while ((c = fgetc(fi)) != EOF)
        bw_write(&bw, codes[c].bits, codes[c].length);

    bw_flush(&bw);
    fclose(fi);
    fclose(fo);
}

/* ---------- Dekompresia ---------- */

void decompress_file(const char* input, const char* output) {
    FILE* fi = fopen(input, "rb");
    FILE* fo = fopen(output, "wb");
    if (!fi || !fo) exit(1);

    char magic[4];
    fread(magic, 1, 4, fi);
    if (memcmp(magic, MAGIC, 4)) exit(1);

    uint64_t size;
    fread(&size, sizeof(uint64_t), 1, fi);

    for (int i = 0; i < SYMBOLS; i++) {
        uint32_t f;
        fread(&f, sizeof(uint32_t), 1, fi);
        frequencies[i] = f;
    }

    Node* root = build_tree();
    BitReader br;
    br_init(&br, fi);

    for (uint64_t i = 0; i < size; i++) {
        Node* n = root;
        while (n->left || n->right) {
            int bit = br_read(&br);
            n = bit ? n->right : n->left;
        }
        fputc(n->symbol, fo);
    }
    fclose(fi);
    fclose(fo);
}