import os
import sys

def update_visual_mods_hpp():
    path = "c:/Projects/AutoPOE2/src/core/memory/visual_mods.hpp"
    with open(path, "r", encoding="utf-8") as f:
        content = f.read()

    # 1. Includes
    old_inc = '#include <cmath>'
    new_inc = '#include <cmath>\n#include <string>\n#include <vector>'
    if '<string>' not in content:
        content = content.replace(old_inc, new_inc, 1)

    # 2. Add CodePatchRecord struct before kVisualZoomMin
    code_patch_struct = '''// ==========================================================
// Bản ghi Code Patch trong phân vùng .text (Native Smooth Zoom)
// ==========================================================
struct CodePatchRecord {
    std::string name;
    uintptr_t address = 0;
    std::vector<uint8_t> originalBytes;
    std::vector<uint8_t> patchedBytes;
    bool isPatched = false;

    bool IsValid() const {
        return address != 0 && !originalBytes.empty() && originalBytes.size() == patchedBytes.size();
    }
};

'''
    if 'struct CodePatchRecord' not in content:
        target = 'inline constexpr float kVisualZoomMin = 1.0f;'
        content = content.replace(target, code_patch_struct + target, 1)

    # 3. Add methods to VisualModsEngine
    methods = '''    // ==========================================================
    // Code Patching trong phân vùng .text (Native 144 FPS, Zero Jitter)
    // ==========================================================
    bool ApplyZoomCodePatch(float maxMultiplier = kVisualZoomMax);
    bool RestoreZoomCodePatch();
    bool IsCodePatchActive() const { return m_zoomCodePatch.isPatched; }
    uintptr_t ZoomClampAddress() const { return m_zoomCodePatch.address; }
    void SetZoomClampAddress(uintptr_t addr) { m_zoomCodePatch.address = addr; }
    const CodePatchRecord& ZoomCodePatchRecord() const { return m_zoomCodePatch; }

    // Khôi phục Camera về mặc định ban đầu'''
    if 'ApplyZoomCodePatch' not in content:
        content = content.replace('    // Khôi phục Camera về mặc định ban đầu', methods, 1)

    # 4. Add member variable m_zoomCodePatch
    member = '''    // Quản lý Code Patching trong .text
    CodePatchRecord m_zoomCodePatch;

    // Lưu trữ giá trị mặc định gốc của game để khôi phục sạch'''
    if 'm_zoomCodePatch;' not in content:
        content = content.replace('    // Lưu trữ giá trị mặc định gốc của game để khôi phục sạch', member, 1)

    with open(path, "w", encoding="utf-8") as f:
        f.write(content)
    print("visual_mods.hpp updated successfully.")

