"""
SFTP client for FHEnom AI file operations.

Handles upload/download of models and datasets to/from TEE machine.
"""

import paramiko
import os
import stat
from pathlib import Path
from typing import Optional, List, Callable
import logging
from tqdm import tqdm

from .exceptions import SFTPError, AuthenticationError

logger = logging.getLogger(__name__)


class SFTPClient:
    """SFTP client for model/dataset transfer operations."""
    
    def __init__(
        self,
        host: str,
        username: str,
        port: int = 22,
        password: Optional[str] = None,
        key_path: Optional[str] = None,
        auto_connect: bool = False
    ):
        """
        Initialize SFTP client.
        
        Args:
            host: TEE machine hostname or IP
            username: SFTP username
            port: SSH port (default 22)
            password: Password (if not using key auth)
            key_path: Path to SSH private key
            auto_connect: Whether to connect immediately
        """
        self.host = host
        self.port = port
        self.username = username
        self.password = password
        self.key_path = key_path
        
        self.client: Optional[paramiko.SSHClient] = None
        self.sftp: Optional[paramiko.SFTPClient] = None
        self._connected = False
        
        logger.info(f"Initialized SFTP client for {username}@{host}:{port}")
        
        if auto_connect:
            self.connect()
    
    def connect(self):
        """Establish SFTP connection."""
        if self._connected:
            logger.debug("Already connected")
            return
        
        try:
            self.client = paramiko.SSHClient()
            self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            connect_kwargs = {
                "hostname": self.host,
                "port": self.port,
                "username": self.username
            }
            
            if self.key_path:
                key_path = os.path.expanduser(self.key_path)
                connect_kwargs["key_filename"] = key_path
                logger.info(f"Using SSH key: {key_path}")
            elif self.password:
                connect_kwargs["password"] = self.password
                logger.info("Using password authentication")
            else:
                raise AuthenticationError("Must provide either password or key_path")
            
            self.client.connect(**connect_kwargs)
            self.sftp = self.client.open_sftp()
            self._connected = True
            logger.info("SFTP connection established")
            
        except paramiko.AuthenticationException as e:
            raise AuthenticationError(f"Authentication failed: {e}")
        except Exception as e:
            raise SFTPError(f"Failed to connect: {e}")
    
    def disconnect(self):
        """Close SFTP connection."""
        if self.sftp:
            self.sftp.close()
            self.sftp = None
        if self.client:
            self.client.close()
            self.client = None
        self._connected = False
        logger.info("SFTP connection closed")
    
    def __enter__(self):
        """Context manager entry."""
        self.connect()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.disconnect()
    
    def _ensure_connected(self):
        """Ensure SFTP connection is active."""
        if not self._connected or not self.sftp:
            raise SFTPError("Not connected. Call connect() first or use context manager.")
    
    def list_directory(self, remote_path: str) -> List[str]:
        """
        List files in remote directory.
        
        Args:
            remote_path: Remote directory path
        
        Returns:
            List of filenames
        """
        self._ensure_connected()
        try:
            files = self.sftp.listdir(remote_path)
            logger.debug(f"Listed {len(files)} files in {remote_path}")
            return files
        except Exception as e:
            raise SFTPError(f"Failed to list directory {remote_path}: {e}")
    
    def exists(self, remote_path: str) -> bool:
        """
        Check if remote path exists.
        
        Args:
            remote_path: Remote path to check
        
        Returns:
            True if exists, False otherwise
        """
        self._ensure_connected()
        try:
            self.sftp.stat(remote_path)
            return True
        except IOError:
            return False
    
    def mkdir(self, remote_path: str, recursive: bool = True):
        """
        Create remote directory.
        
        Args:
            remote_path: Remote directory path
            recursive: Create parent directories if needed
        """
        self._ensure_connected()
        
        if self.exists(remote_path):
            logger.debug(f"Directory already exists: {remote_path}")
            return
        
        if recursive:
            parts = remote_path.split('/')
            current = ''
            for part in parts:
                if not part:
                    continue
                current = f"{current}/{part}" if current else part
                if not self.exists(current):
                    try:
                        self.sftp.mkdir(current)
                        logger.debug(f"Created directory: {current}")
                    except Exception as e:
                        logger.warning(f"Failed to create {current}: {e}")
        else:
            try:
                self.sftp.mkdir(remote_path)
                logger.debug(f"Created directory: {remote_path}")
            except Exception as e:
                raise SFTPError(f"Failed to create directory {remote_path}: {e}")
    
    def upload_file(
        self,
        local_file: str,
        remote_file: str,
        callback: Optional[Callable[[int, int], None]] = None
    ):
        """
        Upload a single file.
        
        Args:
            local_file: Local file path
            remote_file: Remote file path (auto-prefixed with upload/ if needed)
            callback: Optional callback(bytes_transferred, total_bytes)
        """
        self._ensure_connected()
        
        # Normalize path: add upload/ prefix if not present
        if not remote_file.startswith('upload/') and not remote_file.startswith('download/') and not remote_file.startswith('/'):
            remote_file = f"upload/{remote_file}"
        
        try:
            # Ensure remote directory exists
            remote_dir = str(Path(remote_file).parent)
            self.mkdir(remote_dir, recursive=True)
            
            # Upload file
            self.sftp.put(local_file, remote_file, callback=callback)
            logger.info(f"Uploaded: {local_file} -> {remote_file}")
            
        except Exception as e:
            raise SFTPError(f"Failed to upload {local_file}: {e}")
    
    def download_file(
        self,
        remote_file: str,
        local_file: str,
        callback: Optional[Callable[[int, int], None]] = None
    ):
        """
        Download a single file.
        
        Args:
            remote_file: Remote file path (auto-prefixed with download/ if needed)
            local_file: Local file path
            callback: Optional callback(bytes_transferred, total_bytes)
        """
        self._ensure_connected()
        
        # Normalize path: add download/ prefix if not present
        if not remote_file.startswith('upload/') and not remote_file.startswith('download/') and not remote_file.startswith('/'):
            remote_file = f"download/{remote_file}"
        
        try:
            # Ensure local directory exists
            local_dir = os.path.dirname(local_file)
            if local_dir:
                os.makedirs(local_dir, exist_ok=True)
            
            # Download file
            self.sftp.get(remote_file, local_file, callback=callback)
            logger.info(f"Downloaded: {remote_file} -> {local_file}")
            
        except Exception as e:
            raise SFTPError(f"Failed to download {remote_file}: {e}")
    
    def upload_directory(
        self,
        local_path: str,
        remote_path: str,
        exclude_patterns: Optional[List[str]] = None,
        show_progress: bool = True
    ) -> int:
        """
        Upload a directory recursively.
        
        Args:
            local_path: Local directory path
            remote_path: Remote directory path (auto-prefixed with upload/ if needed)
            exclude_patterns: List of patterns to exclude (e.g., ['.git', '__pycache__'])
            show_progress: Whether to show progress bar
        
        Returns:
            Number of files uploaded
        """
        self._ensure_connected()
        
        # 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}"
        
        if exclude_patterns is None:
            exclude_patterns = ['.git', '__pycache__', '.DS_Store', '*.pyc']
        
        local_path_obj = Path(local_path)
        if not local_path_obj.exists():
            raise SFTPError(f"Local path does not exist: {local_path}")
        
        # Create remote directory
        self.mkdir(remote_path, recursive=True)
        
        # Collect all files to upload
        files_to_upload = []
        for file_path in local_path_obj.rglob('*'):
            if file_path.is_file():
                # Check exclusions
                if any(pattern in str(file_path) for pattern in exclude_patterns):
                    continue
                files_to_upload.append(file_path)
        
        # Upload files with progress bar
        uploaded = 0
        iterator = tqdm(files_to_upload, desc="Uploading") if show_progress else files_to_upload
        
        for file_path in iterator:
            relative_path = file_path.relative_to(local_path)
            remote_file_path = f"{remote_path}/{relative_path}".replace('\\\\', '/')
            
            try:
                self.upload_file(str(file_path), remote_file_path)
                uploaded += 1
            except Exception as e:
                logger.error(f"Failed to upload {file_path}: {e}")
        
        logger.info(f"Uploaded {uploaded}/{len(files_to_upload)} files")
        return uploaded
    
    def download_directory(
        self,
        remote_path: str,
        local_path: str,
        show_progress: bool = True
    ) -> int:
        """
        Download a directory recursively.
        
        Args:
            remote_path: Remote directory path (auto-prefixed with download/ if needed)
            local_path: Local directory path
            show_progress: Whether to show progress bar
        
        Returns:
            Number of files downloaded
        """
        self._ensure_connected()
        
        # 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}"
        
        # Create local directory
        os.makedirs(local_path, exist_ok=True)
        
        # Recursively download files
        downloaded = 0
        
        def download_recursive(remote_dir, local_dir):
            nonlocal downloaded
            
            try:
                items = self.sftp.listdir_attr(remote_dir)
            except Exception as e:
                raise SFTPError(f"Failed to list {remote_dir}: {e}")
            
            for item in items:
                if item.filename.startswith('.'):
                    continue
                
                remote_item_path = f"{remote_dir}/{item.filename}"
                local_item_path = os.path.join(local_dir, item.filename)
                
                if stat.S_ISDIR(item.st_mode):
                    # Directory - recurse
                    os.makedirs(local_item_path, exist_ok=True)
                    download_recursive(remote_item_path, local_item_path)
                else:
                    # File - download
                    try:
                        self.download_file(remote_item_path, local_item_path)
                        downloaded += 1
                    except Exception as e:
                        logger.error(f"Failed to download {remote_item_path}: {e}")
        
        download_recursive(remote_path, local_path)
        logger.info(f"Downloaded {downloaded} files")
        return downloaded
    
    def delete(self, remote_path: str, recursive: bool = False):
        """
        Delete remote file or directory.
        
        Args:
            remote_path: Remote path to delete
            recursive: Delete recursively if directory
        """
        self._ensure_connected()
        
        try:
            # Check if it's a directory
            try:
                attrs = self.sftp.stat(remote_path)
                is_dir = stat.S_ISDIR(attrs.st_mode)
            except:
                is_dir = False
            
            if is_dir:
                if recursive:
                    # Delete contents first
                    items = self.sftp.listdir_attr(remote_path)
                    for item in items:
                        item_path = f"{remote_path}/{item.filename}"
                        if stat.S_ISDIR(item.st_mode):
                            self.delete(item_path, recursive=True)
                        else:
                            self.sftp.remove(item_path)
                    # Delete directory
                    self.sftp.rmdir(remote_path)
                else:
                    self.sftp.rmdir(remote_path)
            else:
                # Delete file
                self.sftp.remove(remote_path)
            
            logger.info(f"Deleted: {remote_path}")
            
        except Exception as e:
            raise SFTPError(f"Failed to delete {remote_path}: {e}")
    
    def get_file_size(self, remote_path: str) -> int:
        """
        Get size of remote file in bytes.
        
        Args:
            remote_path: Remote file path
        
        Returns:
            File size in bytes
        """
        self._ensure_connected()
        
        try:
            attrs = self.sftp.stat(remote_path)
            return attrs.st_size
        except Exception as e:
            raise SFTPError(f"Failed to get size of {remote_path}: {e}")
