#include "memory/simd_signature_scanner.hpp"

#include <immintrin.h>
#include <intrin.h>
#include <algorithm>
#include <cstring>
#include <cmath>

namespace memory {

bool SimdSignatureScanner::IsAvx2Supported() {
    int cpuInfo[4] = { 0 };
    __cpuid(cpuInfo, 0);
    int nIds = cpuInfo[0];
    if (nIds >= 7) {
        __cpuidex(cpuInfo, 7, 0);
        // Bit 5 của EBX biểu thị hỗ trợ AVX2
        return (cpuInfo[1] & (1 << 5)) != 0;
    }
    return false;
}

intptr_t SimdSignatureScanner::FindPatternAVX2(
    const uint8_t* data,
    size_t dataSize,
    const AobPattern& pattern
) {
    if (!pattern.Valid() || !data || dataSize < pattern.Size()) return -1;

    const size_t pSize = pattern.Size();

    // Tìm byte đầu tiên và byte cuối cùng không phải là wildcard để lập dual-byte filter
    size_t firstOff = 0;
    while (firstOff < pSize && pattern.wildcard[firstOff]) ++firstOff;
    if (firstOff >= pSize) return 0; // Toàn wildcard

    size_t lastOff = pSize - 1;
    while (lastOff > firstOff && pattern.wildcard[lastOff]) --lastOff;

    const uint8_t firstByte = pattern.bytes[firstOff];
    const uint8_t lastByte = pattern.bytes[lastOff];

    // Nếu CPU không hỗ trợ AVX2 hoặc buffer quá ngắn cho SIMD (cần >= 32 + pSize)
    if (!IsAvx2Supported() || dataSize < 32 + pSize) {
        return AobScanner::Find(data, dataSize, pattern);
    }

    const __m256i targetFirst = _mm256_set1_epi8(static_cast<char>(firstByte));
    const __m256i targetLast = _mm256_set1_epi8(static_cast<char>(lastByte));

    const size_t maxSimdIdx = dataSize - pSize;
    size_t i = 0;

    for (; i + 32 <= maxSimdIdx; i += 32) {
        const __m256i vFirst = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(data + i + firstOff));
        const __m256i vLast = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(data + i + lastOff));

        const __m256i eqFirst = _mm256_cmpeq_epi8(vFirst, targetFirst);
        const __m256i eqLast = _mm256_cmpeq_epi8(vLast, targetLast);
        const __m256i eqBoth = _mm256_and_si256(eqFirst, eqLast);

        uint32_t mask = static_cast<uint32_t>(_mm256_movemask_epi8(eqBoth));

        while (mask != 0) {
            unsigned long bitIdx = 0;
            _BitScanForward(&bitIdx, mask);
            mask &= (mask - 1); // Xóa bit vừa tìm

            const size_t candidate = i + bitIdx;
            if (candidate + pSize <= dataSize) {
                bool matched = true;
                for (size_t j = 0; j < pSize; ++j) {
                    if (!pattern.wildcard[j] && data[candidate + j] != pattern.bytes[j]) {
                        matched = false;
                        break;
                    }
                }
                if (matched) {
                    return static_cast<intptr_t>(candidate);
                }
            }
        }
    }

    // Xử lý nốt phần đuôi (tail) không đủ 32 byte
    for (; i <= maxSimdIdx; ++i) {
        if (data[i + firstOff] == firstByte && data[i + lastOff] == lastByte) {
            bool matched = true;
            for (size_t j = 0; j < pSize; ++j) {
                if (!pattern.wildcard[j] && data[i + j] != pattern.bytes[j]) {
                    matched = false;
                    break;
                }
            }
            if (matched) return static_cast<intptr_t>(i);
        }
    }

    return -1;
}

std::vector<intptr_t> SimdSignatureScanner::FindAllHitsAVX2(
    const uint8_t* data,
    size_t dataSize,
    const AobPattern& pattern,
    size_t maxHits
) {
    std::vector<intptr_t> hits;
    if (!pattern.Valid() || !data || dataSize < pattern.Size() || maxHits == 0) return hits;

    const size_t pSize = pattern.Size();
    size_t firstOff = 0;
    while (firstOff < pSize && pattern.wildcard[firstOff]) ++firstOff;
    if (firstOff >= pSize) return hits;

    size_t lastOff = pSize - 1;
    while (lastOff > firstOff && pattern.wildcard[lastOff]) --lastOff;

    const uint8_t firstByte = pattern.bytes[firstOff];
    const uint8_t lastByte = pattern.bytes[lastOff];

    if (!IsAvx2Supported() || dataSize < 32 + pSize) {
        for (size_t i = 0; i + pSize <= dataSize && hits.size() < maxHits; ++i) {
            if (data[i + firstOff] != firstByte || data[i + lastOff] != lastByte) continue;
            bool matched = true;
            for (size_t j = 0; j < pSize; ++j) {
                if (!pattern.wildcard[j] && data[i + j] != pattern.bytes[j]) {
                    matched = false;
                    break;
                }
            }
            if (matched) hits.push_back(static_cast<intptr_t>(i));
        }
        return hits;
    }

    const __m256i targetFirst = _mm256_set1_epi8(static_cast<char>(firstByte));
    const __m256i targetLast = _mm256_set1_epi8(static_cast<char>(lastByte));

    const size_t maxSimdIdx = dataSize - pSize;
    size_t i = 0;

    for (; i + 32 <= maxSimdIdx && hits.size() < maxHits; i += 32) {
        const __m256i vFirst = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(data + i + firstOff));
        const __m256i vLast = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(data + i + lastOff));

        const __m256i eqBoth = _mm256_and_si256(_mm256_cmpeq_epi8(vFirst, targetFirst),
                                                 _mm256_cmpeq_epi8(vLast, targetLast));

        uint32_t mask = static_cast<uint32_t>(_mm256_movemask_epi8(eqBoth));

        while (mask != 0 && hits.size() < maxHits) {
            unsigned long bitIdx = 0;
            _BitScanForward(&bitIdx, mask);
            mask &= (mask - 1);

            const size_t candidate = i + bitIdx;
            if (candidate + pSize <= dataSize) {
                bool matched = true;
                for (size_t j = 0; j < pSize; ++j) {
                    if (!pattern.wildcard[j] && data[candidate + j] != pattern.bytes[j]) {
                        matched = false;
                        break;
                    }
                }
                if (matched) {
                    hits.push_back(static_cast<intptr_t>(candidate));
                }
            }
        }
    }

    for (; i <= maxSimdIdx && hits.size() < maxHits; ++i) {
        if (data[i + firstOff] == firstByte && data[i + lastOff] == lastByte) {
            bool matched = true;
            for (size_t j = 0; j < pSize; ++j) {
                if (!pattern.wildcard[j] && data[i + j] != pattern.bytes[j]) {
                    matched = false;
                    break;
                }
            }
            if (matched) hits.push_back(static_cast<intptr_t>(i));
        }
    }

    return hits;
}