def update_visual_mods_cpp():
    path = "c:/Projects/AutoPOE2/src/core/memory/visual_mods.cpp"
    with open(path, "r", encoding="utf-8") as f:
        content = f.read()

    # 1. Update LooksLikeCameraFields typical calculation
    old_typical = '''    const bool typical = (maxD >= 12.0f && maxD <= 120.0f && cur >= 6.0f && cur <= 250.0f);'''
    new_typical = '''    const float maxAllowedDistance = unlockedKeepAlive ? 400.0f : 120.0f;
    const bool typical = (maxD >= 12.0f && maxD <= maxAllowedDistance && cur >= 6.0f && cur <= 500.0f);'''
    if old_typical in content:
        content = content.replace(old_typical, new_typical, 1)

    # 2. Add ApplyZoomCodePatch and RestoreZoomCodePatch implementations
    code_patch_impl = '''bool VisualModsEngine::ApplyZoomCodePatch(float maxMultiplier) {
    if (!m_reader.IsAttached()) return false;

    const float cap = ClampVisualZoom(maxMultiplier);
    m_zoomCap = cap;

    // 1. Quét tìm clamp trong .text nếu chưa có địa chỉ
    if (m_zoomCodePatch.address == 0) {
        ModuleInfo mainMod;
        const uint32_t pid = m_reader.AttachedPid();
        if (ResolveMainModule(pid, m_reader, mainMod) && mainMod.base != 0 && mainMod.size != 0) {
            // POE2 clamp pattern: cmp byte ptr [rax+2Fh], 0
            const AobPattern clampPat = AobPattern::Parse("80 78 2F 00");
            uintptr_t match = AobScanner::ScanModule(m_reader, mainMod, clampPat);
            if (match != 0) {
                m_zoomCodePatch.address = match;
                m_zoomCodePatch.name = "CameraClampBypass_cmp2F";
                std::cout << "[VisualMods] Tim thay Zoom Clamp tai .text @ 0x"
                          << std::hex << match << std::dec << std::endl;
            }
        }
    }

    if (m_zoomCodePatch.address != 0 && !m_zoomCodePatch.isPatched) {
        constexpr size_t kPatchSize = 4;
        std::vector<uint8_t> orig(kPatchSize);
        if (m_reader.Read(m_zoomCodePatch.address, orig.data(), kPatchSize)) {
            m_zoomCodePatch.originalBytes = orig;
            // Vá thành 4 NOPs (0x90) để bỏ qua clamp
            m_zoomCodePatch.patchedBytes = {0x90, 0x90, 0x90, 0x90};
            if (m_reader.Write(m_zoomCodePatch.address, m_zoomCodePatch.patchedBytes.data(), kPatchSize)) {
                m_zoomCodePatch.isPatched = true;
                std::cout << "[VisualMods] ApplyZoomCodePatch: Da va clamp tai 0x"
                          << std::hex << m_zoomCodePatch.address << std::dec
                          << " (Native 144 FPS smooth zoom active)." << std::endl;
            } else {
                std::cerr << "[VisualMods] ApplyZoomCodePatch: Ghi byte patch that bai tai 0x"
                          << std::hex << m_zoomCodePatch.address << std::dec << std::endl;
            }
        }
    }

    // Đồng thời cập nhật limits trên heap nếu camera instance đã sẵn sàng
    if (m_cameraAddr != 0) {
        EnsureOrigCameraSaved();
        m_zoomActive = true;
        m_capUnlockOnly = true;
        WriteMaxAndZFar(cap);
    } else if (m_zoomCodePatch.isPatched) {
        m_zoomActive = true;
        m_capUnlockOnly = true;
    }

    return m_zoomCodePatch.isPatched || m_zoomActive;
}

bool VisualModsEngine::RestoreZoomCodePatch() {
    if (!m_zoomCodePatch.isPatched || m_zoomCodePatch.address == 0) return true;

    if (!m_zoomCodePatch.originalBytes.empty()) {
        if (m_reader.Write(m_zoomCodePatch.address,
                           m_zoomCodePatch.originalBytes.data(),
                           m_zoomCodePatch.originalBytes.size())) {
            m_zoomCodePatch.isPatched = false;
            std::cout << "[VisualMods] RestoreZoomCodePatch: Da hoan nguyen byte goc tai 0x"
                      << std::hex << m_zoomCodePatch.address << std::dec
                      << " sach se 100%." << std::endl;
            return true;
        } else {
            std::cerr << "[VisualMods] RestoreZoomCodePatch: Hoan nguyen that bai tai 0x"
                      << std::hex << m_zoomCodePatch.address << std::dec << std::endl;
            return false;
        }
    }
    m_zoomCodePatch.isPatched = false;
    return true;
}

'''
    if 'VisualModsEngine::ApplyZoomCodePatch' not in content:
        content = content.replace('bool VisualModsEngine::RestoreCamera() {', code_patch_impl + 'bool VisualModsEngine::RestoreCamera() {', 1)

    # 3. Update PanicRestoreAll to call RestoreZoomCodePatch
    old_panic = '''void VisualModsEngine::PanicRestoreAll() {
    if (m_zoomActive) {
        RestoreCamera();
    }
    if (m_fogDisabled) {
        RestoreFog();
    }
}'''
    new_panic = '''void VisualModsEngine::PanicRestoreAll() {
    RestoreZoomCodePatch();
    if (m_zoomActive) {
        RestoreCamera();
    }
    if (m_fogDisabled) {
        RestoreFog();
    }
}'''
    if old_panic in content:
        content = content.replace(old_panic, new_panic, 1)

    with open(path, "w", encoding="utf-8") as f:
        f.write(content)
    print("visual_mods.cpp updated successfully.")

def update_visual_tool_main():
    path = "c:/Projects/AutoPOE2/src/visual_tool/main.cpp"
    with open(path, "r", encoding="utf-8") as f:
        content = f.read()

    # In EngineLoop:
    # 1. Check CameraLooksValid with module parameters and debounce
    old_check = '''            const bool camInvalid = engine->CameraAddress() != 0 && !engine->CameraLooksValid();
            if (camInvalid) {
                std::cout << "[VisualTool] Camera heap khong con hop le — AutoDetect lai.\\n";
                engine->SetCameraAddress(0);
                lastDetect = now - std::chrono::seconds(4);
            }'''
    new_check = '''            static int camInvalidStreak = 0;
            const bool camInvalid = engine->CameraAddress() != 0 && !engine->CameraLooksValid();
            if (camInvalid) {
                if (++camInvalidStreak >= 10) {
                    std::cout << "[VisualTool] Camera heap khong con hop le (10 ticks) — AutoDetect lai.\\n";
                    engine->SetCameraAddress(0);
                    lastDetect = now - std::chrono::seconds(4);
                    camInvalidStreak = 0;
                }
            } else {
                camInvalidStreak = 0;
            }'''
    if old_check in content:
        content = content.replace(old_check, new_check, 1)

    # 2. When attached, also trigger ApplyZoomCodePatch
    old_zoom_apply = '''            if (std::abs(engine->ZoomCap() - cap) > 0.01f || !engine->IsCapUnlockActive()) {
                engine->UnlockZoomCap(cap);
            }'''
    new_zoom_apply = '''            if (!engine->IsCodePatchActive()) {
                engine->ApplyZoomCodePatch(cap);
            }
            if (std::abs(engine->ZoomCap() - cap) > 0.01f || !engine->IsCapUnlockActive()) {
                engine->UnlockZoomCap(cap);
            }'''
    if old_zoom_apply in content:
        content = content.replace(old_zoom_apply, new_zoom_apply, 1)

    with open(path, "w", encoding="utf-8") as f:
        f.write(content)
    print("main.cpp updated successfully.")

