"""
Testing helper functions for validating FHEnom AI setup.
"""

import logging
from typing import Dict, Any, Optional, List
from dataclasses import dataclass
from enum import Enum

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

logger = logging.getLogger(__name__)


class TestStatus(Enum):
    """Test result status."""
    PASSED = "passed"
    FAILED = "failed"
    SKIPPED = "skipped"
    WARNING = "warning"


@dataclass
class TestResult:
    """Result of a test."""
    name: str
    status: TestStatus
    message: str
    details: Optional[Dict[str, Any]] = None
    
    def __str__(self) -> str:
        icon = {
            TestStatus.PASSED: "✓",
            TestStatus.FAILED: "✗",
            TestStatus.SKIPPED: "○",
            TestStatus.WARNING: "⚠"
        }[self.status]
        
        return f"{icon} {self.name}: {self.message}"


def validate_connection(
    config: Optional[FHEnomConfig] = None,
    verbose: bool = True
) -> TestResult:
    """
    Validate basic connectivity to FHEnom AI services.
    
    Args:
        config: Configuration (loads from file if not provided)
        verbose: Print detailed output
    
    Returns:
        Test result
    
    Example:
        ```python
        from fhenomai.testing import validate_connection
        
        result = validate_connection()
        print(result)
        ```
    """
    if config is None:
        try:
            config = FHEnomConfig.from_file()
        except FileNotFoundError:
            config = FHEnomConfig.from_env()
    
    try:
        # Try to connect to admin API
        admin = AdminAPI(
            base_url=config.admin_url,
            api_key=config.admin_api_key,
            timeout=config.timeout
        )
        
        # Simple health check - list models
        models = admin.list_models()
        
        if verbose:
            logger.info(f"Connected to Admin API at {config.admin_url}")
            logger.info(f"Found {len(models)} models")
        
        return TestResult(
            name="Connection Test",
            status=TestStatus.PASSED,
            message=f"Successfully connected to {config.admin_url}",
            details={"model_count": len(models)}
        )
    
    except Exception as e:
        return TestResult(
            name="Connection Test",
            status=TestStatus.FAILED,
            message=f"Failed to connect: {str(e)}"
        )


def validate_sftp(
    config: Optional[FHEnomConfig] = None,
    test_upload: bool = True,
    verbose: bool = True
) -> TestResult:
    """
    Validate SFTP connectivity and permissions.
    
    Args:
        config: Configuration (loads from file if not provided)
        test_upload: Test file upload/download
        verbose: Print detailed output
    
    Returns:
        Test result
    
    Example:
        ```python
        from fhenomai.testing import validate_sftp
        
        result = validate_sftp()
        print(result)
        ```
    """
    if config is None:
        try:
            config = FHEnomConfig.from_file()
        except FileNotFoundError:
            config = FHEnomConfig.from_env()
    
    try:
        sftp = SFTPClient(
            host=config.sftp_host,
            username=config.sftp_username,
            port=config.sftp_port,
            password=config.sftp_password,
            key_path=config.sftp_key_path
        )
        
        with sftp:
            # List models directory
            try:
                models = sftp.list_directory(config.sftp_models_dir)
                if verbose:
                    logger.info(f"Connected to SFTP at {config.sftp_host}")
                    logger.info(f"Found {len(models)} items in models directory")
            except Exception as e:
                return TestResult(
                    name="SFTP Test",
                    status=TestStatus.WARNING,
                    message=f"Connected but cannot list models directory: {str(e)}"
                )
            
            # Test upload/download if requested
            if test_upload:
                import tempfile
                import os
                
                test_file_name = ".fhenomai_test.txt"
                test_content = b"FHEnom AI test file"
                
                with tempfile.NamedTemporaryFile(delete=False) as tmp:
                    tmp.write(test_content)
                    tmp_path = tmp.name
                
                try:
                    # Upload
                    remote_path = f"{config.sftp_models_dir}/{test_file_name}"
                    sftp.upload_file(tmp_path, remote_path)
                    
                    # Download
                    download_path = tmp_path + ".download"
                    sftp.download_file(remote_path, download_path)
                    
                    # Verify
                    with open(download_path, 'rb') as f:
                        downloaded_content = f.read()
                    
                    if downloaded_content != test_content:
                        return TestResult(
                            name="SFTP Test",
                            status=TestStatus.FAILED,
                            message="Upload/download verification failed"
                        )
                    
                    # Cleanup
                    sftp.delete(remote_path)
                    os.unlink(tmp_path)
                    os.unlink(download_path)
                    
                    if verbose:
                        logger.info("SFTP upload/download test passed")
                
                except Exception as e:
                    return TestResult(
                        name="SFTP Test",
                        status=TestStatus.FAILED,
                        message=f"Upload/download test failed: {str(e)}"
                    )
        
        return TestResult(
            name="SFTP Test",
            status=TestStatus.PASSED,
            message=f"SFTP connectivity verified",
            details={"host": config.sftp_host, "port": config.sftp_port}
        )
    
    except Exception as e:
        return TestResult(
            name="SFTP Test",
            status=TestStatus.FAILED,
            message=f"SFTP connection failed: {str(e)}"
        )


