import os
import sys
import logging
import pickle
import time
import tempfile
from multiprocessing import Array, Value, Process, Event 
from typing import Any, Dict, Callable
from functools import wraps
import traceback
import hashlib

# Global Dictionaries for internal state management
_EXPORTED_FUNCTIONS: Dict[str, Callable] = {}
_ACTIVE_EVENTS: Dict[str, Event] = {}   
_ACTIVE_FLAGS: Dict[str, Value] = {}    
# GLOBAL DICTIONARIES for path transfer objects (Array and Value)
_ACTIVE_PATH_ARRAYS: Dict[str, Array] = {}
_ACTIVE_PATH_LENGTHS: Dict[str, Value] = {}
# GLOBAL DICTIONARIES for Checksum
_ACTIVE_CHECKSUM_ARRAYS: Dict[str, Array] = {} 
CHECKSUM_BUFFER_SIZE = 64 # Space for a hexadecimal SHA-256 hash (64 chars)


# -----------------------------------------------------
# 1. LogicLock Logging System Configuration
# -----------------------------------------------------

logger = logging.getLogger("logiclock")
# The library only retrieves the logger. 
# Level is set to DEBUG to allow users to capture all messages.
logger.setLevel(logging.DEBUG) 

def set_log_level(level: str):
    """
    Allows the user to change the logging level (e.g., 'DEBUG', 'INFO').
    """
    level_map = {
        'DEBUG': logging.DEBUG,
        'INFO': logging.INFO,
        'WARNING': logging.WARNING,
        'ERROR': logging.ERROR,
        'CRITICAL': logging.CRITICAL
    }
    logger.setLevel(level_map.get(level.upper(), logging.INFO))
    logger.info(f"LogicLock log level adjusted to: {logger.getEffectiveLevel()}")

# -----------------------------------------------------
# 2. Path and Checksum Read/Write Functions
# -----------------------------------------------------

PATH_BUFFER_SIZE = 256

def _write_path(key: str, path_array: Array, path_length_value: Value, path: str):
    """Writes the temporary file path to the shared Array (Guardian)."""
    
    path_bytes = path.encode('utf-8')
    path_len = len(path_bytes)
    
    if path_len >= PATH_BUFFER_SIZE:
        raise ValueError(f"Path is too long: {path_len} bytes.")

    path_array.acquire()
    try:
        path_array.raw = path_bytes
    finally:
        path_array.release()
    
    path_length_value.value = path_len
    logger.debug(f"[{key}] Guardian: Path length written: {path_len}")


def _read_path(key: str) -> str:
    """Reads the path from the shared Array (Parent)."""
    
    path_array = _ACTIVE_PATH_ARRAYS.get(key)
    path_length_value = _ACTIVE_PATH_LENGTHS.get(key)

    if path_array is None or path_length_value is None:
        return "" # Return empty string if objects were already cleaned up

    path_len = path_length_value.value
    if path_len == 0:
        return "" # Path was never written by Guardian

    path_array.acquire()
    try:
        # Read only up to the path length written by the Guardian
        path = path_array.raw[:path_len].decode('utf-8') 
    finally:
        path_array.release()
        
    logger.debug(f"[{key}] Parent: Path read successfully.")
    return path


def _write_checksum(key: str, checksum_array: Array, data: bytes):
    """Calculates SHA-256 and writes the hex checksum to the Array (Guardian)."""
    
    checksum_hex = hashlib.sha256(data).hexdigest()
    checksum_bytes = checksum_hex.encode('utf-8')
    
    if len(checksum_bytes) != CHECKSUM_BUFFER_SIZE:
        raise ValueError("Checksum size is incorrect.")

    checksum_array.acquire()
    try:
        checksum_array.raw = checksum_bytes
    finally:
        checksum_array.release()
    
    logger.debug(f"[{key}] Guardian: Checksum written: {checksum_hex}")
    return checksum_hex

def _read_checksum(key: str) -> str:
    """Reads the checksum from the shared Array (Parent)."""
    checksum_array = _ACTIVE_CHECKSUM_ARRAYS.get(key)
    
    if checksum_array is None:
        raise KeyError(f"Shared checksum objects for '{key}' not found.")
        
    checksum_array.acquire()
    try:
        checksum_bytes = checksum_array.raw[:CHECKSUM_BUFFER_SIZE]
    finally:
        checksum_array.release()
        
    checksum = checksum_bytes.decode('utf-8').strip('\x00') # strip padding
    logger.debug(f"[{key}] Parent: Decoded Checksum: {checksum}")
    return checksum

# -----------------------------------------------------
# 3. Guardian and Access Logic
# -----------------------------------------------------

