import marshal
import mmap
import os
import os.path
import sys
import logging
import threading
import queue
# from profilehooks import profile
from cog.cache import Cache
import xxhash

RECORD_SEP = b'\xFD'
UNIT_SEP = b'\xAC'


class TableMeta:
    __slots__ = ('name', 'namespace', 'db_instance_id', 'column_mode')

    def __init__(self, name, namespace, db_instance_id, column_mode):
        self.name = name
        self.namespace = namespace
        self.db_instance_id = db_instance_id
        self.column_mode = column_mode


class Table:

    def __init__(self, name, namespace, db_instance_id, config, column_mode=False, shared_cache=None,
                 flush_interval=1):
        self.logger = logging.getLogger('table')
        self.config = config
        self.shared_cache = shared_cache
        self.flush_interval = flush_interval
        self.table_meta = TableMeta(name, namespace, db_instance_id, column_mode)
        self.indexer = self.__create_indexer()
        self.store = self.__create_store(shared_cache)

    def __create_indexer(self):
        return Indexer(self.table_meta, self.config, self.logger)

    def __create_store(self, shared_cache):
        return Store(self.table_meta, self.config, self.logger, shared_cache=shared_cache,
                     flush_interval=self.flush_interval)

    def sync(self):
        """Force flush pending writes to disk."""
        self.store.sync()

    def close(self):
        self.indexer.close()
        self.store.close()
        self.logger.info("closed table: " + self.table_meta.name)


class Record:
    '''
    Record is the basic unit of storage in cog.
    value_type: s - string, l - list, u - set
    '''
    __slots__ = ('key', 'value', 'tombstone', 'store_position', 'key_link', 'value_link', 'value_type')
    
    RECORD_LINK_LEN = 16
    RECORD_LINK_NULL = -1
    VALUE_LINK_NULL = -1

    def __init__(self, key, value, tombstone='0', store_position=None, value_type="s", key_link=-1, value_link=-1):
        self.key = key
        self.value = value
        self.tombstone = tombstone
        self.store_position = store_position
        self.key_link = key_link
        self.value_link = value_link
        self.value_type = value_type

    def set_store_position(self, pos):
        if type(pos) is not int:
            raise ValueError("store position must be int but provided : " + str(pos))
        self.store_position = pos

    def set_key_link(self, pos):
        self.key_link = pos

    def set_value_link(self, pos):
        self.value_link = pos

    def set_value(self, value):
        self.value = value

    def is_equal_val(self, other_record):
        return self.key == other_record.key and self.value == other_record.value

    def get_kv_tuple(self):
        return self.key, self.value

    def serialize(self):
        return marshal.dumps((self.key, self.value))

    def marshal(self):
        key_link_bytes = str(self.key_link).encode().rjust(Record.RECORD_LINK_LEN)
        serialized = self.serialize()
        m_record = key_link_bytes \
                   + self.tombstone.encode() \
                   + self.value_type.encode() \
                   + str(len(serialized)).encode() \
                   + UNIT_SEP \
                   + serialized
        if self.value_type == "l" or self.value_type == "u":
            if self.value_link is not None:
                m_record += str(self.value_link).encode()
        m_record += RECORD_SEP
        return m_record

    def is_empty(self):
        return self.key is None and self.value is None

    def __str__(self):
        return "key: {}, value: {}, tombstone: {}, store_position: {}, key_link: {}, value_link: {}, value_type: {}".format(
            self.key, self.value, self.tombstone, self.store_position, self.key_link, self.value_link, self.value_type)

    @classmethod
    def __read_until(cls, start, sbytes, separtor=UNIT_SEP):
        buff = b''
        i = 0  # default
        for i in range(start, len(sbytes)):
            s_byte = sbytes[i: i + 1]
            if s_byte == separtor:
                break
            buff += s_byte
        return buff, i

    @classmethod
    def unmarshal(cls, store_bytes):
        """reads from bytes and creates object
        """
        base_pos = 0
        key_link = int(store_bytes[base_pos: base_pos + Record.RECORD_LINK_LEN])
        next_base_pos = Record.RECORD_LINK_LEN
        tombstone = store_bytes[next_base_pos:next_base_pos + 1].decode()
        value_type = store_bytes[next_base_pos + 1: next_base_pos + 2].decode()
        value_len, end_pos = cls.__read_until(next_base_pos + 2, store_bytes)
        value_len = int(value_len.decode())
        value = store_bytes[end_pos + 1: end_pos + 1 + value_len]
        record = marshal.loads(value)

        value_link = Record.VALUE_LINK_NULL
        if value_type == 'l' or value_type == 'u':
            value_link, end_pos = cls.__read_until(end_pos + value_len + 1, store_bytes, RECORD_SEP)
            value_link = int(value_link.decode())
        return cls(record[0], record[1], tombstone, store_position=None, value_type=value_type, key_link=key_link,
                   value_link=value_link)

    @classmethod
    def __load_value(cls, store_pointer, val_list, store):
        """loads value from the store"""
        while store_pointer != Record.VALUE_LINK_NULL:
            rec = Record.unmarshal(store.read(store_pointer))
            if rec.value_type == 'l':
                val_list.append(rec.value)
            else:
                val_list.add(rec.value)
            store_pointer = rec.value_link
        return val_list

    @classmethod
    # @profile
    def load_from_store(cls, position: int, store):
        record = cls.unmarshal(store.read(position))
        values = None
        if record.value_type == 'l':
            values = cls.__load_value(record.value_link, [record.value], store)
        elif record.value_type == 'u':
            values = cls.__load_value(record.value_link, {record.value}, store)
        if values is not None:
            record.set_value(values)
        return record