def validate_admin_api(
    config: Optional[FHEnomConfig] = None,
    verbose: bool = True
) -> List[TestResult]:
    """
    Validate all Admin API endpoints.
    
    Args:
        config: Configuration (loads from file if not provided)
        verbose: Print detailed output
    
    Returns:
        List of test results
    
    Example:
        ```python
        from fhenomai.testing import validate_admin_api
        
        results = validate_admin_api()
        for result in results:
            print(result)
        ```
    """
    if config is None:
        try:
            config = FHEnomConfig.from_file()
        except FileNotFoundError:
            config = FHEnomConfig.from_env()
    
    results = []
    admin = AdminAPI(
        base_url=config.admin_url,
        api_key=config.admin_api_key,
        timeout=config.timeout
    )
    
    # Test 1: List models
    try:
        models = admin.list_models()
        results.append(TestResult(
            name="List Models",
            status=TestStatus.PASSED,
            message=f"Found {len(models)} models",
            details={"count": len(models)}
        ))
        if verbose:
            logger.info(f"✓ List models: {len(models)} found")
    except Exception as e:
        results.append(TestResult(
            name="List Models",
            status=TestStatus.FAILED,
            message=str(e)
        ))
    
    # Test 2: List online models
    try:
        online_models = admin.list_online_models()
        results.append(TestResult(
            name="List Online Models",
            status=TestStatus.PASSED,
            message=f"Found {len(online_models)} online models",
            details={"count": len(online_models)}
        ))
        if verbose:
            logger.info(f"✓ List online models: {len(online_models)} found")
    except Exception as e:
        results.append(TestResult(
            name="List Online Models",
            status=TestStatus.FAILED,
            message=str(e)
        ))
    
    # Test 3: Get model info (if models exist)
    if models:
        try:
            first_model_id = models[0]["id"]
            info = admin.get_model_info(first_model_id)
            results.append(TestResult(
                name="Get Model Info",
                status=TestStatus.PASSED,
                message=f"Retrieved info for {first_model_id}",
                details={"model_id": first_model_id}
            ))
            if verbose:
                logger.info(f"✓ Get model info: {first_model_id}")
        except Exception as e:
            results.append(TestResult(
                name="Get Model Info",
                status=TestStatus.FAILED,
                message=str(e)
            ))
    else:
        results.append(TestResult(
            name="Get Model Info",
            status=TestStatus.SKIPPED,
            message="No models available to test"
        ))
    
    return results


