"""
Basic usage example for FHEnom AI Python library.
"""

from fhenomai import FHEnomClient, FHEnomConfig

# Load configuration from file or environment
config = FHEnomConfig.from_file()  # or FHEnomConfig.from_env()

# Initialize client
client = FHEnomClient(config)

# List available models
models = client.admin.list_models()
print(f"Found {len(models)} encrypted models:")
for model in models[:5]:  # Show first 5
    print(f"  - {model}")

# Get SFTP manager for file operations
sftp = client.get_sftp_manager()

# Upload a model (example - adjust paths as needed)
print("\nUploading model to TEE...")
# sftp.upload_directory("./my-model", "my-model")  # upload/ prefix added automatically

# Encrypt the model (paths auto-normalized)
print("\nEncrypting model...")
job_id = client.admin.encrypt_model(
    model_name_or_path="my-model",  # Becomes /models/upload/my-model
    out_encrypted_model_path="my-model-encrypted",  # Becomes /models/download/my-model-encrypted
    symbolic_name="my-model-encrypted",  # Symbolic name for the model
    encrypted_model_id="my-model-encrypted",
    encryption_impl="decoder-only-llm",
    dtype="bfloat16"
)
print(f"Encryption job started: {job_id}")

# Wait for completion with progress callback
def progress_callback(status):
    progress = status.get('progress', 0) * 100
    message = status.get('message', 'Processing')
    print(f"\r{message}: {progress:.1f}%", end='', flush=True)

result = client.admin.wait_for_job(
    job_id,
    timeout=3600,
    poll_interval=5,
    callback=progress_callback
)
print(f"\n\nEncryption completed: {result.get('status')}")

# Download encrypted model (download/ prefix added automatically)
if result.get('status') == 'done':
    print("\nDownloading encrypted model...")
    # sftp.download_directory("my-model-encrypted", "./encrypted/my-model")
    print("Model ready for serving!")

# Start serving
print("\nStarting model serving...")
client.admin.start_serving(
    encrypted_model_id="my-model-encrypted",
    server_url="http://YOUR_VLLM_SERVER_IP:8000",  # vLLM server IP/hostname
    display_model_name="my-model"  # Optional: for vLLM --served-model-name
)
print("Model is now serving!")

# Use OpenAI SDK for inference (user-facing API)
from openai import OpenAI

openai_client = OpenAI(
    base_url=config.user_url,
    api_key="not-needed"  # TEE doesn't require authentication
)

response = openai_client.chat.completions.create(
    model="my-model",  # Use display name
    messages=[
        {"role": "user", "content": "Hello! How are you?"}
    ],
    max_tokens=100
)
print(f"\nModel response: {response.choices[0].message.content}")

# Stop serving when done
print("\nStopping serving...")
client.admin.stop_serving("my-model-encrypted")
print("Done!")
