"""
Recent Documents Analysis Module

Analyzes Windows RecentDocs registry keys to extract information about
recently accessed documents and applications.
"""

import os
import struct
from typing import Optional, List, Dict
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 RecentDocsData, RecentDocsEntry, RegistryAnalysisStatus

# LNK file parsing utilities
import struct
from typing import BinaryIO, Dict, Any
from datetime import datetime, timedelta

class LnkFileParser:
    """Parser for Windows LNK files."""
    
    @staticmethod
    def parse_lnk_file(file_path: str) -> Dict[str, Any]:
        """Parse a Windows LNK file and extract metadata."""
        try:
            with open(file_path, 'rb') as f:
                return LnkFileParser._parse_lnk_stream(f)
        except Exception as e:
            return {"error": str(e)}
    
    @staticmethod
    def _parse_lnk_stream(f: BinaryIO) -> Dict[str, Any]:
        """Parse LNK file from binary stream."""
        result = {}
        
        # Read header (4+4+4 bytes: header size, CLSID, flags)
        header = f.read(20)
        if len(header) < 20 or header[:4] != b'L\x00\x00\x00':
            return {"error": "Invalid LNK file header"}
        
        # Parse flags
        flags = struct.unpack('<I', header[16:20])[0]
        
        # Parse file attributes (optional)
        file_attributes = struct.unpack('<I', f.read(4))[0] if flags & 0x20 else 0
        
        # Parse creation, access, and write times
        result["creation_time"] = LnkFileParser._parse_filetime(f)
        result["access_time"] = LnkFileParser._parse_filetime(f)
        result["write_time"] = LnkFileParser._parse_filetime(f)
        
        # Parse file size
        file_size = struct.unpack('<I', f.read(4))[0] if flags & 0x40 else 0
        result["file_size"] = file_size
        
        # Parse icon index
        icon_index = struct.unpack('<I', f.read(4))[0] if flags & 0x80 else 0
        result["icon_index"] = icon_index
        
        # Parse show command
        show_command = struct.unpack('<I', f.read(4))[0] if flags & 0x100 else 0
        result["show_command"] = show_command
        
        # Skip hotkey and reserved bytes
        f.read(8)
        
        # Parse target path
        if flags & 0x01:  # Has shell item id list
            id_list_size = struct.unpack('<H', f.read(2))[0]
            f.seek(id_list_size, 1)
        
        # Parse target path
        if flags & 0x02:  # Has target path
            result["target_path"] = LnkFileParser._parse_string(f)
        
        # Parse working directory
        if flags & 0x04:  # Has working directory
            result["working_directory"] = LnkFileParser._parse_string(f)
        
        # Parse command line arguments
        if flags & 0x08:  # Has command line arguments
            result["command_line"] = LnkFileParser._parse_string(f)
        
        # Parse icon location
        if flags & 0x10:  # Has icon location
            result["icon_location"] = LnkFileParser._parse_string(f)
        
        return result
    
    @staticmethod
    def _parse_filetime(f: BinaryIO) -> Optional[datetime]:
        """Parse Windows FILETIME structure."""
        try:
            low, high = struct.unpack('<II', f.read(8))
            if low == 0 and high == 0:
                return None
            
            # Convert FILETIME to Python datetime
            timestamp = (high << 32) + low
            if timestamp == 0:
                return None
            
            # FILETIME is 100-nanosecond intervals since January 1, 1601
            epoch_start = datetime(1601, 1, 1)
            microseconds = timestamp // 10
            return epoch_start + timedelta(microseconds=microseconds)
        except:
            return None
    
    @staticmethod
    def _parse_string(f: BinaryIO) -> str:
        """Parse null-terminated UTF-16LE string from LNK file."""
        try:
            chars = []
            while True:
                char_bytes = f.read(2)
                if len(char_bytes) < 2:
                    break
                char = struct.unpack('<H', char_bytes)[0]
                if char == 0:
                    break
                chars.append(char)
            
            return bytes(chars).decode('utf-16le', errors='ignore')
        except:
            return ""