def run_full_workflow_test(
    model_path: str,
    config: Optional[FHEnomConfig] = None,
    cleanup: bool = True,
    verbose: bool = True
) -> List[TestResult]:
    """
    Run a complete end-to-end workflow test.
    
    Tests: upload → encrypt → serve → inference → cleanup
    
    Args:
        model_path: Path to test model
        config: Configuration (loads from file if not provided)
        cleanup: Remove test artifacts
        verbose: Print detailed output
    
    Returns:
        List of test results
    
    Example:
        ```python
        from fhenomai.testing import run_full_workflow_test
        
        results = run_full_workflow_test("./test-model")
        for result in results:
            print(result)
        ```
    """
    if config is None:
        try:
            config = FHEnomConfig.from_file()
        except FileNotFoundError:
            config = FHEnomConfig.from_env()
    
    from ..client import FHEnomClient
    import os
    
    results = []
    client = FHEnomClient(config)
    test_model_name = "fhenomai_test_model"
    
    # Test 1: Upload model
    try:
        if verbose:
            logger.info(f"Uploading test model from {model_path}")
        
        remote_path = client.upload_model(
            local_path=model_path,
            model_name=test_model_name,
            show_progress=verbose
        )
        
        results.append(TestResult(
            name="Upload Model",
            status=TestStatus.PASSED,
            message=f"Uploaded to {remote_path}"
        ))
        
        if verbose:
            logger.info(f"✓ Upload completed: {remote_path}")
    
    except Exception as e:
        results.append(TestResult(
            name="Upload Model",
            status=TestStatus.FAILED,
            message=str(e)
        ))
        return results  # Cannot continue without upload
    
    # Test 2: Encrypt model
    try:
        if verbose:
            logger.info("Encrypting test model")
        
        encrypted_name = f"{test_model_name}_encrypted"
        job_id = client.admin.encrypt_model(
            model_name_or_path=test_model_name,
            out_encrypted_model_path=encrypted_name,
            symbolic_name=encrypted_name,
            encrypted_model_id=encrypted_name,
            encryption_impl="decoder-only-llm",
            dtype="bfloat16"
        )
        
        if verbose:
            logger.info(f"Encryption job started: {job_id}")
        
        # Wait for completion
        final_status = client.wait_for_job(job_id, show_progress=verbose)
        
        if final_status["status"] == "completed":
            results.append(TestResult(
                name="Encrypt Model",
                status=TestStatus.PASSED,
                message=f"Encryption completed: {encrypted_name}",
                details={"job_id": job_id}
            ))
            if verbose:
                logger.info(f"✓ Encryption completed")
        else:
            results.append(TestResult(
                name="Encrypt Model",
                status=TestStatus.FAILED,
                message=f"Encryption failed: {final_status.get('error', 'Unknown error')}"
            ))
            return results
    
    except Exception as e:
        results.append(TestResult(
            name="Encrypt Model",
            status=TestStatus.FAILED,
            message=str(e)
        ))
        return results
    
    # Test 3: Start serving
    try:
        if verbose:
            logger.info("Starting model serving")
        
        serve_response = client.serve_model(
            model_id=encrypted_name,
            serving_mode="OpenAI"
        )
        
        results.append(TestResult(
            name="Start Serving",
            status=TestStatus.PASSED,
            message="Model serving started"
        ))
        
        if verbose:
            logger.info("✓ Serving started")
    
    except Exception as e:
        results.append(TestResult(
            name="Start Serving",
            status=TestStatus.FAILED,
            message=str(e)
        ))
    
    # Test 4: Stop serving
    try:
        if verbose:
            logger.info("Stopping model serving")
        
        client.stop_serving()
        
        results.append(TestResult(
            name="Stop Serving",
            status=TestStatus.PASSED,
            message="Model serving stopped"
        ))
        
        if verbose:
            logger.info("✓ Serving stopped")
    
    except Exception as e:
        results.append(TestResult(
            name="Stop Serving",
            status=TestStatus.FAILED,
            message=str(e)
        ))
    
    # Cleanup
    if cleanup:
        try:
            if verbose:
                logger.info("Cleaning up test artifacts")
            
            client.delete_model(test_model_name)
            client.delete_model(encrypted_name)
            
            if verbose:
                logger.info("✓ Cleanup completed")
        
        except Exception as e:
            logger.warning(f"Cleanup failed: {e}")
    
    return results