class Index:

    def __init__(self, table_meta, config, logger, index_id=0):
        self.logger = logging.getLogger('index')
        self.table = table_meta
        self.config = config
        self.name = self.config.cog_index(table_meta.namespace, table_meta.name, table_meta.db_instance_id, index_id)
        self.empty_block = '-1'.zfill(self.config.INDEX_BLOCK_LEN).encode()
        if not os.path.exists(self.name):
            self.logger.info("creating index...")
            f = open(self.name, 'wb+')
            i = 0
            e_blocks = []
            while i < config.INDEX_CAPACITY:
                e_blocks.append(self.empty_block)
                i += 1
            f.write(b''.join(e_blocks))
            self.file_limit = f.tell()
            f.close()
            self.logger.info("new index with capacity" + str(config.INDEX_CAPACITY) + "created: " + self.name)
        else:
            self.logger.info("Index: "+self.name+" already exists.")

        self.db = open(self.name, 'r+b')
        self.db_mem = mmap.mmap(self.db.fileno(), 0)
        self._closed = False

    def close(self):
        if self._closed:
            return
        self._closed = True
        self.db_mem.flush()
        self.db_mem.close()
        self.db.close()

    def get_index_key(self, int_store_position):
        return str(int_store_position).encode().rjust(self.config.INDEX_BLOCK_LEN)

    # @profile
    def put(self, key, store_position, store):
        """
        key chain
        :param key:
        :param store_position:
        :param store:
        :return:
        """

        """
        k5 -> k4 -> k3 -> k2 -> k1
        add: k6
        k6 -> k5 -> k4 -> k3 -> k2 -> k1
        add/update: k4
        1. k4 -> k6 -> k5 -> k4 -> k3 -> k2 -> k1
        2. k4 -> k6 -> k5 -> k3 -> k2 -> k1

        """
        orig_position, orig_hash = self.get_index(key)
        index_value = self.db_mem[orig_position: orig_position + self.config.INDEX_BLOCK_LEN]
        self.logger.debug('writing : ' + str(key) + ' current data at store position: ' + str(index_value))
        if index_value == self.empty_block:
            # First record in the key bucket, point next link to null
            store.update_record_link_inplace(store_position, Record.RECORD_LINK_NULL)
            key_link = store_position
            # write the store position to the index
            self.db_mem[orig_position: orig_position + self.config.INDEX_BLOCK_LEN] = self.get_index_key(store_position)
        else:
            # there are records in the bucket
            # read existing record from the store - use unmarshal, not load_from_store (O(1) vs O(n))
            existing_record = Record.unmarshal(store.read(int(index_value)))
            existing_record.set_store_position(int(index_value))

            if existing_record.key == key:
                # the record at the top of the bucket has the same key, update the record in place.
                # a new entry has been made to the store, update pre and next links.
                store.update_record_link_inplace(store_position, int(existing_record.key_link))
                key_link = int(existing_record.key_link)
            else:
                # this is hash collision.
                # the record at the top of the bucket has a different key, add new record to the top of the bucket.
                # set next link to the record at the top of the bucket, prev is null by default, no need to set.
                store.update_record_link_inplace(store_position, existing_record.store_position)
                key_link = existing_record.store_position

                # check if this record exists in the bucket, if yes remove pointer.
                prev_record = None
                while existing_record.key_link != Record.RECORD_LINK_NULL:
                    # Use unmarshal (O(1)) instead of load_from_store (O(n))
                    existing_record = Record.unmarshal(store.read(existing_record.key_link))
                    existing_record.set_store_position(existing_record.key_link)
                    if existing_record.key == key and prev_record is not None:
                        """
                        if same key found in bucket, update previous record in chain to point to key_link of this record
                        prev_rec -> current rec.key_link
                        curr_rec will not be linked in the bucket anymore.
                        """
                        # update in place the key link pointer of pervios record, ! need to add fixed length padding.
                        store.update_record_link_inplace(prev_record.store_position, existing_record.key_link)
                        key_link = existing_record.key_link
                    prev_record = existing_record

        self.db_mem[orig_position: orig_position + self.config.INDEX_BLOCK_LEN] = self.get_index_key(store_position)
        return key_link

    def get_index(self, key):
        num = cog_hash(key, self.config.INDEX_CAPACITY) % ((sys.maxsize + 1) * 2)
        self.logger.debug("hash for: " + key + " : " + str(num))
        # there may be diff when using mem slice vs write (+1 needed)
        index = (self.config.INDEX_BLOCK_LEN *
                 (max((num % self.config.INDEX_CAPACITY) - 1, 0)))
        self.logger.debug("offset : " + key + " : " + str(index))
        return index, num

    def cog_hash(self, string):
        return xxhash.xxh32(string, seed=2).intdigest() % self.config.INDEX_CAPACITY

    # @profile
    def get(self, key, store):
        self.logger.debug("GET: Reading index: " + self.name)
        index_position, raw_hash = self.get_index(key)
        data_at_index_position = self.db_mem[index_position:index_position + self.config.INDEX_BLOCK_LEN]
        if data_at_index_position == self.empty_block:
            return None
        data_at_index_position = int(data_at_index_position)
        record = Record.load_from_store(data_at_index_position, store)
        record.set_store_position(data_at_index_position)
        self.logger.debug("read record " + str(record))

        if record.key == key:
            return record
        else:
            while record.key_link != Record.RECORD_LINK_NULL:
                self.logger.debug("record.key_link: " + str(record.key_link))
                record = Record.load_from_store(record.key_link, store)
                record.set_store_position(record.key_link)
                if record.key == key:
                    return record
        return None

    def get_head_only(self, key, store):
        """
        Get only the head record without traversing the value chain.
        This is O(1) compared to get() which is O(n) for multi-value keys.
        
        Returns: (record, store_position) or (None, None)
        """
        index_position, raw_hash = self.get_index(key)
        data_at_index_position = self.db_mem[index_position:index_position + self.config.INDEX_BLOCK_LEN]
        if data_at_index_position == self.empty_block:
            return None, None
        store_position = int(data_at_index_position)
        # Only unmarshal, don't load value chain
        record = Record.unmarshal(store.read(store_position))
        record.set_store_position(store_position)
        
        if record.key == key:
            return record, store_position
        else:
            # Hash collision - follow key_link chain
            while record.key_link != Record.RECORD_LINK_NULL:
                store_position = record.key_link
                record = Record.unmarshal(store.read(store_position))
                record.set_store_position(store_position)
                if record.key == key:
                    return record, store_position
        return None, None

    '''
        Iterates through all records in the index, following key_link chains for hash collisions.
    '''

    def scanner(self, store):
        scan_cursor = 0
        while True:
            data_at_position = self.db_mem[scan_cursor:scan_cursor + self.config.INDEX_BLOCK_LEN]
            if len(data_at_position) == 0:  # EOF index
                self.logger.info("Index EOF reached! Scan terminated.")
                return
            if data_at_position == self.empty_block:
                scan_cursor += self.config.INDEX_BLOCK_LEN
                self.logger.debug("GET: skipping empty block during iteration.")
                continue
            
            # Load head record and follow key_link chain to get all records in this bucket
            store_position = int(data_at_position)
            while store_position != Record.RECORD_LINK_NULL:
                record = Record.load_from_store(store_position, store)
                if record is None:  # EOF store
                    self.logger.error("Store EOF reached! Iteration terminated.")
                    return
                yield Record(record.key, record.value, record.tombstone)
                store_position = record.key_link
            
            scan_cursor += self.config.INDEX_BLOCK_LEN

    def delete(self, key, store):
        """
               k5 -> k4 -> k3 -> k2 -> k1
               del: k3
               k6 -> k5 -> k4 -> k2 -> k1

        """
        self.logger.debug("GET: Reading index: " + self.name)
        index_position, raw_hash = self.get_index(key)

        data_at_index_position = self.db_mem[index_position:index_position + self.config.INDEX_BLOCK_LEN]
        if data_at_index_position == self.empty_block:
            return False

        data_at_index_position = int(data_at_index_position)
        record = Record.load_from_store(data_at_index_position, store)
        record.set_store_position(data_at_index_position)
        self.logger.debug("read record " + str(record))
        if record.key == key:
            """delete bucket => map hash table to empty block, or point to next in chain"""
            if record.key_link != Record.RECORD_LINK_NULL:
                # Point index to next record in chain
                self.db_mem[index_position:index_position + self.config.INDEX_BLOCK_LEN] = self.get_index_key(record.key_link)
            else:
                # No more records in chain, clear the bucket
                self.db_mem[index_position:index_position + self.config.INDEX_BLOCK_LEN] = self.empty_block
            return True
        else:
            """search bucket"""
            prev_record = record  # Initialize to the head record
            while record.key_link != Record.RECORD_LINK_NULL:
                next_record = Record.load_from_store(record.key_link, store)
                next_record.set_store_position(record.key_link)
                if next_record.key == key:
                    """
                    if same key found in bucket, update previous record in chain to point to key_link of this record
                    prev_rec -> current rec.key_link
                    curr_rec will not be linked in the bucket anymore.
                    """
                    # update in place the key link pointer of previous record
                    store.update_record_link_inplace(prev_record.store_position, next_record.key_link)
                    return True
                prev_record = next_record
                record = next_record
        return False

    def flush(self):
        self.db_mem.flush()