def update_test_core():
    path = "c:/Projects/AutoPOE2/tests/test_core.cpp"
    with open(path, "r", encoding="utf-8") as f:
        content = f.read()

    test_case = '''// ==========================================================
// 40d. Kiểm thử VisualModsEngine: Code Patching (.text clamp bypass & 1-frame restore)
// ==========================================================
void TestVisualToolCodePatching() {
    std::cout << "[Test 40d] Visual Tool: Code Patching (.text clamp bypass & 1-frame restore)..." << std::endl;

    SimulatedMemoryReader sim;
    CHECK(sim.Initialize(), "sim.Initialize 40d");
    CHECK(sim.Attach(game_layout::kSimulatedPid), "sim.Attach 40d");

    VisualModsEngine visual(sim);

    // Cấy đoạn mã clamp POE2 giả lập trong buffer: 80 78 2F 00 (cmp byte ptr [rax+2Fh], 0)
    const uintptr_t clampAddr = sim.BaseAddress() + 0x150000;
    const uint8_t origClampBytes[4] = { 0x80, 0x78, 0x2F, 0x00 };
    CHECK(sim.Write(clampAddr, origClampBytes, 4), "plant clamp bytes in .text");

    visual.SetZoomClampAddress(clampAddr);
    CHECK(visual.ZoomClampAddress() == clampAddr, "clamp address registered");
    CHECK(!visual.IsCodePatchActive(), "ban dau code patch chua active");

    // 1. Áp dụng Code Patch
    CHECK(visual.ApplyZoomCodePatch(8.0f), "ApplyZoomCodePatch thanh cong");
    CHECK(visual.IsCodePatchActive(), "code patch da active");

    // Đọc byte tại địa chỉ clamp để xác nhận đã được vá 4 NOPs (0x90)
    uint8_t readBytes[4] = {};
    CHECK(sim.Read(clampAddr, readBytes, 4), "doc byte sau khi va");
    CHECK(readBytes[0] == 0x90 && readBytes[1] == 0x90 && readBytes[2] == 0x90 && readBytes[3] == 0x90,
          "byte da duoc va thanh 4 NOPs (0x90)");

    // 2. Hoàn nguyên Code Patch
    CHECK(visual.RestoreZoomCodePatch(), "RestoreZoomCodePatch thanh cong");
    CHECK(!visual.IsCodePatchActive(), "code patch da ngat");

    CHECK(sim.Read(clampAddr, readBytes, 4), "doc byte sau khi hoan nguyen");
    CHECK(readBytes[0] == 0x80 && readBytes[1] == 0x78 && readBytes[2] == 0x2F && readBytes[3] == 0x00,
          "byte goc da duoc phuc hoi sach se 100%");

    // 3. Kiểm chứng PanicRestoreAll hoàn nguyên cả Code Patch
    CHECK(visual.ApplyZoomCodePatch(8.0f), "ApplyZoomCodePatch lan 2");
    CHECK(visual.IsCodePatchActive(), "patch active truoc panic");
    visual.PanicRestoreAll();
    CHECK(!visual.IsCodePatchActive(), "Panic phai tat code patch");
    CHECK(sim.Read(clampAddr, readBytes, 4), "doc byte sau panic");
    CHECK(readBytes[0] == 0x80 && readBytes[1] == 0x78 && readBytes[2] == 0x2F && readBytes[3] == 0x00,
          "Panic da hoan nguyen byte goc 100%");

    std::cout << "  -> Visual Tool Code Patching (.text bypass & clean restore) OK" << std::endl;
}

'''

    if 'TestVisualToolCodePatching' not in content:
        content = content.replace('void TestVisualAutoDetectRelaxedAndAob() {', test_case + 'void TestVisualAutoDetectRelaxedAndAob() {', 1)
        content = content.replace('    TestVisualAutoDetectRelaxedAndAob();', '    TestVisualAutoDetectRelaxedAndAob();\n    TestVisualToolCodePatching();', 1)

    with open(path, "w", encoding="utf-8") as f:
        f.write(content)
    print("test_core.cpp updated successfully.")

if __name__ == "__main__":
    update_visual_mods_hpp()
    update_visual_mods_cpp()
    update_visual_tool_main()
    update_test_core()