"""
Example: Complete workflow with encryption and serving.
"""

from fhenomai import FHEnomClient, FHEnomConfig
import logging

# Enable logging to see progress
logging.basicConfig(level=logging.INFO)

def main():
    # Load configuration
    config = FHEnomConfig.from_file()
    
    # Use context manager for automatic SFTP connection management
    with FHEnomClient(config) as client:
        
        # Step 1: Upload model
        print("=" * 60)
        print("Step 1: Uploading model")
        print("=" * 60)
        
        model_name = "llama-3.2-3b"
        client.upload_model(
            local_path="./models/llama-3.2-3b",
            model_name=model_name,
            show_progress=True
        )
        
        # Step 2: Encrypt model
        print("\n" + "=" * 60)
        print("Step 2: Encrypting model")
        print("=" * 60)
        
        encrypted_model_id = f"{model_name}_encrypted"
        job_id = client.admin.encrypt_model(
            model_id=model_name,
            output_model_id=encrypted_model_id,
            block_size=1024,
            add_residual_term=True,
            quantize_encrypted_model=True,
            quantization_bits=4
        )
        
        print(f"Encryption job started: {job_id}")
        
        # Wait for encryption to complete
        result = client.wait_for_job(job_id, show_progress=True)
        
        if result['status'] == 'completed':
            print("✓ Encryption completed successfully")
        else:
            print(f"✗ Encryption failed: {result.get('error', 'Unknown error')}")
            return
        
        # Step 3: Download encrypted model (optional)
        print("\n" + "=" * 60)
        print("Step 3: Downloading encrypted model")
        print("=" * 60)
        
        client.download_model(
            model_name=encrypted_model_id,
            local_path=f"./models/{encrypted_model_id}",
            show_progress=True
        )
        
        # Step 4: Start serving
        print("\n" + "=" * 60)
        print("Step 4: Starting model serving")
        print("=" * 60)
        
        client.serve_model(
            model_id=encrypted_model_id,
            serving_mode="OpenAI"
        )
        
        print(f"✓ Model {encrypted_model_id} is now serving!")
        print(f"  Use OpenAI SDK with base_url={config.user_url}")
        
        # Step 5: Test inference (using OpenAI SDK)
        print("\n" + "=" * 60)
        print("Step 5: Testing inference")
        print("=" * 60)
        
        try:
            from openai import OpenAI
            
            openai_client = OpenAI(
                base_url=config.user_url,
                api_key=config.user_api_key
            )
            
            response = openai_client.chat.completions.create(
                model=encrypted_model_id,
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": "What is homomorphic encryption?"}
                ],
                max_tokens=100
            )
            
            print("\nInference successful!")
            print(f"Response: {response.choices[0].message.content}")
        
        except ImportError:
            print("OpenAI SDK not installed. Install with: pip install openai")
        except Exception as e:
            print(f"Inference error: {e}")
        
        # Step 6: Stop serving
        print("\n" + "=" * 60)
        print("Step 6: Stopping serving")
        print("=" * 60)
        
        client.stop_serving()
        print("✓ Serving stopped")
        
        print("\n" + "=" * 60)
        print("Workflow completed successfully!")
        print("=" * 60)

if __name__ == '__main__':
    main()
