#include "../include/Containers.h"

#include <cstdint>
#include <cstddef>
#include <stdexcept>

#ifndef __AVX2__
#include <immintrin.h>
#endif

using ht = HashTable;
using trackPtr = std::shared_ptr<Track>;


ht::HashTable( std::size_t capacity ) :ctrl_ ( capacity, EMPTY ), entries_ ( capacity ), size_( 0 ) {
    if  ( !capacity ) throw std::invalid_argument( "Capacity must be > 0" );
}

uint64_t ht::hashString( const std::string& key ) noexcept {
    uint64_t hash = 14695981039346656037ULL;
    for ( char c : key ) {
        hash ^= static_cast<uint8_t> ( c );
        hash *= 1099511628211ULL;
    }
    return hash;
}

uint8_t ht::fingerprint ( uint64_t hash ) noexcept {
    return static_cast<uint8_t> ( hash & 0x7F );  // Use lower 7 bits for fingerprint
}

std::ptrdiff_t ht::probe ( const std::string &key, uint64_t hash, bool insert ) const noexcept {
    const std::size_t capacity  = ctrl_.size();
    const uint8_t fp            = fingerprint ( hash );

#ifdef __AVX2__
    std::size_t group = ( hash % capacity ) & ~( GROUP_SIZE - 1 );  // Align to group boundary
    __m256i fpVec       = _mm256_set1_epi8 ( fp );
    __m256i emptyVec    = _mm256_set1_epi8 ( static_cast<char> ( EMPTY ) );

    std::ptrdiff_t firstTombstone = -1;

    for ( ;; ) {
        __m256i ctrlVec     = _mm256_loadu_si256 ( reinterpret_cast<const __m256i *> ( &ctrl_[ group ] ) );
        uint32_t matchMask  = _mm256_movemask_epi8 ( _mm256_cmpeq_epi8 ( ctrlVec, fpVec ) );
        uint32_t emptyMask  = _mm256_movemask_epi8 ( _mm256_cmpeq_epi8 ( ctrlVec, emptyVec ) );

        // check candidates
        for ( uint32_t mask = matchMask; mask != 0; mask &= ( mask - 1 ) ) {
            std::ptrdiff_t idx = ( group + static_cast<std::size_t> ( __builtin_ctz ( mask ) ) ) % capacity;

            if ( entries_[ idx ].hash == hash && entries_[ idx ].key == key ) {
                return static_cast<std::ptrdiff_t> ( idx );  // Found
            }
        }
        // record first tombstone for this group
        if ( insert && firstTombstone == -1 ) {
            for ( std::size_t i = 0; i < GROUP_SIZE; ++i ) {
                std::size_t idx = ( group + i ) % capacity;
                if ( ctrl_[ idx ] == DELETED ) {
                    firstTombstone = static_cast<std::ptrdiff_t> ( idx );
                    break;
                }
            }
        }

        if ( emptyMask ) {
            if ( !insert )
                return -1;  // Not found
            std::size_t slot = ( group + static_cast<std::size_t> ( __builtin_ctz ( emptyMask ) ) ) % capacity;
            return firstTombstone != -1 ? firstTombstone : static_cast<std::ptrdiff_t> ( slot );  // Insert here
        }
        group = ( group + GROUP_SIZE ) % capacity;  // Move to next group
    }
#else  // Fallback to scalar probing
    std::size_t group = ( hash % capacity ) & ~( GROUP_SIZE - 1 );  // Align to group boundary
    std::ptrdiff_t firstTombstone = -1;

    for ( ;; ) {
        bool foundEmpty = false;
        std::size_t emptyIdx = 0;

        for ( std::size_t i = 0; i < GROUP_SIZE; ++i ) {
            std::size_t idx = ( group + i ) % capacity;
            uint8_t c = ctrl_.at( idx );

            if ( c == EMPTY ) {
                if ( !foundEmpty ) {
                    foundEmpty  = true;
                    emptyIdx    = idx;
                }
            } else if ( c == DELETED ) {
                if ( insert && firstTombstone == -1 ) {
                    firstTombstone = static_cast<std::ptrdiff_t> ( idx );
                }
            } else if ( c == fp && entries_.at( idx ).hash == hash && entries_.at( idx ).key == key ) {
                return static_cast<std::ptrdiff_t> ( idx );  // Found
            }
        }

        if ( foundEmpty ) {
            if ( !insert )
                return -1;  // Not found
            return firstTombstone != -1 ? firstTombstone : static_cast<std::ptrdiff_t> ( emptyIdx );  // Insert here
        }
        group = ( group + GROUP_SIZE ) % capacity;  // Move to next group
    }  // end probeloop
#endif
}  // probe

void ht::ht_insert( const std::string& key, std::shared_ptr<Track> t ) {
    //double loadFactor = static_cast<double> ( size_ ) / static_cast<double> ( ctrl_.size());

    if ( loadFactor_ > MAX_LOAD ) rehash( ctrl_.size() * 2 );

    uint64_t hash = hashString ( key );
    auto idx = static_cast<std::size_t> ( probe ( key, hash, true ) );

    if ( ctrl_.at( idx ) == EMPTY || ctrl_.at( idx ) == DELETED ) {
        Entry tmp;
        tmp.key             = key;
        tmp.hash            = hash;
        tmp.track           = t;
        entries_.at( idx )  = std::move(tmp);
        ++size_;
    } else {
        entries_.at( idx ).track = t;
    }
    ctrl_.at( idx ) = fingerprint( hash );
}

std::shared_ptr<Track> ht::ht_lookup( const std::string& key ) const {
    uint64_t hash = hashString( key );
    auto idx = probe( key, hash, false );
    if ( idx != -1 ) return entries_.at( static_cast<std::size_t> ( idx )).track;
    return nullptr;
}

bool ht::ht_delete( const std::string& key ) {
    uint64_t hash = hashString( key );
    auto raw = probe( key, hash, false );

    if ( raw == -1 ) return false;

    auto idx = static_cast<size_t>( raw );
    ctrl_.at( idx )     = DELETED;
    entries_.at( idx )  = Entry();
    --size_;
    return true;
}

void ht::rehash( std::size_t newCapacity ) {
    ht newTable( newCapacity );
    for ( std::size_t i = 0; i < ctrl_.size(); ++i ) {
        if ( ctrl_.at( i ) != EMPTY && ctrl_.at( i ) != DELETED ) {
            const Entry& entry = entries_.at( i );
            newTable.ht_insert( entry.key, entry.track );
        }
    }
    ctrl_       = std::move( newTable.ctrl_ );
    entries_    = std::move( newTable.entries_ );
    size_       = newTable.size_;
}