"""
ShellBags Analysis Module

Analyzes Windows ShellBags registry keys to extract information about
folder access history and window positions.
"""

import struct
from typing import Optional, List, Dict, Any
import winreg
from datetime import datetime

from diona.system.sensitive.core.info_store import InfoStore
from diona.system.sensitive.runcontext.models import ContextInfo
from .models import ShellBagsData, ShellBagEntry, RegistryAnalysisStatus


class ShellBagsAnalyzer:
    """Analyzer for Windows ShellBags registry data."""
    
    SHELL_BAGS_PATHS = [
        r"Software\\Microsoft\\Windows\\Shell\\Bags",
        r"Software\\Microsoft\\Windows\\Shell\\BagMRU", 
        r"Software\\Microsoft\\Windows\\Shell\\BagMRU_Backup"
    ]
    
    def __init__(self):
        self._info_store = InfoStore()
        self._context_info: Optional[ContextInfo] = None
        
    def _initialize_context(self) -> bool:
        """Initialize context information and check permissions."""
        try:
            self._info_store.initialize()
            self._context_info = self._info_store._context_info
            
            if not self._context_info:
                return False
                
            # Check if we have sufficient permissions for registry access
            return True
            
        except Exception as e:
            print(f"Error initializing context: {e}")
            return False
    
    def _parse_binary_data(self, data: bytes, data_type: str) -> Any:
        """Parse binary registry data based on data type."""
        try:
            if not isinstance(data, bytes):
                return data
                
            if data_type == "uint32":
                return struct.unpack('<I', data[:4])[0] if len(data) >= 4 else None
            elif data_type == "uint16":
                return struct.unpack('<H', data[:2])[0] if len(data) >= 2 else None
            elif data_type == "string":
                # Try to decode as UTF-16LE
                try:
                    return data.decode('utf-16le', errors='ignore').rstrip('\\0')
                except:
                    return data.decode('utf-8', errors='ignore').rstrip('\\0')
            else:
                return data
        except Exception as e:
            print(f"Error parsing binary data: {e}")
            return None
    
    def _parse_item_data(self, data: bytes) -> Dict[str, Any]:
        """Parse ShellBag item data (MRU entries)."""
        item_info = {}
        
        try:
            # Basic parsing - this is a simplified version
            # ShellBags binary data is complex and varies by Windows version
            if len(data) >= 4:
                item_info['size'] = len(data)
                item_info['hex_dump'] = data.hex()[:32] + "..." if len(data) > 16 else data.hex()
                
                # Try to extract potential path information
                # This is a heuristic approach
                for encoding in ['utf-16le', 'utf-8']:
                    try:
                        decoded = data.decode(encoding, errors='ignore')
                        # Look for common path patterns
                        if '\\\\' in decoded or ':' in decoded:
                            path_parts = [part for part in decoded.split('\\0') if part.strip()]
                            if path_parts:
                                item_info['potential_paths'] = path_parts
                                break
                    except:
                        continue
                        
        except Exception as e:
            print(f"Error parsing item data: {e}")
            
        return item_info
    
    def _parse_bag_key_recursive(self, key_handle, current_path: str = "") -> List[ShellBagEntry]:
        """Recursively parse ShellBags registry key structure."""
        entries: List[ShellBagEntry] = []
        
        try:
            # Read values from current key
            value_count = winreg.QueryInfoKey(key_handle)[1]
            
            slot_index = None
            view_mode = None
            item_data = {}
            
            for i in range(value_count):
                try:
                    value_name, value_data, value_type = winreg.EnumValue(key_handle, i)
                    
                    if value_name == "NodeSlot":
                        slot_index = self._parse_binary_data(value_data, "uint32")
                    elif value_name == "LogicalViewMode":
                        view_mode = self._parse_binary_data(value_data, "uint32")
                    elif value_name.startswith("Item_") and value_type == winreg.REG_BINARY:
                        item_info = self._parse_item_data(value_data)
                        item_data[value_name] = item_info
                        
                except Exception as e:
                    print(f"Error processing value {value_name}: {e}")
                    continue
            
            # Create entry for current node if we have meaningful data
            if slot_index is not None or view_mode is not None or item_data:
                entry = ShellBagEntry(
                    path=current_path if current_path else "/",
                    slot_index=slot_index,
                    view_mode=view_mode,
                    item_data=item_data if item_data else None
                )
                entries.append(entry)
            
            # Recursively process subkeys
            subkey_count = winreg.QueryInfoKey(key_handle)[0]
            
            for i in range(subkey_count):
                subkey_name = winreg.EnumKey(key_handle, i)
                
                # Only process numeric subkeys (typical ShellBags structure)
                if subkey_name.isdigit():
                    try:
                        subkey_path = f"{current_path}/{subkey_name}" if current_path else subkey_name
                        subkey_handle = winreg.OpenKey(key_handle, subkey_name)
                        
                        child_entries = self._parse_bag_key_recursive(subkey_handle, subkey_path)
                        entries.extend(child_entries)
                        
                        winreg.CloseKey(subkey_handle)
                        
                    except Exception as e:
                        print(f"Error processing subkey {subkey_name}: {e}")
                        continue
                        
        except Exception as e:
            print(f"Error parsing bag key at {current_path}: {e}")
            
        return entries
    
    def analyze_location(self, bag_path: str) -> List[ShellBagEntry]:
        """Analyze ShellBags data at a specific registry location."""
        entries: List[ShellBagEntry] = []
        
        try:
            root_key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, bag_path)
            entries = self._parse_bag_key_recursive(root_key, bag_path)
            winreg.CloseKey(root_key)
            
        except FileNotFoundError:
            print(f"ShellBags location not found: {bag_path}")
        except PermissionError:
            print(f"Permission denied accessing: {bag_path}")
        except Exception as e:
            print(f"Error analyzing {bag_path}: {e}")
            
        return entries
    
    def analyze(self) -> ShellBagsData:
        """
        Analyze ShellBags registry data from all known locations.
        
        Returns:
            ShellBagsData object containing analysis results
        """
        if not self._initialize_context():
            return ShellBagsData(
                status=RegistryAnalysisStatus.PERMISSION_DENIED,
                error_message="Failed to initialize context or insufficient permissions"
            )
        
        bag_locations: Dict[str, List[ShellBagEntry]] = {}
        
        try:
            # Analyze each ShellBags location
            for bag_path in self.SHELL_BAGS_PATHS:
                entries = self.analyze_location(bag_path)
                if entries:
                    bag_locations[bag_path] = entries
            
            if bag_locations:
                return ShellBagsData(
                    bag_locations=bag_locations,
                    status=RegistryAnalysisStatus.SUCCESS
                )
            else:
                return ShellBagsData(
                    status=RegistryAnalysisStatus.KEY_NOT_FOUND,
                    error_message="No ShellBags data found in registry"
                )
                
        except Exception as e:
            return ShellBagsData(
                status=RegistryAnalysisStatus.ERROR,
                error_message=f"Error analyzing ShellBags: {str(e)}"
            )