"""
Progressive search implementation with persistent state management.
"""

import json
import os
import uuid
from datetime import datetime
from pathlib import Path
from typing import List, Optional, Dict, Any

from diona.project.storage import ProjectStorage

from .models import SearchConfig, SearchResult, SearchState, SearchStatus, DirectoryState
from .file_rules import FileSelectionRules


class ProgressiveSearch:
    """Progressive search engine with persistent state management."""
    
    def __init__(self):
        self.storage = ProjectStorage()
        self.search_module = "windows"
        self.search_dir = "search"
        
        # Initialize storage directories
        self._ensure_storage_directories()
    
    def _ensure_storage_directories(self):
        """Ensure required storage directories exist."""
        # Create search directory in ~/.diona/project/windows/search/
        search_path = self.storage.get_path(
            self.search_module, 
            self.search_dir, 
            "project"
        )
        os.makedirs(search_path, exist_ok=True)
    
    def create_search_session(self, config: SearchConfig, task_name: str = "sensitive_search") -> str:
        """Create a new progressive search session."""
        session_id = str(uuid.uuid4())
        
        # Create initial search state
        search_state = SearchState(
            session_id=session_id,
            task_name=task_name,
            config=config,
            status=SearchStatus.PENDING
        )
        
        # Initialize directory discovery
        self._discover_directories(search_state)
        
        # Save initial state
        self._save_search_state(session_id, search_state)
        
        return session_id
    
    def _discover_directories(self, search_state: SearchState):
        """Discover all directories to be scanned."""
        target_dir = search_state.config.target_directory
        max_depth = search_state.config.max_depth
        
        # Start with the target directory
        directories_to_process = [(target_dir, 0)]
        discovered_directories = []
        
        while directories_to_process:
            current_dir, current_depth = directories_to_process.pop(0)
            
            # Skip if depth exceeds limit
            if current_depth > max_depth:
                continue
            
            discovered_directories.append(current_dir)
            
            try:
                # Discover subdirectories
                for entry in os.scandir(current_dir):
                    if entry.is_dir():
                        # Check if we should include hidden directories
                        if entry.name.startswith('.') and not search_state.config.include_hidden_files:
                            continue
                        directories_to_process.append((entry.path, current_depth + 1))
            except (OSError, PermissionError):
                # Skip directories with access issues
                continue
        
        # Initialize directory states - 确保目标目录优先处理
        # 按深度排序，浅层目录优先处理
        discovered_directories.sort(key=lambda d: self._calculate_depth(d, target_dir))
        
        for directory in discovered_directories:
            depth = self._calculate_depth(directory, target_dir)
            search_state.scanned_directories[directory] = DirectoryState(
                directory_path=directory,
                depth=depth,
                status="pending"
            )
            search_state.pending_directories.append(directory)
        
        search_state.total_directories = len(discovered_directories)
    
    def load_search_session(self, session_id: str) -> Optional[SearchState]:
        """Load an existing search session."""
        return self._load_search_state(session_id)
    
    def execute_search_session(self, session_id: str) -> SearchState:
        """Execute a search session with file limit constraints."""
        search_state = self._load_search_state(session_id)
        
        if not search_state:
            raise ValueError(f"Search session {session_id} not found")
        
        if search_state.status == SearchStatus.COMPLETED:
            return search_state
        
        # Update status
        search_state.status = SearchStatus.IN_PROGRESS
        if not search_state.start_time:
            search_state.start_time = datetime.now()
        
        # Initialize file rules
        file_rules = FileSelectionRules(search_state.config)
        
        # Calculate remaining file limit for this session
        files_remaining = search_state.config.max_files_per_session
        
        if files_remaining <= 0:
            search_state.status = SearchStatus.PAUSED
            self._save_search_state(session_id, search_state)
            return search_state
        
        # Process directories
        files_scanned_this_session = 0
        
        while search_state.pending_directories and files_scanned_this_session < files_remaining:
            current_dir = search_state.pending_directories.pop(0)
            search_state.current_directory = current_dir
            
            # Update directory state
            if current_dir in search_state.scanned_directories:
                search_state.scanned_directories[current_dir].status = "scanning"
            
            try:
                # Scan directory
                files_scanned, results_found = self._scan_directory(
                    current_dir, search_state, file_rules, files_remaining - files_scanned_this_session
                )
                
                files_scanned_this_session += files_scanned
                
                # Update directory state
                if current_dir in search_state.scanned_directories:
                    search_state.scanned_directories[current_dir].status = "completed"
                    search_state.scanned_directories[current_dir].files_scanned = files_scanned
                    search_state.scanned_directories[current_dir].last_scan_time = datetime.now()
                
                # Update progress
                search_state.update_progress(
                    files_scanned=files_scanned, 
                    files_found=results_found, 
                    directories_scanned=1
                )
                
                # Move to completed directories
                search_state.completed_directories.append(current_dir)
                
            except (OSError, PermissionError) as e:
                # Mark directory as completed with error
                if current_dir in search_state.scanned_directories:
                    search_state.scanned_directories[current_dir].status = "error"
                search_state.completed_directories.append(current_dir)
                search_state.update_progress(directories_scanned=1)
                continue
        
        # Check if search is complete
        if search_state.is_complete():
            search_state.status = SearchStatus.COMPLETED
            search_state.end_time = datetime.now()
        else:
            search_state.status = SearchStatus.PAUSED
        
        # Save updated state
        self._save_search_state(session_id, search_state)
        
        return search_state
    
    def _scan_directory(self, directory: str, search_state: SearchState, 
                       file_rules: FileSelectionRules, max_files: int) -> tuple[int, int]:
        """Scan a single directory for sensitive files."""
        files_scanned = 0
        results_found = 0
        
        try:
            # Get all files in directory
            entries = list(os.scandir(directory))
            
            # Count total files in directory
            total_files = sum(1 for entry in entries if entry.is_file())
            if directory in search_state.scanned_directories:
                search_state.scanned_directories[directory].total_files = total_files
            
            # Process files
            for entry in entries:
                if files_scanned >= max_files:
                    break
                    
                if entry.is_file():
                    # Check if we've already scanned this file
                    if entry.path in search_state.scanned_files:
                        continue
                    
                    # Check if file should be scanned (including hidden file check)
                    if not file_rules.should_scan_file(entry.path):
                        continue
                    
                    # Evaluate file
                    result = file_rules.evaluate_file(entry.path)
                    
                    if result:
                        search_state.add_result(result)
                        results_found += 1
                    
                    # Mark file as scanned
                    search_state.add_scanned_file(entry.path)
                    files_scanned += 1
        
        except (OSError, PermissionError):
            # Skip directories with access issues
            pass
        
        return files_scanned, results_found
    
    def _calculate_depth(self, path: str, base_path: str) -> int:
        """Calculate directory depth relative to base path."""
        try:
            relative_path = os.path.relpath(path, base_path)
            if relative_path == '.':
                return 0
            depth = len(relative_path.split(os.sep))
            return depth
        except ValueError:
            # Paths are on different drives (Windows)
            return 999
    
    def get_search_results(self, session_id: str) -> List[SearchResult]:
        """Get search results for a session."""
        search_state = self._load_search_state(session_id)
        
        if not search_state:
            return []
        
        return search_state.search_results
    
    def get_search_progress(self, session_id: str) -> Dict[str, Any]:
        """Get detailed search progress information."""
        search_state = self._load_search_state(session_id)
        
        if not search_state:
            return {}
        
        return {
            "session_id": session_id,
            "status": search_state.status.value,
            "progress_percentage": search_state.get_progress_percentage(),
            "total_directories": search_state.total_directories,
            "directories_scanned": search_state.total_directories_scanned,
            "total_files_scanned": search_state.total_files_scanned,
            "files_found": search_state.total_files_found,
            "pending_directories": len(search_state.pending_directories),
            "completed_directories": len(search_state.completed_directories)
        }
    
    def save_search_results(self, session_id: str, results_file: str = "results.json"):
        """Save search results to a file."""
        search_state = self._load_search_state(session_id)
        
        if not search_state:
            raise ValueError(f"Search session {session_id} not found")
        
        # Convert results to serializable format
        results_data = []
        for result in search_state.search_results:
            results_data.append({
                "file_path": result.file_path,
                "file_name": result.file_name,
                "file_size": result.file_size,
                "file_extension": result.file_extension,
                "last_modified": result.last_modified.isoformat(),
                "file_type": result.file_type,
                "matched_keywords": result.matched_keywords,
                "risk_score": result.risk_score,
                "confidence": result.confidence,
                "scan_timestamp": result.scan_timestamp.isoformat()
            })
        
        # Save results
        results_path = self.storage.get_path(
            self.search_module, 
            f"{self.search_dir}/{results_file}", 
            "project"
        )
        
        with open(results_path, 'w', encoding='utf-8') as f:
            json.dump(results_data, f, indent=2, ensure_ascii=False)
    
    def _save_search_state(self, session_id: str, search_state: SearchState):
        """Save search state to persistent storage."""
        # Convert to serializable format
        state_data = {
            "session_id": search_state.session_id,
            "task_name": search_state.task_name,
            "status": search_state.status.value,
            "start_time": search_state.start_time.isoformat() if search_state.start_time else None,
            "end_time": search_state.end_time.isoformat() if search_state.end_time else None,
            "total_files_scanned": search_state.total_files_scanned,
            "total_files_found": search_state.total_files_found,
            "total_directories": search_state.total_directories,
            "total_directories_scanned": search_state.total_directories_scanned,
            "current_directory": search_state.current_directory,
            "scanned_files": search_state.scanned_files,
            "pending_directories": search_state.pending_directories,
            "completed_directories": search_state.completed_directories,
            "search_results": [],
            "scanned_directories": {},
            "config": {
                "target_directory": search_state.config.target_directory,
                "max_files_per_session": search_state.config.max_files_per_session,
                "max_depth": search_state.config.max_depth,
                "document_extensions": search_state.config.document_extensions,
                "key_extensions": search_state.config.key_extensions,
                "config_extensions": search_state.config.config_extensions,
                "archive_extensions": search_state.config.archive_extensions,
                "wallet_extensions": search_state.config.wallet_extensions,
                "sensitive_keywords": search_state.config.sensitive_keywords,
                "min_file_size": search_state.config.min_file_size,
                "max_file_size": search_state.config.max_file_size,
                "include_hidden_files": search_state.config.include_hidden_files,
                "follow_symlinks": search_state.config.follow_symlinks
            },
            "created_at": search_state.created_at.isoformat(),
            "updated_at": search_state.updated_at.isoformat()
        }
        
        # Convert search results to serializable format
        for result in search_state.search_results:
            state_data["search_results"].append({
                "file_path": result.file_path,
                "file_name": result.file_name,
                "file_size": result.file_size,
                "file_extension": result.file_extension,
                "last_modified": result.last_modified.isoformat(),
                "file_type": result.file_type,
                "matched_keywords": result.matched_keywords,
                "risk_score": result.risk_score,
                "confidence": result.confidence,
                "scan_timestamp": result.scan_timestamp.isoformat()
            })
        
        # Convert directory states
        for dir_path, dir_state in search_state.scanned_directories.items():
            state_data["scanned_directories"][dir_path] = {
                "depth": dir_state.depth,
                "status": dir_state.status,
                "files_scanned": dir_state.files_scanned,
                "total_files": dir_state.total_files,
                "last_scan_time": dir_state.last_scan_time.isoformat() if dir_state.last_scan_time else None
            }
        
        # Save state file
        state_path = self.storage.get_path(
            self.search_module, 
            f"{self.search_dir}/{session_id}.json", 
            "project"
        )
        
        with open(state_path, 'w', encoding='utf-8') as f:
            json.dump(state_data, f, indent=2, ensure_ascii=False)
    
    def _load_search_state(self, session_id: str) -> Optional[SearchState]:
        """Load search state from persistent storage."""
        state_path = self.storage.get_path(
            self.search_module, 
            f"{self.search_dir}/{session_id}.json", 
            "project"
        )
        
        if not os.path.exists(state_path):
            return None
        
        try:
            with open(state_path, 'r', encoding='utf-8') as f:
                state_data = json.load(f)
            
            # Reconstruct SearchConfig
            config_data = state_data["config"]
            config = SearchConfig(
                target_directory=config_data["target_directory"],
                max_files_per_session=config_data["max_files_per_session"],
                max_depth=config_data["max_depth"],
                document_extensions=config_data["document_extensions"],
                key_extensions=config_data["key_extensions"],
                config_extensions=config_data["config_extensions"],
                archive_extensions=config_data["archive_extensions"],
                wallet_extensions=config_data["wallet_extensions"],
                sensitive_keywords=config_data["sensitive_keywords"],
                min_file_size=config_data["min_file_size"],
                max_file_size=config_data["max_file_size"],
                include_hidden_files=config_data["include_hidden_files"],
                follow_symlinks=config_data["follow_symlinks"]
            )
            
            # Reconstruct SearchState
            search_state = SearchState(
                session_id=state_data["session_id"],
                task_name=state_data["task_name"],
                status=SearchStatus(state_data["status"]),
                config=config,
                total_files_scanned=state_data["total_files_scanned"],
                total_files_found=state_data["total_files_found"],
                total_directories=state_data["total_directories"],
                total_directories_scanned=state_data["total_directories_scanned"],
                current_directory=state_data["current_directory"],
                scanned_files=state_data["scanned_files"],
                pending_directories=state_data["pending_directories"],
                completed_directories=state_data["completed_directories"],
                created_at=datetime.fromisoformat(state_data["created_at"]),
                updated_at=datetime.fromisoformat(state_data["updated_at"])
            )
            
            # Reconstruct search results
            from diona.system.sensitive.search.models import SearchResult
            for result_data in state_data.get("search_results", []):
                result = SearchResult(
                    file_path=result_data["file_path"],
                    file_name=result_data["file_name"],
                    file_size=result_data["file_size"],
                    file_extension=result_data["file_extension"],
                    last_modified=datetime.fromisoformat(result_data["last_modified"]),
                    file_type=result_data["file_type"],
                    matched_keywords=result_data["matched_keywords"],
                    risk_score=result_data["risk_score"],
                    confidence=result_data["confidence"],
                    scan_timestamp=datetime.fromisoformat(result_data["scan_timestamp"])
                )
                search_state.search_results.append(result)
            
            # Reconstruct directory states
            for dir_path, dir_data in state_data.get("scanned_directories", {}).items():
                dir_state = DirectoryState(
                    directory_path=dir_path,
                    depth=dir_data["depth"],
                    status=dir_data["status"],
                    files_scanned=dir_data["files_scanned"],
                    total_files=dir_data["total_files"]
                )
                if dir_data["last_scan_time"]:
                    dir_state.last_scan_time = datetime.fromisoformat(dir_data["last_scan_time"])
                search_state.scanned_directories[dir_path] = dir_state
            
            # Handle optional datetime fields
            if state_data["start_time"]:
                search_state.start_time = datetime.fromisoformat(state_data["start_time"])
            if state_data["end_time"]:
                search_state.end_time = datetime.fromisoformat(state_data["end_time"])
            
            return search_state
            
        except (json.JSONDecodeError, KeyError, ValueError) as e:
            print(f"Error loading search state: {e}")
            return None
    
    def list_search_sessions(self) -> List[Dict[str, Any]]:
        """List all available search sessions."""
        search_path = self.storage.get_path(
            self.search_module, 
            self.search_dir, 
            "project"
        )
        
        sessions = []
        
        try:
            if os.path.exists(search_path):
                for file_name in os.listdir(search_path):
                    # Only process files that look like session state files (UUID format)
                    if (file_name.endswith('.json') and 
                        len(file_name) == 41 and  # UUID length + .json
                        file_name != 'results.json' and 
                        not file_name.startswith('test_')):
                        
                        session_id = file_name[:-5]  # Remove .json extension
                        
                        # Load basic session info without full state
                        state_path = os.path.join(search_path, file_name)
                        try:
                            with open(state_path, 'r', encoding='utf-8') as f:
                                state_data = json.load(f)
                            
                            # Only include if it's a dictionary (session state)
                            if isinstance(state_data, dict):
                                sessions.append({
                                    "session_id": session_id,
                                    "task_name": state_data.get("task_name", "unknown"),
                                    "status": state_data.get("status", "unknown"),
                                    "total_files_scanned": state_data.get("total_files_scanned", 0),
                                    "total_files_found": state_data.get("total_files_found", 0),
                                    "total_directories": state_data.get("total_directories", 0),
                                    "directories_scanned": state_data.get("total_directories_scanned", 0),
                                    "progress_percentage": round((state_data.get("total_directories_scanned", 0) / max(state_data.get("total_directories", 1), 1) * 100), 1),
                                    "created_at": state_data.get("created_at", "unknown"),
                                    "updated_at": state_data.get("updated_at", "unknown")
                                })
                            
                        except (json.JSONDecodeError, KeyError, OSError):
                            continue
        except OSError:
            # Directory doesn't exist or can't be accessed
            pass
        
        return sessions