class RecentDocsAnalyzer:
    """Analyzer for Windows RecentDocs registry data."""
    
    RECENT_DOCS_KEY_PATH = r"Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\RecentDocs"
    
    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
            # For now, we assume standard user permissions are sufficient for HKCU access
            return True
            
        except Exception as e:
            print(f"Error initializing context: {e}")
            return False
    
    def _decode_registry_string(self, data: bytes) -> str:
        """Decode UTF-16LE encoded registry string data with proper Windows registry format handling."""
        try:
            # Windows RecentDocs registry format:
            # The data contains the filename in UTF-16LE, but may be followed by additional metadata
            # We need to find the actual filename part
            
            # Method 1: Look for common patterns in RecentDocs data
            # RecentDocs entries often contain the actual filename followed by .lnk and other metadata
            
            # Try to find the actual filename by looking for common patterns
            patterns = [
                # Pattern 1: Look for file extensions and extract reasonable length
                self._extract_by_extension(data),
                # Pattern 2: Try UTF-16LE decoding with intelligent truncation
                self._decode_utf16le_intelligent(data),
                # Pattern 3: Extract printable ASCII characters
                self._extract_printable_ascii(data),
                # Pattern 4: Try UTF-8 decoding
                self._decode_utf8_safe(data)
            ]
            
            # Return the first non-empty result
            for result in patterns:
                if result and len(result.strip()) > 1:
                    return result.strip()
            
            return ""
            
        except Exception as e:
            return ""
    
    def _extract_by_extension(self, data: bytes) -> str:
        """Extract filename by looking for common file extensions."""
        common_extensions = [
            b'.txt', b'.doc', b'.docx', b'.pdf', b'.jpg', b'.jpeg', b'.png', 
            b'.zip', b'.rar', b'.7z', b'.exe', b'.msi', b'.py', b'.js',
            b'.html', b'.css', b'.xml', b'.json', b'.yml', b'.yaml', b'.ini',
            b'.config', b'.log', b'.bat', b'.cmd', b'.ps1', b'.sh', b'.md'
        ]
        
        for ext in common_extensions:
            if ext in data:
                ext_index = data.find(ext)
                if ext_index > 0:
                    # Extract a reasonable chunk around the extension
                    start_index = max(0, ext_index - 200)  # Look back 200 bytes
                    end_index = min(len(data), ext_index + len(ext) + 50)  # Look forward 50 bytes
                    
                    chunk = data[start_index:end_index]
                    
                    # Try UTF-16LE decoding
                    try:
                        decoded = chunk.decode('utf-16le', errors='ignore')
                        # Clean up: remove nulls and control characters, keep printable
                        cleaned = ''.join(c for c in decoded if 32 <= ord(c) <= 126 or c in '.-_()[]{}')
                        if len(cleaned) > 3:  # Reasonable filename length
                            return cleaned
                    except:
                        pass
        
        return ""
    
    def _decode_utf16le_intelligent(self, data: bytes) -> str:
        """Intelligent UTF-16LE decoding with pattern recognition."""
        try:
            # Try decoding the entire data
            decoded = data.decode('utf-16le', errors='ignore')
            
            # Look for patterns that indicate actual filenames
            # Filenames typically contain alphanumeric characters, dots, hyphens, underscores
            import re
            
            # Pattern for reasonable filenames
            filename_pattern = r'[\w\-\.\s\(\)\[\]{}]+\.[\w]{1,10}'
            matches = re.findall(filename_pattern, decoded)
            
            if matches:
                # Return the longest match (likely the actual filename)
                longest_match = max(matches, key=len)
                if len(longest_match) > 3:
                    return longest_match
            
            # If no pattern match, try to extract the most reasonable substring
            # Split by null characters and take the longest non-empty part
            parts = decoded.split('\\x00')
            if parts:
                longest_part = max(parts, key=len)
                # Clean up: keep only printable characters
                cleaned = ''.join(c for c in longest_part if 32 <= ord(c) <= 126 or c in '.-_()[]{}')
                if len(cleaned) > 3:
                    return cleaned
            
            return ""
            
        except:
            return ""
    
    def _extract_printable_ascii(self, data: bytes) -> str:
        """Extract printable ASCII characters from the data."""
        try:
            ascii_chars = []
            for byte in data:
                if 32 <= byte <= 126:  # Printable ASCII range
                    ascii_chars.append(chr(byte))
            
            if ascii_chars:
                result = ''.join(ascii_chars)
                # Look for filename-like patterns in ASCII
                if '.' in result and len(result) > 3:
                    return result
            
            return ""
        except:
            return ""
    
    def _decode_utf8_safe(self, data: bytes) -> str:
        """Safe UTF-8 decoding with cleanup."""
        try:
            decoded = data.decode('utf-8', errors='ignore')
            # Remove control characters, keep printable
            cleaned = ''.join(c for c in decoded if 32 <= ord(c) <= 126 or c in '.-_()[]{}')
            if len(cleaned) > 3:
                return cleaned
            return ""
        except:
            return ""
    
    def _parse_recent_docs_key(self, key_handle) -> Dict[str, List[RecentDocsEntry]]:
        """Parse RecentDocs registry key and extract file entries."""
        entries: Dict[str, List[RecentDocsEntry]] = {}
        
        try:
            # Enumerate all subkeys (file extension groups)
            subkey_count = winreg.QueryInfoKey(key_handle)[0]
            
            for i in range(subkey_count):
                subkey_name = winreg.EnumKey(key_handle, i)
                
                # Only process extension groups (starting with .)
                if subkey_name.startswith('.'):
                    try:
                        subkey_handle = winreg.OpenKey(key_handle, subkey_name)
                        extension_entries = self._parse_extension_group(subkey_handle, subkey_name)
                        entries[subkey_name] = extension_entries
                        winreg.CloseKey(subkey_handle)
                    except Exception as e:
                        print(f"Error processing extension group {subkey_name}: {e}")
                        continue
                        
        except Exception as e:
            print(f"Error parsing RecentDocs key: {e}")
            
        return entries
    
    def _extract_file_extension(self, file_name: str) -> str:
        """Extract file extension from filename."""
        try:
            # Extract the last extension from filename
            # For "gradle-8.14.3.lnk" -> ".lnk"
            # For "4.0.lnk" -> ".lnk"
            if '.' in file_name:
                # Split by dots and get the last part
                parts = file_name.split('.')
                if len(parts) > 1:
                    last_part = parts[-1].lower()
                    # Check if it's a valid file extension (not too long, alphanumeric)
                    if last_part and len(last_part) <= 10 and last_part.isalnum():
                        return '.' + last_part
            return ""
        except:
            return ""
    
    def _resolve_lnk_file_path(self, file_name: str) -> Optional[str]:
        """Resolve actual file path for .lnk shortcut files."""
        try:
            # Check if this is a .lnk file
            if not file_name.lower().endswith('.lnk'):
                return None
            
            # Try to find the .lnk file in common locations
            common_locations = [
                os.path.expanduser("~\\AppData\\Roaming\\Microsoft\\Windows\\Recent"),
                os.path.expanduser("~\\Recent"),
                os.path.expanduser("~\\AppData\\Roaming\\Microsoft\\Windows\\Recent Documents"),
            ]
            
            for location in common_locations:
                if os.path.exists(location):
                    lnk_path = os.path.join(location, file_name)
                    if os.path.exists(lnk_path):
                        # Parse the .lnk file to get target path
                        result = LnkFileParser.parse_lnk_file(lnk_path)
                        target_path = result.get('target_path')
                        if target_path and os.path.exists(target_path):
                            return target_path
                        
                        # If target path doesn't exist, try working directory
                        working_dir = result.get('working_directory')
                        if working_dir and os.path.exists(working_dir):
                            return working_dir
            
            return None
            
        except Exception as e:
            print(f"Error resolving .lnk file {file_name}: {e}")
            return None
    
    def _parse_extension_group(self, key_handle, registry_key_name: str) -> List[RecentDocsEntry]:
        """Parse individual extension group subkey."""
        entries: List[RecentDocsEntry] = []
        
        try:
            value_count = winreg.QueryInfoKey(key_handle)[1]
            
            for i in range(value_count):
                try:
                    value_name, value_data, value_type = winreg.EnumValue(key_handle, i)
                    
                    # Only process string values (REG_BINARY or REG_SZ)
                    if value_type in (winreg.REG_BINARY, winreg.REG_SZ):
                        if value_type == winreg.REG_BINARY:
                            file_name = self._decode_registry_string(value_data)
                        else:
                            file_name = str(value_data).rstrip('\\0')
                        
                        if file_name and file_name.strip():
                            # Extract actual file extension from filename, not from registry key name
                            actual_extension = self._extract_file_extension(file_name)
                            
                            # Resolve actual file path for .lnk files
                            actual_file_path = None
                            if actual_extension == '.lnk':
                                actual_file_path = self._resolve_lnk_file_path(file_name)
                            
                            entry = RecentDocsEntry(
                                file_name=file_name,
                                value_name=value_name,
                                file_extension=actual_extension,
                                registry_key_name=registry_key_name,  # Store the original registry key for reference
                                actual_file_path=actual_file_path  # Actual file path for shortcuts
                            )
                            entries.append(entry)
                            
                except Exception as e:
                    print(f"Error processing value {i} in extension group {registry_key_name}: {e}")
                    continue
                    
        except Exception as e:
            print(f"Error parsing extension group {registry_key_name}: {e}")
            
        return entries
    
    def _parse_application_key(self, key_handle) -> List[RecentDocsEntry]:
        """Parse Application subkey for recent application documents."""
        entries: List[RecentDocsEntry] = []
        
        try:
            value_count = winreg.QueryInfoKey(key_handle)[1]
            
            for i in range(value_count):
                try:
                    value_name, value_data, value_type = winreg.EnumValue(key_handle, i)
                    
                    if value_type in (winreg.REG_BINARY, winreg.REG_SZ):
                        if value_type == winreg.REG_BINARY:
                            app_name = self._decode_registry_string(value_data)
                        else:
                            app_name = str(value_data).rstrip('\\0')
                        
                        if app_name and app_name.strip():
                            entry = RecentDocsEntry(
                                file_name=app_name,
                                value_name=value_name,
                                file_extension="Application"
                            )
                            entries.append(entry)
                            
                except Exception as e:
                    print(f"Error processing application value {i}: {e}")
                    continue
                    
        except Exception as e:
            print(f"Error parsing application key: {e}")
            
        return entries
    
    def analyze(self) -> RecentDocsData:
        """
        Analyze RecentDocs registry data.
        
        Returns:
            RecentDocsData object containing analysis results
        """
        if not self._initialize_context():
            return RecentDocsData(
                status=RegistryAnalysisStatus.PERMISSION_DENIED,
                error_message="Failed to initialize context or insufficient permissions"
            )
        
        try:
            # Open the main RecentDocs key
            recent_docs_key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, self.RECENT_DOCS_KEY_PATH)
            
            # Parse extension groups
            extension_groups = self._parse_recent_docs_key(recent_docs_key)
            
            # Try to parse Application subkey
            applications: List[RecentDocsEntry] = []
            try:
                app_key = winreg.OpenKey(recent_docs_key, "Application")
                applications = self._parse_application_key(app_key)
                winreg.CloseKey(app_key)
            except FileNotFoundError:
                # Application key doesn't exist, which is normal
                pass
            except Exception as e:
                print(f"Error processing Application key: {e}")
            
            winreg.CloseKey(recent_docs_key)
            
            return RecentDocsData(
                extension_groups=extension_groups,
                applications=applications,
                status=RegistryAnalysisStatus.SUCCESS
            )
            
        except FileNotFoundError:
            return RecentDocsData(
                status=RegistryAnalysisStatus.KEY_NOT_FOUND,
                error_message="RecentDocs registry key not found"
            )
            
        except PermissionError:
            return RecentDocsData(
                status=RegistryAnalysisStatus.PERMISSION_DENIED,
                error_message="Permission denied accessing registry"
            )
            
        except Exception as e:
            return RecentDocsData(
                status=RegistryAnalysisStatus.ERROR,
                error_message=f"Error analyzing RecentDocs: {str(e)}"
            )