def _target_function(key: str, args: tuple, kwargs: dict):
    """
    The Target function executed in the isolated process (Guardian). 
    Ensures resource cleanup regardless of user logic outcome.
    """
    
    release_event = kwargs.pop('__release_event', None)
    shared_flag = kwargs.pop('__shared_flag', None) 
    path_array = kwargs.pop('__path_array', None)
    path_length_value = kwargs.pop('__path_length', None)
    checksum_array = kwargs.pop('__checksum_array', None) 

    try:
        user_func = _EXPORTED_FUNCTIONS[key]
    except KeyError:
        logger.error(f"[{key}] CRITICAL ERROR: Exported function not found in subprocess (KeyError).")
        return
        
    logger.info(f"[{key}] Guardian Process STARTED. PID: {os.getpid()}")
    
    temp_file_path = None
    data_payload = None
    
    try:
        # 1. Execute user function
        result = user_func(*args, **kwargs)
        
        # 🛑 2. CREATE and WRITE to TEMPORARY FILE
        fd, temp_file_path = tempfile.mkstemp(prefix='logiclock_')
        
        with os.fdopen(fd, 'wb') as tmp:
            data_payload = pickle.dumps(result)
            tmp.write(data_payload)
        
        logger.info(f"[{key}] Result written to Temp File: {temp_file_path}")

        # 🛑 3. WRITE PATH and CHECKSUM to SHARED ARRAY
        _write_path(key, path_array, path_length_value, temp_file_path)
        _write_checksum(key, checksum_array, data_payload)
        
        time.sleep(0.1) 
        
        # 🛑 RAISE THE FLAG 
        if shared_flag:
            shared_flag.value = True 
            
        logger.info(f"[{key}] Write complete. Guardian BLOCKED, awaiting Parent's signal.")
        if release_event:
            release_event.wait() 

    except Exception as e:
        logger.error(f"[{key}] Guardian Runtime Error: {type(e).__name__}: {e}")
        # Traceback is suppressed in production environment for stability but logged here.
        # traceback.print_exc(file=sys.stderr)
        
    finally:
        # GUARANTEED CLEANUP: If the file was created before the crash, clean it up.
        if temp_file_path and os.path.exists(temp_file_path):
            try:
                os.remove(temp_file_path)
                logger.debug(f"[{key}] Guardian: Cleaned up temp file: {temp_file_path}")
            except OSError as e:
                logger.warning(f"[{key}] Guardian cleanup failed: {e}")
        
        logger.info(f"[{key}] Guardian Process FINALIZED (Signal received).")


def isolate_export(key: str) -> Callable:
    """
    Decorator to isolate a function into a subprocess (Guardian). 
    Prepares all shared synchronization objects.
    """
    def decorator(func: Callable):
        
        if key in _EXPORTED_FUNCTIONS:
            logger.warning(f"Export key '{key}' already registered. Overwriting.")
        _EXPORTED_FUNCTIONS[key] = func
        
        @wraps(func)
        def wrapper(*args, **kwargs):
            
            logger.info(f"[{key}] Invoking decorator. Launching Guardian process...")
            
            # CREATE SYNCHRONIZATION OBJECTS
            release_event = Event() 
            shared_flag = Value('b', False) 
            path_array = Array('c', PATH_BUFFER_SIZE)
            path_length_value = Value('i', 0)
            checksum_array = Array('c', CHECKSUM_BUFFER_SIZE)

            # REGISTER ACTIVE OBJECTS
            _ACTIVE_EVENTS[key] = release_event
            _ACTIVE_FLAGS[key] = shared_flag
            _ACTIVE_PATH_ARRAYS[key] = path_array
            _ACTIVE_PATH_LENGTHS[key] = path_length_value
            _ACTIVE_CHECKSUM_ARRAYS[key] = checksum_array 
            
            # PASS ALL OBJECTS TO CHILD
            process_kwargs = {
                '__release_event': release_event, 
                '__shared_flag': shared_flag, 
                '__path_array': path_array, 
                '__path_length': path_length_value, 
                '__checksum_array': checksum_array,
                **kwargs
            }
            
            p = Process(target=_target_function, args=(key, args, process_kwargs))
            p.start()
            
            logger.info(f"[{key}] Guardian process launched successfully. Child PID: {p.pid}")
            return p 
        return wrapper
    return decorator


