"""
Main FHEnom AI client - unified interface combining Admin API and SFTP.

Provides high-level convenience methods for common workflows.
"""

import logging
from typing import Optional, Dict, Any, List, Callable
from pathlib import Path

from .admin_api import AdminAPI
from .sftp import SFTPClient
from .sftp_manager import SFTPManager
from .config import FHEnomConfig
from .exceptions import FHEnomError

logger = logging.getLogger(__name__)


class FHEnomClient:
    """
    Unified FHEnom AI client.
    
    Combines Admin API and SFTP client for seamless workflow management.
    
    Example:
        ```python
        from fhenomai import FHEnomClient, FHEnomConfig
        
        # Initialize from config
        config = FHEnomConfig.from_file()
        client = FHEnomClient(config)
        
        # Or use context manager
        with FHEnomClient(config) as client:
            # Upload and encrypt model
            model_id = client.encrypt_and_upload_model(
                local_path="./my-model",
                model_name="my-model-encrypted"
            )
            
            # Start serving
            client.start_serving(model_id)
        ```
    """
    
    def __init__(
        self,
        config: Optional[FHEnomConfig] = None,
        admin_api: Optional[AdminAPI] = None,
        sftp_client: Optional[SFTPClient] = None,
        auto_connect_sftp: bool = False
    ):
        """
        Initialize FHEnom client.
        
        Args:
            config: Configuration object (loads from file if not provided)
            admin_api: Pre-configured AdminAPI instance (creates from config if not provided)
            sftp_client: Pre-configured SFTPClient instance (creates from config if not provided)
            auto_connect_sftp: Whether to connect SFTP immediately
        """
        # Load config if not provided
        if config is None:
            try:
                config = FHEnomConfig.from_file()
            except FileNotFoundError:
                config = FHEnomConfig.from_env()
        
        self.config = config
        
        # Initialize SFTP client first (needed by AdminAPI)
        if sftp_client:
            self.sftp = sftp_client
        else:
            self.sftp = SFTPClient(
                host=config.sftp_host,
                username=config.sftp_username,
                port=config.sftp_port,
                password=config.sftp_password,
                key_path=config.sftp_key_path,
                auto_connect=auto_connect_sftp
            )
        
        # Create SFTP manager for admin operations
        self._sftp_manager = SFTPManager(self.sftp)
        
        # Initialize Admin API with SFTP manager
        if admin_api:
            self.admin = admin_api
            self.admin._sftp_manager = self._sftp_manager
        else:
            self.admin = AdminAPI(
                base_url=config.admin_url,
                timeout=config.timeout,
                max_retries=config.max_retries,
                verify_ssl=config.verify_ssl,
                sftp_manager=self._sftp_manager,
                auth_token=config.auth_token
            )
        
        logger.info("Initialized FHEnomClient with SFTP manager")
    
    def __enter__(self):
        """Context manager entry - connects SFTP."""
        self.sftp.connect()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - disconnects SFTP."""
        self.sftp.disconnect()
    
    def get_sftp_manager(self) -> SFTPManager:
        """
        Get the SFTP manager instance.
        
        Returns:
            SFTPManager instance
        """
        return self._sftp_manager
    
    # ========================
    # High-Level Workflows
    # ========================
    
    def encrypt_and_upload_model(
        self,
        local_path: str,
        model_name: str,
        output_path: Optional[str] = None,
        upload_output: bool = True,
        wait: bool = True,
        show_progress: bool = True,
        **encrypt_kwargs
    ) -> str:
        """
        Complete workflow: upload model, encrypt, optionally download encrypted version.
        
        Args:
            local_path: Local model directory
            model_name: Name for the model on TEE
            output_path: Optional local path to download encrypted model
            upload_output: Whether to keep encrypted model on TEE
            wait: Whether to wait for encryption to complete
            show_progress: Show progress bars
            **encrypt_kwargs: Additional arguments for encrypt_model
        
        Returns:
            Model ID of encrypted model
        
        Example:
            ```python
            model_id = client.encrypt_and_upload_model(
                local_path="./llama-3.2-3b",
                model_name="llama-3.2-3b-encrypted",
                block_size=1024,
                add_residual_term=True
            )
            ```
        """
        logger.info(f"Starting encrypt_and_upload workflow for {model_name}")
        
        # 1. Upload model to TEE
        remote_models_dir = self.config.sftp_models_dir
        remote_model_path = f"{remote_models_dir}/{model_name}"
        
        logger.info(f"Uploading model to {remote_model_path}")
        with self.sftp:
            self.sftp.upload_directory(
                local_path=local_path,
                remote_path=remote_model_path,
                show_progress=show_progress
            )
        
        # 2. Encrypt model
        logger.info(f"Encrypting model {model_name}")
        
        encrypted_name = f"{model_name}_encrypted"
        job_id = self.admin.encrypt_model(
            model_id=model_name,
            output_model_id=encrypted_name,
            **encrypt_kwargs
        )
        
        logger.info(f"Encryption job started: {job_id}")
        
        # 3. Wait for encryption
        if wait:
            self.admin.wait_for_job(job_id, show_progress=show_progress)
            logger.info("Encryption completed")
        
        # 4. Download encrypted model if requested
        if output_path and wait:
            remote_encrypted_path = f"{remote_models_dir}/{encrypted_name}"
            logger.info(f"Downloading encrypted model to {output_path}")
            
            with self.sftp:
                self.sftp.download_directory(
                    remote_path=remote_encrypted_path,
                    local_path=output_path,
                    show_progress=show_progress
                )
        
        return encrypted_name
    
    def encrypt_and_upload_dataset(
        self,
        local_file: str,
        dataset_name: str,
        text_fields: List[str],
        output_path: Optional[str] = None,
        wait: bool = True,
        show_progress: bool = True,
        **encrypt_kwargs
    ) -> str:
        """
        Complete workflow: upload dataset, encrypt, optionally download encrypted version.
        
        Args:
            local_file: Local dataset file (Parquet)
            dataset_name: Name for dataset on TEE
            text_fields: List of column names to encrypt
            output_path: Optional local path to download encrypted dataset
            wait: Whether to wait for encryption to complete
            show_progress: Show progress bars
            **encrypt_kwargs: Additional arguments for encrypt_dataset
        
        Returns:
            Dataset ID of encrypted dataset
        
        Example:
            ```python
            dataset_id = client.encrypt_and_upload_dataset(
                local_file="./my_dataset.parquet",
                dataset_name="my_dataset",
                text_fields=["text", "summary"],
                block_size=1024
            )
            ```
        """
        logger.info(f"Starting encrypt_and_upload_dataset workflow for {dataset_name}")
        
        # 1. Upload dataset to TEE
        remote_datasets_dir = self.config.sftp_datasets_dir
        remote_file_path = f"{remote_datasets_dir}/{dataset_name}.parquet"
        
        logger.info(f"Uploading dataset to {remote_file_path}")
        with self.sftp:
            self.sftp.upload_file(
                local_file=local_file,
                remote_file=remote_file_path
            )
        
        # 2. Encrypt dataset
        logger.info(f"Encrypting dataset {dataset_name}")
        
        encrypted_name = f"{dataset_name}_encrypted"
        job_id = self.admin.encrypt_dataset(
            dataset_id=dataset_name,
            output_dataset_id=encrypted_name,
            text_fields=text_fields,
            **encrypt_kwargs
        )
        
        logger.info(f"Encryption job started: {job_id}")
        
        # 3. Wait for encryption
        if wait:
            self.admin.wait_for_job(job_id, show_progress=show_progress)
            logger.info("Encryption completed")
        
        # 4. Download encrypted dataset if requested
        if output_path and wait:
            remote_encrypted_path = f"{remote_datasets_dir}/{encrypted_name}.parquet"
            logger.info(f"Downloading encrypted dataset to {output_path}")
            
            with self.sftp:
                self.sftp.download_file(
                    remote_file=remote_encrypted_path,
                    local_file=output_path
                )
        
        return encrypted_name
    
    def serve_model(
        self,
        model_id: str,
        serving_mode: str = "OpenAI",
        server_url: Optional[str] = None,
        api_key: Optional[str] = None,
        **start_kwargs
    ) -> Dict[str, Any]:
        """
        Start serving a model.
        
        Convenience wrapper around admin.start_serving().
        
        Args:
            model_id: Model to serve
            serving_mode: Serving mode (OpenAI, vLLM, etc.)
            server_url: Optional custom server URL
            api_key: Optional custom API key
            **start_kwargs: Additional arguments for start_serving
        
        Returns:
            Serving status response
        
        Example:
            ```python
            response = client.serve_model(
                model_id="llama-3.2-3b-encrypted",
                serving_mode="OpenAI"
            )
            ```
        """
        return self.admin.start_serving(
            model_id=model_id,
            serving_mode=serving_mode,
            server_url=server_url,
            api_key=api_key,
            **start_kwargs
        )
    
    def stop_serving(self) -> Dict[str, Any]:
        """
        Stop the currently serving model.
        
        Returns:
            Stop status response
        """
        return self.admin.stop_serving()
    
    # ========================
    # Model Management
    # ========================
    
    def upload_model(
        self,
        local_path: str,
        model_name: str,
        show_progress: bool = True
    ) -> str:
        """
        Upload a model to TEE.
        
        Args:
            local_path: Local model directory
            model_name: Name for the model on TEE
            show_progress: Show upload progress
        
        Returns:
            Remote model path
        """
        remote_path = f"{self.config.sftp_models_dir}/{model_name}"
        
        with self.sftp:
            self.sftp.upload_directory(
                local_path=local_path,
                remote_path=remote_path,
                show_progress=show_progress
            )
        
        return remote_path
    
    def download_model(
        self,
        model_name: str,
        local_path: str,
        show_progress: bool = True
    ) -> str:
        """
        Download a model from TEE.
        
        Args:
            model_name: Name of model on TEE
            local_path: Local directory to download to
            show_progress: Show download progress
        
        Returns:
            Local model path
        """
        remote_path = f"{self.config.sftp_models_dir}/{model_name}"
        
        with self.sftp:
            self.sftp.download_directory(
                remote_path=remote_path,
                local_path=local_path,
                show_progress=show_progress
            )
        
        return local_path
    
    def list_models(self, include_offline: bool = True) -> List[Dict[str, Any]]:
        """
        List available models.
        
        Args:
            include_offline: Include offline models
        
        Returns:
            List of model information
        """
        if include_offline:
            return self.admin.list_models()
        else:
            return self.admin.list_online_models()
    
    def list_online_models(self) -> List[str]:
        """
        Get list of currently serving models.
        
        Returns:
            List of online model IDs
        """
        return self.admin.list_online_models()
    
    def get_model_info(self, model_id: str) -> Dict[str, Any]:
        """
        Get information about a specific model.
        
        Args:
            model_id: Model identifier
        
        Returns:
            Model information
        """
        return self.admin.get_model_info(model_id)
    
    def encrypt_model(
        self,
        model_name_or_path: str,
        out_encrypted_model_path: str,
        server_ip: str = "fhenom_ai_server",
        server_port: int = 9100,
        encryption_impl: str = "decoder-only-llm",
        preprocessing_impl: str = "default",
        inference_mode: str = "double_tokenizer",
        driver_mode: str = "safetensor",
        dtype: str = "bfloat16",
        model_key_password: Optional[str] = None,
        encrypted_model_id: Optional[str] = None,
        preprocessing_args: Optional[List[str]] = None,
        check_duplicate: bool = True
    ) -> Dict[str, Any]:
        """
        Submit model encryption job.
        
        Args:
            model_name_or_path: Model name or path (upload/ prefix added automatically if missing)
            out_encrypted_model_path: Output name or path (download/ prefix added automatically if missing)
            server_ip: FHEnom AI Server hostname/IP
            server_port: Server admin port
            encryption_impl: Encryption method
            preprocessing_impl: Preprocessing strategy
            inference_mode: Inference mode
            driver_mode: Model driver mode
            dtype: Data type
            model_key_password: Optional password
            encrypted_model_id: Optional custom model ID
            preprocessing_args: Optional preprocessing arguments
            check_duplicate: Check for duplicate output names and raise warning (default: True)
        
        Returns:
            Response with job_id
            
        Raises:
            ValueError: If duplicate name found and check_duplicate=True
        """
        # Normalize paths: add upload/ or download/ if not present
        if not model_name_or_path.startswith('upload/') and not model_name_or_path.startswith('/models/upload/'):
            model_name_or_path = f"upload/{model_name_or_path}"
        
        if not out_encrypted_model_path.startswith('download/') and not out_encrypted_model_path.startswith('/models/download/'):
            out_encrypted_model_path = f"download/{out_encrypted_model_path}"
        
        # Convert to full /models/ paths if needed
        if model_name_or_path.startswith('upload/') and not model_name_or_path.startswith('/models/'):
            model_name_or_path = f"/models/{model_name_or_path}"
        
        if out_encrypted_model_path.startswith('download/') and not out_encrypted_model_path.startswith('/models/'):
            out_encrypted_model_path = f"/models/{out_encrypted_model_path}"
        
        # Extract model name for duplicate check
        output_model_name = out_encrypted_model_path.replace('download/', '').replace('/models/download/', '')
        
        # Check for duplicate names if requested
        if check_duplicate:
            try:
                model_ids = self.admin.list_models()
                existing_names = []
                
                for model_id in model_ids:
                    try:
                        info = self.admin.get_model_info(model_id)
                        server_path = info.get('encrypted_model_server_path', '')
                        if server_path.startswith('/models/download/'):
                            existing_name = server_path.replace('/models/download/', '', 1)
                            if existing_name == output_model_name:
                                existing_names.append(model_id)
                    except:
                        pass
                
                if existing_names:
                    raise ValueError(
                        f"The output name '{output_model_name}' is already used by {len(existing_names)} model(s): "
                        f"{', '.join(existing_names[:3])}{'...' if len(existing_names) > 3 else ''}. "
                        f"The new model will have a unique ID, but this may cause confusion. "
                        f"Use check_duplicate=False to skip this check or choose a different name."
                    )
            except ValueError:
                raise  # Re-raise ValueError
            except Exception:
                pass  # Ignore other errors during duplicate check
        
        return self.admin.encrypt_model(
            model_name_or_path=model_name_or_path,
            out_encrypted_model_path=out_encrypted_model_path,
            server_ip=server_ip,
            server_port=server_port,
            encryption_impl=encryption_impl,
            preprocessing_impl=preprocessing_impl,
            inference_mode=inference_mode,
            driver_mode=driver_mode,
            dtype=dtype,
            model_key_password=model_key_password,
            encrypted_model_id=encrypted_model_id,
            preprocessing_args=preprocessing_args
        )
    
    def get_job_status(self, job_id: str) -> Dict[str, Any]:
        """
        Get encryption job status.
        
        Args:
            job_id: Job identifier
        
        Returns:
            Job status information
        """
        return self.admin.get_job_status(job_id)
    
    def delete_model(self, model_name: str, local_only: bool = False):
        """
        Delete a model from TEE.
        
        Args:
            model_name: Name of model to delete
            local_only: Only delete local files, keep SFTP
        """
        if not local_only:
            remote_path = f"{self.config.sftp_models_dir}/{model_name}"
            with self.sftp:
                self.sftp.delete(remote_path, recursive=True)
            logger.info(f"Deleted model from TEE: {model_name}")
    
    # ========================
    # Dataset Management
    # ========================
    
    def upload_dataset(
        self,
        local_file: str,
        dataset_name: str
    ) -> str:
        """
        Upload a dataset to TEE.
        
        Args:
            local_file: Local dataset file
            dataset_name: Name for dataset on TEE
        
        Returns:
            Remote dataset path
        """
        remote_path = f"{self.config.sftp_datasets_dir}/{dataset_name}.parquet"
        
        with self.sftp:
            self.sftp.upload_file(
                local_file=local_file,
                remote_file=remote_path
            )
        
        return remote_path
    
    def download_dataset(
        self,
        dataset_name: str,
        local_file: str
    ) -> str:
        """
        Download a dataset from TEE.
        
        Args:
            dataset_name: Name of dataset on TEE
            local_file: Local file path to download to
        
        Returns:
            Local file path
        """
        remote_path = f"{self.config.sftp_datasets_dir}/{dataset_name}.parquet"
        
        with self.sftp:
            self.sftp.download_file(
                remote_file=remote_path,
                local_file=local_file
            )
        
        return local_file
    
    # ========================
    # Job Management
    # ========================
    
    def get_job_status(self, job_id: str) -> Dict[str, Any]:
        """
        Get status of a job.
        
        Args:
            job_id: Job identifier
        
        Returns:
            Job status information
        """
        return self.admin.get_job_status(job_id)
    
    def wait_for_job(
        self,
        job_id: str,
        poll_interval: float = 2.0,
        timeout: Optional[float] = None,
        callback: Optional[Callable[[Dict[str, Any]], None]] = None,
        show_progress: bool = True
    ) -> Dict[str, Any]:
        """
        Wait for a job to complete.
        
        Args:
            job_id: Job identifier
            poll_interval: Seconds between status checks
            timeout: Maximum seconds to wait
            callback: Optional callback for each status check
            show_progress: Show progress updates
        
        Returns:
            Final job status
        """
        return self.admin.wait_for_job(
            job_id=job_id,
            poll_interval=poll_interval,
            timeout=timeout,
            callback=callback,
            show_progress=show_progress
        )
    
    # ========================
    # Health & Monitoring
    # ========================
    
    def get_health(self) -> Dict[str, Any]:
        """
        Get TEE system health status.
        
        Returns:
            Health information (placeholder)
        """
        return self.admin.get_health()
    
    def get_metrics(self) -> Dict[str, Any]:
        """
        Get TEE performance metrics.
        
        Returns:
            Metrics information (placeholder)
        """
        return self.admin.get_metrics()
    
    def get_tee_info(self) -> Dict[str, Any]:
        """
        Get TEE version and configuration information.
        
        Returns:
            TEE information (placeholder)
        """
        return self.admin.get_tee_info()