class Store:
    """
    Store manages persistence of records to disk with configurable flush behavior.
    
    Args:
        flush_interval: Number of writes before auto-flush. 
                       1 = flush every write (safest, default)
                       0 = manual flush only (fastest, use sync())
                       N>1 = flush every N writes with async background thread
    """

    def __init__(self, tablemeta, config, logger, caching_enabled=True, shared_cache=None, 
                 flush_interval=1):
        self.caching_enabled = caching_enabled
        self.batch_mode = False  # When True, defers flush() until end_batch()
        self.logger = logging.getLogger('store')
        self.tablemeta = tablemeta
        self.config = config
        self.flush_interval = flush_interval
        self.write_count = 0
        self._closed = False
        self.empty_block = '-1'.zfill(self.config.INDEX_BLOCK_LEN).encode()
        self.store = self.config.cog_store(
            tablemeta.namespace, tablemeta.name, tablemeta.db_instance_id)
        self.store_cache = Cache(self.store, shared_cache)
        temp = open(self.store, 'a')  # create if not exist
        temp.close()
        self.store_file = open(self.store, 'rb+')
        
        # Thread safety
        self._lock = threading.Lock()
        
        # Auto-enable async flush when interval > 1
        self._use_async = flush_interval > 1
        if self._use_async:
            self._flush_queue = queue.Queue()
            self._flush_thread = threading.Thread(target=self._flush_worker, daemon=True)
            self._flush_thread.start()
            self._shutdown = False
        
        logger.info(f"Store init: {self.store} (flush_interval={flush_interval})")

    def _flush_worker(self):
        """Background thread that processes flush requests."""
        while True:
            try:
                # Wait for flush signal (blocks until item available)
                item = self._flush_queue.get(timeout=1.0)
                if item == "SHUTDOWN":
                    # Drain remaining items and shutdown
                    while not self._flush_queue.empty():
                        try:
                            self._flush_queue.get_nowait()
                            self._flush_queue.task_done()
                        except queue.Empty:
                            break
                    self._flush_queue.task_done()
                    break
                # Perform actual flush (check if not closed)
                if not self._closed:
                    with self._lock:
                        if not self._closed:
                            self.store_file.flush()
                self._flush_queue.task_done()
            except queue.Empty:
                # Timeout - check if we should continue
                if getattr(self, '_shutdown', False):
                    break
                continue

    def _request_flush(self):
        """Request a flush - async if interval > 1, sync otherwise."""
        if self._closed:
            return
        if self._use_async:
            self._flush_queue.put("FLUSH")
        else:
            self.store_file.flush()

    def _handle_write_flush(self):
        """Increment write count and trigger flush if threshold reached."""
        if not self.batch_mode:
            self.write_count += 1
            if self.flush_interval > 0 and self.write_count >= self.flush_interval:
                self._request_flush()
                self.write_count = 0

    def sync(self):
        """
        Force flush all pending writes to disk.
        Blocks until flush is complete.
        """
        if self._closed:
            return
        with self._lock:
            if not self._closed:
                self.store_file.flush()
        if self._use_async:
            # Wait for async queue to drain
            self._flush_queue.join()

    def close(self):
        """Close the store, ensuring all data is flushed."""
        if self._closed:
            return
            
        # Mark as closed first
        self._closed = True
        
        if self._use_async:
            self._shutdown = True
            self._flush_queue.put("SHUTDOWN")
            self._flush_thread.join(timeout=5.0)
        
        with self._lock:
            try:
                self.store_file.flush()
                self.store_file.close()
            except ValueError:
                pass  # File already closed

    def begin_batch(self):
        """
        Enable batch mode - defers flush() until end_batch() is called.
        Use this when inserting many records for significantly better performance.
        """
        self.batch_mode = True

    def end_batch(self):
        """
        End batch mode and flush all pending writes to disk.
        """
        with self._lock:
            self.store_file.flush()
        self.batch_mode = False

    def save(self, record):
        """
        Store data with configurable flush behavior.
        """
        with self._lock:
            self.store_file.seek(0, 2)
            store_position = self.store_file.tell()
            record.set_store_position(store_position)
            marshalled_record = record.marshal()
            self.store_file.write(marshalled_record)
            
            if self.caching_enabled:
                self.store_cache.put(store_position, marshalled_record)
            
            # Handle flush based on interval
            self._handle_write_flush()
        
        return store_position

    def update_record_link_inplace(self, start_pos, int_value):
        """updates record link in store file in place"""
        if type(int_value) is not int:
            raise ValueError("store position must be int but provided : " + str(start_pos))

        byte_value = str(int_value).encode().rjust(Record.RECORD_LINK_LEN)
        self.logger.debug('update_record_link_inplace: ' + str(byte_value))
        
        with self._lock:
            self.store_file.seek(start_pos)
            self.store_file.write(byte_value)

            if self.caching_enabled:
                self.store_cache.partial_update_from_zero_index(start_pos, byte_value)
            
            self._handle_write_flush()

    # @profile
    def read(self, position):
        self.logger.debug("store read request at position: " + str(position))
        if self.caching_enabled:
            cached_record = self.store_cache.get(position)
            if cached_record is not None:
                return cached_record

        self.store_file.seek(position)
        record = self.__read_record()

        if self.caching_enabled:
            self.store_cache.put(position, record)

        return record

    def __read_record(self):
        """Read a single record using structured/length-aware parsing.

        Record layout (written by Record.marshal):
            [key_link: 16 bytes]
            [tombstone: 1 byte]
            [value_type: 1 byte]
            [value_len digits: variable ASCII]
            [UNIT_SEP: 1 byte]
            [serialized payload: exactly value_len bytes]
            [value_link digits: variable ASCII, only if value_type in ('l','u')]
            [RECORD_SEP: 1 byte]
        """
        # --- fixed-size header: key_link(16) + tombstone(1) + value_type(1) = 18 bytes
        header = self.__read_exactly(18)
        if header is None or len(header) < 18:
            return None

        value_type = chr(header[17])       # 's', 'l', or 'u'

        # --- variable-length value_len digits terminated by UNIT_SEP
        len_buf = b''
        while True:
            b = self.__read_exactly(1)
            if b is None:
                return None
            if b == UNIT_SEP:
                break
            len_buf += b
        value_len = int(len_buf.decode())

        # --- serialized payload: exactly value_len bytes
        payload = self.__read_exactly(value_len)
        if payload is None or len(payload) < value_len:
            return None

        # --- optional value_link + RECORD_SEP terminator
        tail = b''
        if value_type in ('l', 'u'):
            # value_link digits (ASCII) terminated by RECORD_SEP
            while True:
                b = self.__read_exactly(1)
                if b is None:
                    break
                tail += b
                if b == RECORD_SEP:
                    break
        else:
            # just the RECORD_SEP terminator
            sep = self.__read_exactly(1)
            if sep is not None:
                tail = sep

        record = header + len_buf + UNIT_SEP + payload + tail
        self.logger.debug("store __read_record: " + str(record))
        return record

    def __read_exactly(self, n):
        """Read exactly *n* bytes from the store file, or return None on EOF."""
        data = self.store_file.read(n)
        if len(data) == 0:
            return None
        while len(data) < n:
            chunk = self.store_file.read(n - len(data))
            if len(chunk) == 0:
                return data          # partial read (EOF mid-record)
            data += chunk
        return data


