#include <stdbool.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

#include "packet.h"
#include "buffer.h"
#include "header.h"
#include "question.h"
#include "record.h"
#include "../io/log.h"

void read_packet(PacketBuffer* buffer, Packet* packet) {
    read_header(buffer, &packet->header);
    
    packet->questions = malloc(sizeof(Question) * packet->header.questions);
    for(uint16_t i = 0; i < packet->header.questions; i++) {
        if (!read_question(buffer, &packet->questions[i])) {
            i--;
            packet->header.questions--;
        }
    }

    packet->answers = malloc(sizeof(Record) * packet->header.answers);
    for(uint16_t i = 0; i < packet->header.answers; i++) {
        if (!read_record(buffer, &packet->answers[i])) {
            i--;
            packet->header.answers--;
        }
    }

    packet->authorities = malloc(sizeof(Record) * packet->header.authoritative_entries);
    for(uint16_t i = 0; i < packet->header.authoritative_entries; i++) {
        if (!read_record(buffer, &packet->authorities[i])) {
            i--;
            packet->header.authoritative_entries--;
        }
    }

    packet->resources = malloc(sizeof(Record) * packet->header.resource_entries);
    for(uint16_t i = 0; i < packet->header.resource_entries; i++) {
        if (!read_record(buffer, &packet->resources[i])) {
            i--;
            packet->header.resource_entries--;
        }
    }
}

void write_packet(PacketBuffer* buffer, Packet* packet) {
    write_header(buffer, &packet->header);
    
    for(uint16_t i = 0; i < packet->header.questions; i++) {
        write_question(buffer, &packet->questions[i]);
    }

    for(uint16_t i = 0; i < packet->header.answers; i++) {
        write_record(buffer, &packet->answers[i]);
    }

    for(uint16_t i = 0; i < packet->header.authoritative_entries; i++) {
        write_record(buffer, &packet->authorities[i]);
    }

    for(uint16_t i = 0; i < packet->header.resource_entries; i++) {
        write_record(buffer, &packet->resources[i]);
    }
}

void free_packet(Packet* packet) {
    
    for(uint16_t i = 0; i < packet->header.questions; i++) {
        free_question(&packet->questions[i]);
    }
    free(packet->questions);

    for(uint16_t i = 0; i < packet->header.answers; i++) {
        free_record(&packet->answers[i]);
    }
    free(packet->answers);

    for(uint16_t i = 0; i < packet->header.authoritative_entries; i++) {
        free_record(&packet->authorities[i]);
    }
    free(packet->authorities);

    for(uint16_t i = 0; i < packet->header.resource_entries; i++) {
        free_record(&packet->resources[i]);
    }
    free(packet->resources);
}

bool get_random_a(Packet* packet, IpAddr* addr) {
    for (uint16_t i = 0; i < packet->header.answers; i++) {
        Record record = packet->answers[i];
        if (record.type == A) {
            create_ip_addr(record.data.a.addr, addr);
            return true;
        } else if (record.type == AAAA) {
            create_ip_addr6(record.data.aaaa.addr, addr);
            return true;
        }
    }
    return false;
}

static bool ends_with(uint8_t* full, uint8_t* end) {
    uint8_t check = end[0];
    uint8_t len = full[0];

    if (check > len) {
        return false;
    }

    for(uint8_t i = 0; i < check; i++) {
        if (end[check - 1 - i] != full[len - 1 - i]) {
            return false;
        }
    }
    
    return true;
}

static bool equals(uint8_t* a, uint8_t* b) {
    if (a[0] != b[0]) {
        return false;
    }

    for(uint8_t i = 1; i < a[0] + 1; i++) {
        if(a[i] != b[i]) {
            return false;
        }
    }

    return true;
}

bool get_resolved_ns(Packet* packet, uint8_t* qname, IpAddr* addr) {
    for (uint16_t i = 0; i < packet->header.authoritative_entries; i++) {
        Record record = packet->authorities[i];
        if (record.type != NS) {
            continue;
        }

        if(!ends_with(qname, record.domain)) {
            continue;
        }

        for (uint16_t i = 0; i < packet->header.resource_entries; i++) {
            Record resource = packet->resources[i];
            if (!equals(record.data.ns.host, resource.domain)) {
                continue;
            }

            if (resource.type == A) {
                create_ip_addr(record.data.a.addr, addr);
                return true;
            } else if (resource.type == AAAA) {
                create_ip_addr6(record.data.aaaa.addr, addr);
                return true;
            }
        }
    }
    return false;
}

bool get_unresoled_ns(Packet* packet, uint8_t* qname, Question* question) {
    for (uint16_t i = 0; i < packet->header.authoritative_entries; i++) {
        Record record = packet->authorities[i];
        if (record.type != NS) {
            continue;
        }

        if(!ends_with(qname, record.domain)) {
            continue;
        }

        uint8_t* host = record.data.ns.host;

        question->qtype = NS;
        question->domain = malloc(host[0] + 1);
        memcpy(question->domain, host, host[0] + 1);
        return true;
    }
    return false;
}