class LogicLockExports:
    """
    Access class that implements the synchronous lock, read, and integrity 
    validation (Checksum) for the isolated process result.
    """
    def __getattr__(self, key: str) -> Any:
        
        # 1. VALIDATION AND OBJECT RETRIEVAL (EXISTING)
        # Assuming objects like release_event, shared_flag, etc. are retrieved here...
        release_event = _ACTIVE_EVENTS.get(key)
        shared_flag = _ACTIVE_FLAGS.get(key)
        # (Other objects retrieved similarly)
        
        if not all([release_event, shared_flag]):
            raise KeyError(f"Export '{key}' is not active or was already retrieved.")

        # --- LOCAL STATE INITIALIZATION ---
        # Variables that must be defined before the main try block 
        # to ensure they are available in the 'finally' or post-loop checks.
        data = None         
        temp_file_path = None
        
        RETRY_DELAY = 0.1 
        MAX_WAIT_SECONDS = 10 
        start_time = time.time()
        
        # 🛑 CRITICAL RESILIENCE BLOCK
        try:
            # 1. SYNCHRONOUS LOCK
            logger.info(f"[{key}] Access requested. Reading data (Guardian is alive).")
            logger.info(f"[{key}] Activating synchronous lock. Waiting for Flag to be TRUE...")

            while not shared_flag.value:
                if (time.time() - start_time) > MAX_WAIT_SECONDS:
                    raise TimeoutError(f"Guardian failed to set flag within {MAX_WAIT_SECONDS} seconds.")
                time.sleep(RETRY_DELAY)

            # 2. READ AND VALIDATE INTEGRITY
            logger.info(f"[{key}] Flag is TRUE. Guardian finished writing. Attempting final read...")
            
            MAX_READ_ATTEMPTS = 3 
            
            # STEP 1: Read the checksum from shared ARRAY (Security First!)
            expected_checksum = _read_checksum(key)
            
            for attempt in range(MAX_READ_ATTEMPTS):
                try:
                    # STEP 2: Read file path
                    temp_file_path = _read_path(key)
                    if not temp_file_path:
                        logger.critical(f"[{key}] Path is empty. Guardian crashed before writing the file path.")
                        # Force a clean exception to trigger the 'finally' block
                        raise FileNotFoundError("Guardian failed to write temp file path.")
                    
                    # STEP 3: Read data from temporary file (AS BYTES)
                    with open(temp_file_path, 'rb') as tmp:
                        data_payload = tmp.read()
                        
                    # 🛑 STEP 4: VERIFY CHECKSUM (THE INTEGRITY CHECK)
                    read_file_checksum = hashlib.sha256(data_payload).hexdigest()
                    
                    if read_file_checksum != expected_checksum:
                        logger.critical(f"[{key}] INTEGRITY CHECK FAILED! Expected: {expected_checksum[:8]}... Found: {read_file_checksum[:8]}...")
                        raise RuntimeError("Data integrity check failed. Possible tampering on disk.")
                        
                    # STEP 5: DESERIALIZE DATA (Only after integrity is verified!)
                    data = pickle.loads(data_payload)

                    # SUCCESS: Break the loop and proceed to cleanup
                    break 

                except FileNotFoundError as e:
                    if attempt < MAX_READ_ATTEMPTS - 1:
                        logger.warning(f"[{key}] Read attempt {attempt+1} failed ({type(e).__name__}). Retrying in {RETRY_DELAY}s...")
                        time.sleep(RETRY_DELAY)
                    else:
                        raise e # Re-raise on last attempt
                        
                except Exception as e:
                    logger.error(f"[{key}] FATAL ERROR during read/validation: {e}")
                    raise

            # 3. POST-LOOP VERIFICATION AND RETURN
            # If the loop finished without success, data remains None.
            if data is None:
                raise KeyError(f"Export '{key}' not available. Unknown failure.")
                
            logger.info(f"[{key}] Read (File and Path) successful. Preparing for cleanup.")
            return data

        except Exception as e:
            # Re-raise any critical exception from locking or reading
            raise
            
        finally:
            # 🛑 GUARANTEED CLEANUP OF ALL RESOURCES
            
            # 1. Signal the Guardian to terminate
            if release_event:
                release_event.set()
                
            # 2. Clean up the Temporary File (Only if the path was read)
            if temp_file_path and os.path.exists(temp_file_path):
                try:
                    os.remove(temp_file_path)
                    logger.debug(f"[{key}] Parent: Cleaned up temp file: {temp_file_path}")
                except OSError as e:
                    logger.warning(f"[{key}] Parent cleanup failed: {e}")
                    
            # 3. Clean up the Global Synchronization Resources
            _ACTIVE_EVENTS.pop(key, None)
            _ACTIVE_FLAGS.pop(key, None) 
            _ACTIVE_PATH_ARRAYS.pop(key, None) 
            _ACTIVE_PATH_LENGTHS.pop(key, None) 
            _ACTIVE_CHECKSUM_ARRAYS.pop(key, None)
            
            logger.debug(f"[{key}] Global synchronization resources cleaned up for final exit.")
            
exports = LogicLockExports()