class Indexer:
    '''
    Manages indexes. Creates new index when an index is full.
    Searches all indexes for get requests.
    Provides same get/put/del method as single index but over multuple files.
    '''

    def __init__(self, tablemeta, config, logger):
        self.tablemeta = tablemeta
        self.config = config
        self.logger = logging.getLogger('indexer')
        self.index_list = []  # future range index.
        self.index_id = 0
        self.load_indexes()
        # if no index currenlty exist, create new live index.
        if len(self.index_list) == 0:
            self.index_list.append(Index(tablemeta, config, logger, self.index_id))
            self.live_index = self.index_list[self.index_id]

    def close(self):
        for idx in self.index_list:
            idx.close()

    def load_indexes(self):
        for f in os.listdir(self.config.cog_data_dir(self.tablemeta.namespace)):
            if self.config.INDEX in f:
                if self.tablemeta.name == self.config.get_table_name(f):
                    self.logger.info("loading index file: " + f)
                    id = self.config.index_id(f)
                    index = Index(self.tablemeta, self.config, self.logger, id)
                    self.index_list.append(index)
                    # make the latest index the live index.
                    if id >= self.index_id:
                        self.index_id = id
                        self.live_index = index

    def put(self, key, store_position, store):
        resp = self.live_index.put(key, store_position, store)
        self.logger.debug("Key: " + key + " indexed in: " + self.live_index.name)
        return resp

    # @profile
    def get(self, key, store):
        idx = self.index_list[0]  # only one index file.
        return idx.get(key, store)

    def get_head_only(self, key, store):
        """Get head record only, O(1) - doesn't traverse value chain."""
        idx = self.index_list[0]
        return idx.get_head_only(key, store)

    def scanner(self, store):
        for idx in self.index_list:
            self.logger.debug("SCAN: index: " + idx.name)
            for r in idx.scanner(store):
                yield r

    def delete(self, key, store):
        for idx in self.index_list:
            if idx.delete(key, store):
                return True
            else:
                return False

def cog_hash(string, index_capacity):
    return xxhash.xxh32(string, seed=2).intdigest() % index_capacity
