"""
SFTP Manager for TEE directory operations.

Provides high-level tools for managing TEE machine directories through SFTP.
This is an admin-only interface since only admins have SFTP access.
"""

import stat
from typing import List, Dict, Optional, Tuple, Any
from pathlib import Path
import logging

from .sftp import SFTPClient
from .exceptions import SFTPError

logger = logging.getLogger(__name__)


class FileInfo:
    """Information about a remote file."""
    
    def __init__(self, name: str, path: str, size: int, is_dir: bool, mtime: int):
        self.name = name
        self.path = path
        self.size = size
        self.is_dir = is_dir
        self.mtime = mtime
    
    @property
    def size_mb(self) -> float:
        """File size in MB."""
        return self.size / (1024 ** 2)
    
    @property
    def size_gb(self) -> float:
        """File size in GB."""
        return self.size / (1024 ** 3)
    
    def __repr__(self) -> str:
        type_str = "dir" if self.is_dir else "file"
        size_str = f"{self.size_mb:.2f} MB" if not self.is_dir else "-"
        return f"FileInfo({self.name}, {type_str}, {size_str})"


class SFTPManager:
    """
    High-level SFTP manager for TEE directory operations.
    
    Provides convenient methods for:
    - Listing files in standard TEE directories
    - Clearing directories
    - Managing models and datasets
    - Disk usage monitoring
    
    This is designed for admin operations only.
    
    Example:
        ```python
        from fhenomai import FHEnomClient
        
        client = FHEnomClient()
        with client.admin.sftp as sftp:
            # List uploaded models
            models = sftp.list_upload_directory()
            
            # Clear download directory
            sftp.clear_download_directory()
            
            # Get disk usage
            usage = sftp.get_directory_size("upload")
            print(f"Upload directory: {usage:.2f} GB")
        ```
    """
    
    # Standard TEE directory paths
    BASE_PATH = "/var/lib/fhenomai/FHEnomAI-server/admin"
    UPLOAD_DIR = "upload"
    DOWNLOAD_DIR = "download"
    
    def __init__(self, sftp_client: SFTPClient):
        """
        Initialize SFTP manager.
        
        Args:
            sftp_client: Connected SFTP client instance
        """
        self.sftp = sftp_client
        
        # Use absolute paths for the base locations
        self.upload_path = f"{self.BASE_PATH}/{self.UPLOAD_DIR}"
        self.download_path = f"{self.BASE_PATH}/{self.DOWNLOAD_DIR}"
    
    def _ensure_connected(self):
        """Ensure SFTP connection is active."""
        if not self.sftp._connected:
            self.sftp.connect()
    
    def list_files(
        self,
        remote_path: str,
        recursive: bool = False,
        include_hidden: bool = False
    ) -> List[FileInfo]:
        """
        List files in a remote directory.
        
        Args:
            remote_path: Remote directory path
            recursive: List files recursively
            include_hidden: Include hidden files (starting with .)
        
        Returns:
            List of FileInfo objects
        """
        self._ensure_connected()
        
        try:
            files = []
            items = self.sftp.sftp.listdir_attr(remote_path)
            
            for item in items:
                # Skip hidden files unless requested
                if not include_hidden and item.filename.startswith('.'):
                    continue
                
                item_path = f"{remote_path}/{item.filename}"
                is_dir = stat.S_ISDIR(item.st_mode)
                
                file_info = FileInfo(
                    name=item.filename,
                    path=item_path,
                    size=item.st_size,
                    is_dir=is_dir,
                    mtime=item.st_mtime
                )
                files.append(file_info)
                
                # Recurse into directories if requested
                if recursive and is_dir:
                    files.extend(self.list_files(item_path, recursive=True, include_hidden=include_hidden))
            
            return files
        
        except Exception as e:
            raise SFTPError(f"Failed to list files in {remote_path}: {e}")
    
    def list_upload_directory(self, recursive: bool = True) -> List[FileInfo]:
        """
        List files in the TEE upload directory.
        
        Args:
            recursive: List files recursively
        
        Returns:
            List of FileInfo objects
        """
        logger.info("Listing upload directory")
        return self.list_files(self.upload_path, recursive=recursive)
    
    def list_download_directory(self, recursive: bool = True) -> List[FileInfo]:
        """
        List files in the TEE download directory.
        
        Args:
            recursive: List files recursively
        
        Returns:
            List of FileInfo objects
        """
        logger.info("Listing download directory")
        return self.list_files(self.download_path, recursive=recursive)
    
    def list_models_in_upload(self) -> List[str]:
        """
        List model directories in upload folder.
        
        Returns:
            List of model directory names
        """
        files = self.list_upload_directory(recursive=False)
        models = [f.name for f in files if f.is_dir]
        logger.info(f"Found {len(models)} models in upload directory")
        return models
    
    def list_models_in_download(self) -> List[str]:
        """
        List encrypted model directories in download folder.
        
        Returns:
            List of encrypted model directory names
        """
        files = self.list_download_directory(recursive=False)
        models = [f.name for f in files if f.is_dir]
        logger.info(f"Found {len(models)} encrypted models in download directory")
        return models
    
    def clear_directory(
        self,
        remote_path: str,
        exclude_patterns: Optional[List[str]] = None,
        dry_run: bool = False
    ) -> Tuple[int, int]:
        """
        Clear all contents of a directory.
        
        Args:
            remote_path: Remote directory to clear
            exclude_patterns: List of filename patterns to exclude from deletion
            dry_run: If True, only count files without deleting
        
        Returns:
            Tuple of (files_deleted, dirs_deleted)
        """
        self._ensure_connected()
        
        if exclude_patterns is None:
            exclude_patterns = []
        
        files_deleted = 0
        dirs_deleted = 0
        
        def should_exclude(filename: str) -> bool:
            return any(pattern in filename for pattern in exclude_patterns)
        
        def clear_recursive(path: str):
            nonlocal files_deleted, dirs_deleted
            
            try:
                items = self.sftp.sftp.listdir_attr(path)
            except Exception as e:
                logger.error(f"Error listing {path}: {e}")
                return
                
            for item in items:
                item_path = f"{path}/{item.filename}"
                
                # Check exclusions
                if should_exclude(item.filename):
                    logger.debug(f"Skipping excluded: {item_path}")
                    continue
                
                try:
                    if stat.S_ISDIR(item.st_mode):
                        # Recursively clear directory
                        clear_recursive(item_path)
                        # Delete empty directory
                        if not dry_run:
                            try:
                                self.sftp.sftp.rmdir(item_path)
                                dirs_deleted += 1
                                logger.debug(f"Deleted directory: {item_path}")
                            except Exception as e:
                                logger.warning(f"Could not delete directory {item_path}: {e}")
                        else:
                            dirs_deleted += 1
                    else:
                        # Delete file
                        if not dry_run:
                            try:
                                self.sftp.sftp.remove(item_path)
                                files_deleted += 1
                                logger.debug(f"Deleted file: {item_path}")
                            except Exception as e:
                                logger.warning(f"Could not delete file {item_path}: {e}")
                        else:
                            files_deleted += 1
                except Exception as e:
                    logger.warning(f"Error processing {item_path}: {e}")
        
        clear_recursive(remote_path)
        
        action = "Would delete" if dry_run else "Deleted"
        logger.info(f"{action} {files_deleted} files and {dirs_deleted} directories from {remote_path}")
        
        return files_deleted, dirs_deleted
    
    def clear_upload_directory(self, exclude_patterns: Optional[List[str]] = None, dry_run: bool = False) -> Tuple[int, int]:
        """
        Clear the TEE upload directory.
        
        Args:
            exclude_patterns: List of filename patterns to exclude from deletion
            dry_run: If True, only count files without deleting
        
        Returns:
            Tuple of (files_deleted, dirs_deleted)
        """
        logger.info(f"Clearing upload directory (dry_run={dry_run})")
        return self.clear_directory(self.upload_path, exclude_patterns, dry_run)
    
    def clear_download_directory(self, exclude_patterns: Optional[List[str]] = None, dry_run: bool = False) -> Tuple[int, int]:
        """
        Clear the TEE download directory.
        
        Args:
            exclude_patterns: List of filename patterns to exclude from deletion
            dry_run: If True, only count files without deleting
        
        Returns:
            Tuple of (files_deleted, dirs_deleted)
        """
        logger.info(f"Clearing download directory (dry_run={dry_run})")
        return self.clear_directory(self.download_path, exclude_patterns, dry_run)
    
    def clear_all_directories(self, dry_run: bool = False) -> Dict[str, Tuple[int, int]]:
        """
        Clear both upload and download directories.
        
        Args:
            dry_run: If True, only count files without deleting
        
        Returns:
            Dictionary mapping directory names to (files_deleted, dirs_deleted) tuples
        """
        logger.info(f"Clearing all TEE directories (dry_run={dry_run})")
        
        results = {
            "upload": self.clear_upload_directory(dry_run=dry_run),
            "download": self.clear_download_directory(dry_run=dry_run)
        }
        
        total_files = sum(r[0] for r in results.values())
        total_dirs = sum(r[1] for r in results.values())
        action = "Would delete" if dry_run else "Deleted"
        logger.info(f"{action} {total_files} files and {total_dirs} directories total")
        
        return results
    
    def get_directory_size(self, remote_path: str) -> float:
        """
        Calculate total size of a directory in GB.
        
        Args:
            remote_path: Remote directory path
        
        Returns:
            Total size in GB
        """
        files = self.list_files(remote_path, recursive=True)
        total_bytes = sum(f.size for f in files if not f.is_dir)
        total_gb = total_bytes / (1024 ** 3)
        
        logger.info(f"Directory {remote_path}: {total_gb:.2f} GB")
        return total_gb
    
    def get_disk_usage(self) -> Dict[str, float]:
        """
        Get disk usage for all TEE directories.
        
        Returns:
            Dictionary mapping directory names to sizes in GB
        """
        logger.info("Calculating disk usage")
        
        usage = {
            "upload": self.get_directory_size(self.upload_path),
            "download": self.get_directory_size(self.download_path)
        }
        
        usage["total"] = sum(usage.values())
        
        return usage
    
    def delete_model_from_upload(self, model_name: str, dry_run: bool = False) -> bool:
        """
        Delete a specific model from the upload directory.
        
        Args:
            model_name: Name of the model directory to delete
            dry_run: If True, only check if model exists without deleting
        
        Returns:
            True if deleted (or would be deleted in dry_run), False if not found
        """
        model_path = f"{self.upload_path}/{model_name}"
        
        if not self.sftp.exists(model_path):
            logger.warning(f"Model not found in upload: {model_name}")
            return False
        
        if not dry_run:
            files, dirs = self.clear_directory(model_path, dry_run=False)
            self.sftp.delete(model_path, recursive=False)
            logger.info(f"Deleted model from upload: {model_name} ({files} files, {dirs} dirs)")
        else:
            logger.info(f"Would delete model from upload: {model_name}")
        
        return True
    
    def delete_model_from_download(self, model_id: str, dry_run: bool = False) -> bool:
        """
        Delete a specific encrypted model from the download directory.
        
        Args:
            model_id: ID/UUID of the encrypted model directory to delete
            dry_run: If True, only check if model exists without deleting
        
        Returns:
            True if deleted (or would be deleted in dry_run), False if not found
        """
        model_path = f"{self.download_path}/{model_id}"
        
        if not self.sftp.exists(model_path):
            logger.warning(f"Encrypted model not found in download: {model_id}")
            return False
        
        if not dry_run:
            files, dirs = self.clear_directory(model_path, dry_run=False)
            self.sftp.delete(model_path, recursive=False)
            logger.info(f"Deleted encrypted model from download: {model_id} ({files} files, {dirs} dirs)")
        else:
            logger.info(f"Would delete encrypted model from download: {model_id}")
        
        return True
    
    def get_directory_summary(self) -> Dict[str, Any]:
        """
        Get a comprehensive summary of TEE directories.
        
        Returns:
            Dictionary with directory information including:
            - File counts
            - Directory sizes
            - Model lists
        """
        logger.info("Generating directory summary")
        
        upload_files = self.list_upload_directory(recursive=False)
        download_files = self.list_download_directory(recursive=False)
        
        usage = self.get_disk_usage()
        
        summary = {
            "upload": {
                "path": self.upload_path,
                "size_gb": usage["upload"],
                "file_count": len([f for f in upload_files if not f.is_dir]),
                "dir_count": len([f for f in upload_files if f.is_dir]),
                "models": [f.name for f in upload_files if f.is_dir]
            },
            "download": {
                "path": self.download_path,
                "size_gb": usage["download"],
                "file_count": len([f for f in download_files if not f.is_dir]),
                "dir_count": len([f for f in download_files if f.is_dir]),
                "models": [f.name for f in download_files if f.is_dir]
            },
            "total_size_gb": usage["total"]
        }
        
        return summary
    
    def upload_directory(
        self,
        local_path: str,
        remote_path: str,
        show_progress: bool = True
    ):
        """
        Upload a local directory to remote path.
        
        Args:
            local_path: Local directory path
            remote_path: Remote directory path (auto-prefixed with upload/ if needed)
            show_progress: Show upload progress
        """
        # Normalize path: add upload/ prefix if not present
        if not remote_path.startswith('upload/') and not remote_path.startswith('download/') and not remote_path.startswith('/'):
            remote_path = f"upload/{remote_path}"
        
        self._ensure_connected()
        return self.sftp.upload_directory(
            local_path=local_path,
            remote_path=remote_path,
            show_progress=show_progress
        )
    
    def download_directory(
        self,
        remote_path: str,
        local_path: str,
        show_progress: bool = True
    ):
        """
        Download a remote directory to local path.
        
        Args:
            remote_path: Remote directory path (auto-prefixed with download/ if needed)
            local_path: Local directory path
            show_progress: Show download progress
        """
        # Normalize path: add download/ prefix if not present
        if not remote_path.startswith('upload/') and not remote_path.startswith('download/') and not remote_path.startswith('/'):
            remote_path = f"download/{remote_path}"
        
        self._ensure_connected()
        return self.sftp.download_directory(
            remote_path=remote_path,
            local_path=local_path,
            show_progress=show_progress
        )
    
    def list_directory(self, remote_path: str) -> List[str]:
        """
        List contents of a directory.
        
        Args:
            remote_path: Remote directory path
        
        Returns:
            List of filenames
        """
        self._ensure_connected()
        return self.sftp.list_directory(remote_path)
    
    def __enter__(self):
        """Context manager entry."""
        self._ensure_connected()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        # Don't disconnect - let the parent SFTP client handle that
        pass
