#include "nfc_engine/nfc_brute.h"
#include "esp_log.h"
#include "esp_task_wdt.h"
#include "esp_timer.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include <string.h>

static const char *TAG = "nfc_brute";

/* Community default keys (subset from public Proxmark3 / MCT-style lists). */
static const uint8_t k_builtin[][6] = {
    {0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF}, {0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
    {0xA0, 0xA1, 0xA2, 0xA3, 0xA4, 0xA5}, {0xA5, 0xA4, 0xA3, 0xA2, 0xA1, 0xA0},
    {0x89, 0xEC, 0xA9, 0x7F, 0x8C, 0x2A}, {0x5C, 0x8F, 0xF9, 0x99, 0x0D, 0xA2},
    {0x75, 0xCC, 0xB5, 0x9C, 0x9B, 0xED}, {0xD0, 0x1A, 0xFE, 0xEB, 0x89, 0x0A},
    {0x4B, 0x79, 0x1B, 0xEA, 0x7B, 0xCC}, {0x26, 0x12, 0xC6, 0xDE, 0x84, 0xCA},
    {0x70, 0x7B, 0x11, 0xFC, 0x14, 0x81}, {0x03, 0xF9, 0x06, 0x76, 0x46, 0xAE},
    {0x23, 0x52, 0xC5, 0xB5, 0x6D, 0x85}, {0xB0, 0xB1, 0xB2, 0xB3, 0xB4, 0xB5},
    {0xC0, 0xC1, 0xC2, 0xC3, 0xC4, 0xC5}, {0xD0, 0xD1, 0xD2, 0xD3, 0xD4, 0xD5},
    {0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF}, {0x4D, 0x3A, 0x99, 0xC3, 0x51, 0xDD},
    {0x1A, 0x98, 0x2C, 0x7E, 0x45, 0x9A}, {0xFA, 0xFA, 0xFA, 0xFA, 0xFA, 0xFA},
    {0xFB, 0xFB, 0xFB, 0xFB, 0xFB, 0xFB}, {0xD3, 0xF7, 0xD3, 0xF7, 0xD3, 0xF7},
    {0x5A, 0x1B, 0x85, 0xFC, 0xE2, 0x0A}, {0x71, 0x4C, 0x5C, 0x88, 0x6E, 0x97},
    {0x58, 0x7E, 0xE5, 0xF9, 0x35, 0x0F}, {0xA0, 0x47, 0x8C, 0xC3, 0x90, 0x91},
    {0x53, 0x3C, 0xB6, 0xC7, 0x23, 0xF6}, {0x8F, 0xD0, 0xA4, 0xF2, 0x56, 0xE9},
    {0xE0, 0x00, 0x00, 0x00, 0x00, 0x00}, {0xE7, 0xD6, 0x06, 0x4C, 0x58, 0x60},
    {0xB2, 0x7C, 0xCA, 0xB3, 0x0D, 0xBD}, {0xD2, 0xEC, 0xE8, 0xB9, 0x39, 0x5E},
    {0x14, 0x94, 0xE8, 0x16, 0x63, 0xD7}, {0x7C, 0x9F, 0xB8, 0x47, 0x42, 0x42},
    {0x56, 0x93, 0x69, 0xC5, 0xA0, 0xE5}, {0x63, 0x21, 0x93, 0xBE, 0x1C, 0x3C},
    {0x8E, 0x26, 0x5B, 0xE2, 0x45, 0xBF}, {0xF4, 0x6B, 0x6D, 0xC0, 0xD6, 0xC4},
    {0x2A, 0xA0, 0x5E, 0xD1, 0x85, 0x6F}, {0xAE, 0x3F, 0xF4, 0xEE, 0xA0, 0xDB},
};

#define NBUILTIN (sizeof(k_builtin) / sizeof(k_builtin[0]))
#define MAX_VARIANTS_PER_KEY 14
#define MAX_TRIES_BEFORE_WDT 48

static bool add_sector_hit(cJSON *hits, uint8_t sector, const uint8_t key[6], const char *key_type)
{
    if (!hits || !key || !key_type) {
        return false;
    }
    char hx[16];
    for (int i = 0; i < 6; i++) {
        snprintf(hx + i * 2, 3, "%02X", key[i]);
    }
    hx[12] = 0;
    cJSON *h = cJSON_CreateObject();
    if (!h) {
        return false;
    }
    cJSON_AddNumberToObject(h, "sector", sector);
    cJSON_AddStringToObject(h, "keyHex", hx);
    cJSON_AddStringToObject(h, "keyType", key_type);
    cJSON_AddItemToArray(hits, h);
    return true;
}

static int push_variant(const uint8_t base[6], int idx, uint8_t out[6])
{
    memcpy(out, base, 6);
    switch (idx) {
    case 0:
        return 0;
    case 1:
        out[5] ^= 0xFF;
        return 0;
    case 2:
        out[0] ^= 0xFF;
        return 0;
    case 3:
        out[5] ^= 0xAA;
        return 0;
    case 4:
        out[5] ^= 0x55;
        return 0;
    default: {
        int n = idx - 5;
        if (n >= 0 && n < 10) {
            out[5] = (uint8_t)((out[5] & 0xF0) | (uint8_t)n);
            return 0;
        }
    }
        return -1;
    }
}

static bool try_key_on_trailer(nfc_tag_info_t *tag, uint8_t trailer, const uint8_t key[6], bool key_b)
{
    nfc_mifare_key_t k;
    memcpy(k.key, key, 6);
    k.key_b = key_b;
    return nfc_mifare_authenticate_block(tag, trailer, &k) == ESP_OK;
}

/** Trailer block for MIFARE Classic sector index (0..15 for 1K, 0..39 for 4K). */
static uint8_t classic_trailer_for_sector(const nfc_tag_info_t *tag, uint8_t sec)
{
    if (nfc_tag_is_mifare_classic_4k(tag)) {
        if (sec <= 31) {
            return (uint8_t)(sec * 4 + 3);
        }
        if (sec <= 39) {
            return (uint8_t)(128 + (sec - 32) * 16 + 15);
        }
        return 0xFF;
    }
    return (uint8_t)(sec * 4 + 3);
}

cJSON *nfc_mifare_dictionary_attack(nfc_tag_info_t *tag, uint8_t sector_first, uint8_t sector_last,
                                    const uint8_t *extra, size_t extra_n, bool variations,
                                    int *attempts_out)
{
    int64_t t0_us = esp_timer_get_time();
    int attempts = 0;
    cJSON *root = cJSON_CreateObject();
    cJSON *hits = cJSON_CreateArray();
    if (!root || !hits) {
        cJSON_Delete(root);
        cJSON_Delete(hits);
        return NULL;
    }

    if (!tag) {
        cJSON_Delete(root);
        cJSON_Delete(hits);
        return NULL;
    }

    if (!nfc_tag_is_mifare_classic(tag)) {
        cJSON_AddStringToObject(root, "error", "not_classic_sak_hint");
        cJSON_AddItemToObject(root, "sectorHits", hits);
        if (attempts_out) {
            *attempts_out = 0;
        }
        return root;
    }

    uint8_t sector_max = nfc_tag_is_mifare_classic_4k(tag) ? 39 : 15;

    if (sector_first > sector_last) {
        uint8_t t = sector_first;
        sector_first = sector_last;
        sector_last = t;
    }
    if (sector_first > sector_max) {
        sector_first = sector_max;
    }
    if (sector_last > sector_max) {
        sector_last = sector_max;
    }

    for (int sec_i = sector_first; sec_i <= sector_last; sec_i++) {
        uint8_t sec = (uint8_t)sec_i;
        uint8_t trailer = classic_trailer_for_sector(tag, sec);
        if (trailer == 0xFF) {
            continue;
        }
        bool got = false;

        for (size_t bi = 0; bi < NBUILTIN && !got; bi++) {
            uint8_t trial[6];
            int maxv = variations ? MAX_VARIANTS_PER_KEY : 1;
            for (int vi = 0; vi < maxv; vi++) {
                if (push_variant(k_builtin[bi], vi, trial) != 0) {
                    break;
                }
                attempts++;
                if (try_key_on_trailer(tag, trailer, trial, false)) {
                    if (!add_sector_hit(hits, sec, trial, "A")) {
                        cJSON_Delete(root);
                        return NULL;
                    }
                    got = true;
                    break;
                }
                if (try_key_on_trailer(tag, trailer, trial, true)) {
                    if (!add_sector_hit(hits, sec, trial, "B")) {
                        cJSON_Delete(root);
                        return NULL;
                    }
                    got = true;
                    break;
                }
                if ((attempts % MAX_TRIES_BEFORE_WDT) == 0) {
                    esp_task_wdt_reset();
                    vTaskDelay(pdMS_TO_TICKS(1));
                }
            }
        }

        for (size_t ei = 0; ei < extra_n && !got; ei++) {
            const uint8_t *ek = extra + ei * 6;
            uint8_t trial[6];
            int maxv = variations ? MAX_VARIANTS_PER_KEY : 1;
            for (int vi = 0; vi < maxv; vi++) {
                if (push_variant(ek, vi, trial) != 0) {
                    break;
                }
                attempts++;
                if (try_key_on_trailer(tag, trailer, trial, false)) {
                    if (!add_sector_hit(hits, sec, trial, "A")) {
                        cJSON_Delete(root);
                        return NULL;
                    }
                    got = true;
                    break;
                }
                if (try_key_on_trailer(tag, trailer, trial, true)) {
                    if (!add_sector_hit(hits, sec, trial, "B")) {
                        cJSON_Delete(root);
                        return NULL;
                    }
                    got = true;
                    break;
                }
                if ((attempts % MAX_TRIES_BEFORE_WDT) == 0) {
                    esp_task_wdt_reset();
                    vTaskDelay(pdMS_TO_TICKS(1));
                }
            }
        }

        if (!got) {
            cJSON *h = cJSON_CreateObject();
            if (!h) {
                cJSON_Delete(root);
                return NULL;
            }
            cJSON_AddNumberToObject(h, "sector", sec);
            cJSON_AddBoolToObject(h, "miss", true);
            cJSON_AddItemToArray(hits, h);
        }
    }

    cJSON_AddItemToObject(root, "sectorHits", hits);
    cJSON_AddNumberToObject(root, "attempts", attempts);
    cJSON_AddNumberToObject(root, "durationMs", (double)((esp_timer_get_time() - t0_us) / 1000));
    cJSON_AddStringToObject(root, "note",
                           "Dictionary + bounded variants only — not exhaustive 48-bit keyspace.");
    if (attempts_out) {
        *attempts_out = attempts;
    }
    ESP_LOGI(TAG, "dictionary attack attempts=%d", attempts);
    return root;
}