size_t SimdSignatureScanner::ScanFloatRangeAVX2(
    const uint8_t* buffer,
    size_t bufferSize,
    uintptr_t baseAddress,
    float fmin,
    float fmax,
    std::vector<std::pair<uintptr_t, float>>& outFloats
) {
    if (!buffer || bufferSize < 4) return 0;

    const size_t initialCount = outFloats.size();

    if (!IsAvx2Supported() || bufferSize < 32) {
        // Scalar fallback
        const size_t floatCount = bufferSize / 4;
        const float* floats = reinterpret_cast<const float*>(buffer);
        for (size_t i = 0; i < floatCount; ++i) {
            const float f = floats[i];
            if (std::isfinite(f) && f >= fmin && f <= fmax) {
                outFloats.emplace_back(baseAddress + i * 4, f);
            }
        }
        return outFloats.size() - initialCount;
    }

    const __m256 vmin = _mm256_set1_ps(fmin);
    const __m256 vmax = _mm256_set1_ps(fmax);

    // Mỗi bước quét 8 float = 32 bytes
    const size_t simdBytes = (bufferSize / 32) * 32;
    size_t byteOffset = 0;

    for (; byteOffset < simdBytes; byteOffset += 32) {
        const __m256 vf = _mm256_loadu_ps(reinterpret_cast<const float*>(buffer + byteOffset));

        // So sánh: vf >= vmin VÀ vf <= vmax
        const __m256 cmpGe = _mm256_cmp_ps(vf, vmin, _CMP_GE_OQ);
        const __m256 cmpLe = _mm256_cmp_ps(vf, vmax, _CMP_LE_OQ);
        const __m256 inRange = _mm256_and_ps(cmpGe, cmpLe);

        int mask = _mm256_movemask_ps(inRange);
        if (mask != 0) {
            const float* chunk = reinterpret_cast<const float*>(buffer + byteOffset);
            for (int k = 0; k < 8; ++k) {
                if ((mask & (1 << k)) != 0) {
                    const float f = chunk[k];
                    if (std::isfinite(f)) {
                        outFloats.emplace_back(baseAddress + byteOffset + k * 4, f);
                    }
                }
            }
        }
    }

    // Xử lý nốt các float còn lại ở tail
    const size_t remainingFloats = (bufferSize - byteOffset) / 4;
    const float* tailFloats = reinterpret_cast<const float*>(buffer + byteOffset);
    for (size_t i = 0; i < remainingFloats; ++i) {
        const float f = tailFloats[i];
        if (std::isfinite(f) && f >= fmin && f <= fmax) {
            outFloats.emplace_back(baseAddress + byteOffset + i * 4, f);
        }
    }

    return outFloats.size() - initialCount;
}

bool SimdSignatureScanner::ScanStaticRootsBatched(
    const uint8_t* peSectionData,
    size_t sectionSize,
    uintptr_t sectionBaseRva,
    std::vector<SignatureRoot>& roots,
    float* outDurationMs
) {
    if (!peSectionData || sectionSize == 0 || roots.empty()) return false;

    const auto t0 = std::chrono::high_resolution_clock::now();
    bool allFound = true;

    for (auto& root : roots) {
        const intptr_t hit = FindPatternAVX2(peSectionData, sectionSize, root.pattern);
        if (hit >= 0) {
            const size_t ripOff = root.ripOffset;
            if (static_cast<size_t>(hit) + ripOff + 4 <= sectionSize) {
                const int32_t disp = *reinterpret_cast<const int32_t*>(peSectionData + hit + ripOff);
                root.resolvedAddress = sectionBaseRva + hit + ripOff + 4 + disp;
            } else {
                root.resolvedAddress = 0;
                allFound = false;
            }
        } else {
            root.resolvedAddress = 0;
            allFound = false;
        }
    }

    const auto t1 = std::chrono::high_resolution_clock::now();
    if (outDurationMs) {
        *outDurationMs = std::chrono::duration<float, std::milli>(t1 - t0).count();
    }

    return allFound;
}

} // namespace memory