import yaml
import os
import itertools
import datetime
import decimal
import importlib
import types
import array
import inspect
import io
import msgpack
import zstandard as zstd
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
from .input import secure_input


class Formatstore():
    to_yaml = 1
    to_binary = 2
    to_dict = 3


class Store:
    def __init__(self, filename=None, password=None, protected=False, format=Formatstore.to_binary, compressed=False, cached=False):
        self.filename = filename
        self.id_counter = itertools.count(1)  # Générateur d'identifiants uniques
        self.format = format

        self.enable_compression = compressed
        self.enable_cache = cached

        self._cache = {}
        self._cache_full = False  # True si le cache contient l'arborescence complète

        if protected and not password:
            password = secure_input("Password store: ")
        self.password = password.encode() if password else None

        self._serialize_handlers = {
            str: self._serialize_str,
            bool: self._serialize_bool,
            int: self._serialize_int,
            float: self._serialize_float,
            bytes: self._serialize_bytes,
            bytearray: self._serialize_bytearray,
            tuple: self._serialize_tuple,
            list: self._serialize_list,
            dict: self._serialize_dict,
            array.array: self._serialize_array,
            datetime.datetime: self._serialize_datetime,
            decimal.Decimal: self._serialize_decimal,
            io.BytesIO: self._serialize_BytesIO,
            io.StringIO: self._serialize_StringIO,
            io.TextIOWrapper: self._serialize_TextIOWrapper,
            io.BufferedReader: self._serialize_BufferedReader,
            types.BuiltinFunctionType: self._serialize_builtin,
            types.BuiltinMethodType: self._serialize_builtin,
            types.FunctionType: self._serialize_functiontype,
            complex: self._serialize_complex,
            set: self._serialize_frozenset,
            frozenset: self._serialize_frozenset,
        }

        self._deserialize_handlers = {
            'ref': self._deserialize_ref,
            'BytesIO': self._deserialize_BytesIO,
            'StringIO': self._deserialize_StringIO,
            'TextIOWrapper': self._deserialize_TextIOWrapper,
            'BufferedReader': self._deserialize_BufferedReader,
            'object': self._deserialize_object,
            'datetime': self._deserialize_datetime,
            'decimal': self._deserialize_decimal,
            'complex': self._deserialize_complex,
            'bytes': self._deserialize_bytes,
            'bytearray': self._deserialize_bytearray,
            'array': self._deserialize_array,
            'list': self._deserialize_list,
            'dict': self._deserialize_dict,
            'tuple': self._deserialize_tuple,
            'function': self._deserialize_function,
            'builtin': self._deserialize_builtin,
            'str': self._deserialize_str,
            'int': self._deserialize_int,
            'float': self._deserialize_float,
            'bool': self._deserialize_bool,
            'set': self._deserialize_frozenset,
            'frozenset': self._deserialize_frozenset,
        }

        # Si le fichier n'existe pas, créez-le avec une liste vide
        if self.filename:
            if os.path.dirname(self.filename) and not os.path.exists(os.path.dirname(self.filename)):
                os.makedirs(os.path.dirname(self.filename), exist_ok=True)

            #Impossible d'écrire directement un dictionnaire dans un fichier
            if self.format is Formatstore.to_dict:
                raise ValueError(f"Format {self.format} unsupported with output file")

            if not os.path.exists(filename):
                self._write_file({})
        else:
            self._data = {}
            self._write_file({})


    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        pass

    def get_content(self):
        return self._data

    def encrypt(self, data):
        salt = os.urandom(16)  # Salt pour le KDF

        kdf = Scrypt(
            salt=salt,
            length=32,
            n=2**14,
            r=8,
            p=1,
            backend=default_backend()
        )
        key = kdf.derive(self.password)
        iv = os.urandom(16)
        cipher = Cipher(algorithms.AES(key), modes.CFB(iv), backend=default_backend())
        encryptor = cipher.encryptor()
        encrypted_data = encryptor.update(data) + encryptor.finalize()
        return salt + iv + encrypted_data

    def decrypt(self, data):
        salt = data[:16]
        iv = data[16:32]
        encrypted_data = data[32:]
        kdf = Scrypt(
            salt=salt,
            length=32,
            n=2**14,
            r=8,
            p=1,
            backend=default_backend()
        )
        key = kdf.derive(self.password)
        cipher = Cipher(algorithms.AES(key), modes.CFB(iv), backend=default_backend())
        decryptor = cipher.decryptor()
        return decryptor.update(encrypted_data) + decryptor.finalize()

    def push(self, data_name, data, overwrite=True):
        existing_data = self._read_file()  # Données SÉRIALISÉES

        keys = data_name.split('.')
        d = existing_data

        # Naviguer dans le dictionnaire en créant les sous-dicts si besoin
        for key in keys[:-1]:
            if key not in d or not isinstance(d[key], dict):
                d[key] = {}
            d = d[key]

        last_key = keys[-1]

        if last_key not in d or overwrite:
            if isinstance(data, dict):
                if last_key not in d or not isinstance(d[last_key], dict):
                    d[last_key] = {}
                for k, v in data.items():
                    d[last_key][k] = v  # ← Fusion directe, comme avant
            else:
                d[last_key] = self._serialize(data)
        else:
            raise NameError(f"data_name '{data_name}' already exists")

        self._write_file(existing_data)

        # Ajouter juste le cache
        if self.enable_cache:
            # Désérialiser juste ce qu'on a modifié pour le cache
            deserialized = self._deserialize(d[last_key])
            self._store_in_cache(data_name, deserialized)


    def pull(self, data_name=None):
        if self.enable_cache and data_name is not None:
            cached_value = self._get_from_cache(data_name)
            if cached_value is not None:
                return cached_value

        existing_data = self._read_file()
        
        if data_name is None:
            result = self._deserialize(existing_data)
            if self.enable_cache:
                self._cache = result  # Remplacer complètement le cache
                self._cache_full = True

        else:
            keys = data_name.split('.')
            d = existing_data
            for key in keys:
                if not isinstance(d, dict) or key not in d:
                    return None
                d = d[key]
            result = self._deserialize(d)
            if self.enable_cache:
                self._store_in_cache(data_name, result)

        return result


    def erase(self, data_name):
        # Si on a le cache complet, l'utiliser directement
        if self.enable_cache and self._cache_full:
            # Si cache full, il faut resérialiser pour modifier
            existing_data = self._serialize(self._cache)
        else:
            existing_data = self._read_file()
        
        keys = data_name.split('.')
        d = existing_data

        # Naviguer dans le dict pour atteindre la clé à supprimer
        for key in keys[:-1]:
            if key not in d or not isinstance(d[key], dict):
                return False  # chemin inexistant
            d = d[key]

        last_key = keys[-1]
        if last_key in d:
            del d[last_key]
            self._write_file(existing_data)

            # Supprimer du cache aussi
            if self.enable_cache:
                self._remove_from_cache(data_name)

            return True
        return False


    def list(self):
        # Si enable_cache est activé et le cache est complet, l'utiliser
        if self.enable_cache and self._cache_full:
            existing_data = self._cache
        else:
            existing_data = self._read_file()

        def collect_keys(d, prefix=''):
            keys = []
            for k, v in d.items():
                # Ignorer les métadonnées internes
                if isinstance(k, str) and k.startswith('_'):
                    continue

                path = f"{prefix}.{k}" if prefix else k
                if isinstance(v, dict):
                    keys.extend(collect_keys(v, path))
                else:
                    keys.append(path)
            return keys

        deserialized = self._deserialize(existing_data)
        return collect_keys(deserialized)


    def _read_file(self):
        # Lire les données existantes depuis un fichier ou self._data
        if self.filename:
            with open(self.filename, 'rb') as f:
                file_content = f.read()
        else:
            file_content = self._data

        #Déchiffrer le fichier si besoin
        if self.password:
            file_content = self.decrypt(file_content)

        if self.enable_compression and isinstance(file_content, bytes):
            if file_content.endswith(b'\x01'):
                dctx = zstd.ZstdDecompressor()
                file_content = dctx.decompress(file_content[:-1])
            elif file_content.endswith(b'\x00'):
                file_content = file_content[:-1]

        try:
            if self.format is Formatstore.to_yaml:
                existing_data = yaml.load(file_content, Loader=yaml.SafeLoader)
            elif self.format is Formatstore.to_binary:
                existing_data = msgpack.unpackb(file_content, raw=False, strict_map_key=False)
            elif self.format is Formatstore.to_dict:
                existing_data = self._data
        except Exception as e:
            raise ValueError("Data unpacking failed. Possibly wrong password or corrupted data.") from e

        return existing_data


    def _write_file(self, data):
        # Écrire les données mises à jour
        if self.format is Formatstore.to_yaml:
            data_pack = yaml.dump(data)
            data_pack = data_pack.encode()
        elif self.format is Formatstore.to_binary:
            data_pack = msgpack.packb(data)
        else:
            data_pack = data

        if self.enable_compression and isinstance(data_pack, bytes):
            if len(data_pack) > 1024:  # Compresse seulement si > 1KB
                cctx = zstd.ZstdCompressor(level=22)  # 🔥 Compression MAX
                data_pack = cctx.compress(data_pack) + b'\x01'  # Flag Zstandard
            else:
                data_pack = data_pack + b'\x00'

        if self.password:
            data_pack = self.encrypt(data_pack)

        if self.format is Formatstore.to_yaml or self.format is Formatstore.to_binary:
            if self.filename:
                with open(self.filename, 'wb') as f:
                    f.write(data_pack)
            else:
                self._data = data_pack
        else:
            self._data = data_pack


    ######################
    ##     Serialize    ##
    ######################

    def _serialize_str(self, data, unique_id, obj_cache):
        return {'_type': 'str', '_id': unique_id, 'value': data}

    def _serialize_bool(self, data, unique_id, obj_cache):
        return {'_type': 'bool', '_id': unique_id, 'value': data}

    def _serialize_int(self, data, unique_id, obj_cache):
        return {'_type': 'int', '_id': unique_id, 'value': data}

    def _serialize_float(self, data, unique_id, obj_cache):
        return {'_type': 'float', '_id': unique_id, 'value': data}

    def _serialize_bytes(self, data, unique_id, obj_cache):
        return {'_type': 'bytes', '_id': unique_id, 'data': data.hex()}

    def _serialize_bytearray(self, data, unique_id, obj_cache):
        return {'_type': 'bytearray', '_id': unique_id, 'data': data.hex()}

    def _serialize_tuple(self, data, unique_id, obj_cache):
        return {'_type': 'tuple', '_id': unique_id, 'data': [self._serialize(item, obj_cache) for item in data]}

    def _serialize_list(self, data, unique_id, obj_cache):
        return {'_type': 'list', '_id': unique_id, 'data': [self._serialize(item, obj_cache) for item in data]}

    def _serialize_dict(self, data, unique_id, obj_cache):
        return {'_type': 'dict', '_id': unique_id, 'data': {key: self._serialize(value, obj_cache) for key, value in data.items()}}

    def _serialize_array(self, data, unique_id, obj_cache):
        return {'_type': 'array', '_id': unique_id, 'typecode': data.typecode, 'data': data.tolist()}

    def _serialize_datetime(self, data, unique_id, obj_cache):
        return {'_type': 'datetime', '_id': unique_id, 'data': data.timestamp()}

    def _serialize_decimal(self, data, unique_id, obj_cache):
        return {'_type': 'decimal', '_id': unique_id, 'value': str(data)}

    def _serialize_BytesIO(self, data, unique_id, obj_cache):
        return {
            '_type': 'BytesIO',
            '_id': unique_id,
            'data': data.getvalue().hex(),
            'position': data.tell(),
            'mode': 'rb'
        }


    def _serialize_StringIO(self, data, unique_id, obj_cache):
        return {
            '_type': 'StringIO',
            '_id': unique_id,
            'data': data.getvalue(),
            'position': data.tell(),
            'mode': 'r'
        }

    def _serialize_TextIOWrapper(self, data, unique_id, obj_cache):
        original_position = data.tell()
        data.seek(0)
        content = data.read()
        data.seek(original_position)
        return {
            '_type': 'TextIOWrapper',
            '_id': unique_id,
            'data': content,
            'position': original_position,
            'encoding': data.encoding,
            'newline': data.newlines
        }


    def _serialize_BufferedReader(self, data, unique_id, obj_cache):
        original_position = data.tell()
        data.seek(0)
        content = data.read()
        data.seek(original_position)
        return {
            '_type': 'BufferedReader',
            '_id': unique_id,
            'data': content.hex(),
            'position': original_position,
            'mode': 'rb'
        }

    def _serialize_builtin(self, data, unique_id, obj_cache):
        return {'_type': 'builtin', '_id': unique_id, 'name': data.__name__}


    def _serialize_functiontype(self, data, unique_id, obj_cache):
        return {
            '_type': 'function',
            '_id': unique_id,
            'name': data.__name__,
            'code': inspect.getsource(data),
            'globals': {k: self._serialize(v, obj_cache) for k, v in data.__globals__.items() if k in data.__code__.co_names},
            'defaults': self._serialize(data.__defaults__, obj_cache),
            'closure': self._serialize([c.cell_contents for c in (data.__closure__ or [])], obj_cache)
        }

    def _serialize_complex(self, data, unique_id, obj_cache):
        return {'_type': 'complex', '_id': unique_id, 'real': data.real, 'imag': data.imag}

    def _serialize_frozenset(self, data, unique_id, obj_cache):
        return {'_type': type(data).__name__, '_id': unique_id, 'data': [self._serialize(item, obj_cache) for item in data]}


    def _serialize(self, data, obj_cache=None):
        if obj_cache is None:
            obj_cache = {}
        
        # Gérer les références circulaires
        obj_id = id(data)
        if obj_id in obj_cache:
            return {'_type': 'ref', '_id': obj_cache[obj_id]}

        # Attribuer un identifiant unique pour chaque nouvel objet
        unique_id = next(self.id_counter)
        obj_cache[obj_id] = unique_id

        if data is None:
            return {'_type': 'none', '_id': unique_id}

        handler = self._serialize_handlers.get(type(data))
        if handler:
            return handler(data, unique_id, obj_cache)
        
        if hasattr(data, '__dict__') and not isinstance(data, types.FunctionType):
            serialized_data = {
                '_type': 'object',
                '_class': data.__class__.__name__,
                '_module': data.__class__.__module__,
                '_id': unique_id,
                'data': {key: self._serialize(value, obj_cache) for key, value in data.__dict__.items()},
                'methods': {name: method.__code__.co_code.hex() for name, method in data.__class__.__dict__.items() if callable(method)},
            }
            return serialized_data
        else:
            raise TypeError(f"Unsupported data type: {type(data)}")

    ######################
    ##    Deserialize   ##
    ######################
    def _deserialize_ref(self, data, obj_cache):
        if data['_id'] in obj_cache:
            return obj_cache[data['_id']]
        else:
            raise KeyError(f"Reference with id {data['_id']} not found")
    
    def _deserialize_BytesIO(self, data, obj_cache):
        buf = io.BytesIO(bytes.fromhex(data['data']))
        buf.seek(data['position'])
        obj_cache[data['_id']] = buf
        return buf

    def _deserialize_StringIO(self, data, obj_cache):
        buf = io.StringIO(data['data'])
        buf.seek(data['position'])
        obj_cache[data['_id']] = buf
        return buf
    
    def _deserialize_TextIOWrapper(self, data, obj_cache):
        buf = io.StringIO(data['data'])
        buf.seek(data['position'])
        wrapped = io.TextIOWrapper(buf, encoding=data['encoding'], newline=data['newline'])
        obj_cache[data['_id']] = wrapped
        return wrapped

    def _deserialize_BufferedReader(self, data, obj_cache):
        buf = io.BytesIO(bytes.fromhex(data['data']))
        buf.seek(data['position'])
        wrapped = io.BufferedReader(buf)
        obj_cache[data['_id']] = wrapped
        return wrapped

    def _deserialize_object(self, data, obj_cache):
        class_name = data['_class']
        module_name = data['_module']
        module = importlib.import_module(module_name)
        cls = getattr(module, class_name)
        obj = cls.__new__(cls)
        obj.__dict__.update({key: self._deserialize(value, obj_cache) for key, value in data['data'].items()})
        obj_cache[data['_id']] = obj
        return obj
    
    def _deserialize_datetime(self, data, obj_cache):
        return datetime.datetime.fromtimestamp(data['data'])
    
    def _deserialize_decimal(self, data, obj_cache):
        return decimal.Decimal(data['value'])
    
    def _deserialize_complex(self, data, obj_cache):
        return complex(data['real'], data['imag'])
    
    def _deserialize_bytes(self, data, obj_cache):
        return bytes.fromhex(data['data'])
    
    def _deserialize_bytearray(self, data, obj_cache):
        return bytearray.fromhex(data['data'])
    
    def _deserialize_array(self, data, obj_cache):
        return array.array(data['typecode'], data['data'])
    
    def _deserialize_frozenset(self, data, obj_cache):
        collection = {self._deserialize(item, obj_cache) for item in data['data']} if data['_type'] == 'set' else frozenset(self._deserialize(item, obj_cache) for item in data['data'])
        obj_cache[data['_id']] = collection
        return collection
    
    def _deserialize_list(self, data, obj_cache):
        collection = [self._deserialize(item, obj_cache) for item in data['data']]
        obj_cache[data['_id']] = collection
        return collection

    def _deserialize_dict(self, data, obj_cache):
        collection = {key: self._deserialize(value, obj_cache) for key, value in data['data'].items()}
        obj_cache[data['_id']] = collection
        return collection

    def _deserialize_tuple(self, data, obj_cache):
        collection = tuple(self._deserialize(item, obj_cache) for item in data['data'])
        obj_cache[data['_id']] = collection
        return collection

    def _deserialize_function(self, data, obj_cache):
        globals_dict = {}
        for name, value in data['globals'].items():
            try:
                #globals_dict[name] = eval(value)  # ❌ SÉCURITÉ : Exécute du code !
                globals_dict[name] = self._deserialize(value, obj_cache)  # Safe
            except:
                globals_dict[name] = value
        
        exec(data['code'], globals_dict)
        func_name = data['name']
        func = globals_dict[func_name]
        obj_cache[data['_id']] = func
        return func
    
    def _deserialize_builtin(self, data, obj_cache):
        collection = getattr(__builtins__, data['name'])
        obj_cache[data['_id']] = collection
        return collection
    
    def _deserialize_str(self, data, obj_cache):
        obj_cache[data['_id']] = data['value']
        return data['value']
    
    def _deserialize_int(self, data, obj_cache):
        obj_cache[data['_id']] = int(data['value'])
        return int(data['value'])
    
    def _deserialize_float(self, data, obj_cache):
        obj_cache[data['_id']] = float(data['value'])
        return float(data['value'])
    
    def _deserialize_bool(self, data, obj_cache):
        obj_cache[data['_id']] = bool(data['value'])
        return bool(data['value'])

    def _deserialize(self, data, obj_cache=None):
        if obj_cache is None:
            obj_cache = {}

        # Gérer les références circulaires
        if isinstance(data, dict) and '_type' in data:
            if data['_type'] == 'none':
                return None

            handler = self._deserialize_handlers.get(data['_type'])
            if handler:
                return handler(data, obj_cache)

        elif isinstance(data, list):
            return [self._deserialize(item, obj_cache) for item in data]
        elif isinstance(data, dict):
            return {key: self._deserialize(value, obj_cache) for key, value in data.items()}
        else:
            return data

    ######################
    ##       Cache      ##
    ######################
    def _store_in_cache(self, data_name, data):
        """Stocke une donnée DÉSÉRIALISÉE dans le cache hiérarchique"""
        keys = data_name.split('.')
        d = self._cache
        
        # Créer la structure hiérarchique
        for key in keys[:-1]:
            if key not in d:
                d[key] = {}
            elif not isinstance(d[key], dict):
                d[key] = {}
            d = d[key]
        
        # Stocker la donnée DÉSÉRIALISÉE au dernier niveau
        last_key = keys[-1]
        d[last_key] = data


    def _get_from_cache(self, data_name):
        """Récupère une donnée du cache hiérarchique (déjà désérialisée)"""
        if data_name is None:
            if self._cache_full:
                return self._cache
            return None
        
        keys = data_name.split('.')
        d = self._cache
        
        # Naviguer dans le cache hiérarchique
        for key in keys:
            if isinstance(d, dict) and key in d:
                d = d[key]
            else:
                return None
        
        # Retourner directement (déjà désérialisée)
        return d


    def _remove_from_cache(self, data_name):
        """Supprime une donnée du cache hiérarchique"""
        keys = data_name.split('.')
        d = self._cache
        
        # Naviguer jusqu'au parent de la clé à supprimer
        for key in keys[:-1]:
            if key not in d:
                return
            d = d[key]
        
        # Supprimer la clé
        last_key = keys[-1]
        if last_key in d:
            del d[last_key]