#!/usr/bin/env python3
"""
Setup script to configure nanobot to use Llama models with AirLLM.
This script will:
1. Check/create the config file
2. Set up Llama model configuration
3. Guide you through getting a Hugging Face token if needed
"""

import json
import os
from pathlib import Path

CONFIG_PATH = Path.home() / ".nanobot" / "config.json"

def get_hf_token_instructions():
    """Print instructions for getting a Hugging Face token."""
    print("\n" + "="*70)
    print("GETTING A HUGGING FACE TOKEN")
    print("="*70)
    print("\nTo use Llama models (which are gated), you need a Hugging Face token:")
    print("\n1. Go to: https://huggingface.co/settings/tokens")
    print("2. Click 'New token'")
    print("3. Give it a name (e.g., 'nanobot')")
    print("4. Select 'Read' permission")
    print("5. Click 'Generate token'")
    print("6. Copy the token (starts with 'hf_...')")
    print("\nThen accept the Llama model license:")
    print("1. Go to: https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct")
    print("2. Click 'Agree and access repository'")
    print("3. Accept the license terms")
    print("\n" + "="*70 + "\n")

def load_existing_config():
    """Load existing config or return default."""
    if CONFIG_PATH.exists():
        try:
            with open(CONFIG_PATH) as f:
                return json.load(f)
        except Exception as e:
            print(f"Warning: Could not read existing config: {e}")
            return {}
    return {}

def create_llama_config():
    """Create or update config for Llama with AirLLM."""
    config = load_existing_config()
    
    # Ensure providers section exists
    if "providers" not in config:
        config["providers"] = {}
    
    # Ensure agents section exists
    if "agents" not in config:
        config["agents"] = {}
    if "defaults" not in config["agents"]:
        config["agents"]["defaults"] = {}
    
    # Choose Llama model
    print("\n" + "="*70)
    print("CHOOSE LLAMA MODEL")
    print("="*70)
    print("\nAvailable models:")
    print("  1. Llama-3.2-3B-Instruct (Recommended - fast, minimal memory)")
    print("  2. Llama-3.1-8B-Instruct (Good balance of performance and speed)")
    print("  3. Custom (enter model path)")
    
    choice = input("\nChoose model (1-3, default: 1): ").strip() or "1"
    
    model_map = {
        "1": "meta-llama/Llama-3.2-3B-Instruct",
        "2": "meta-llama/Llama-3.1-8B-Instruct",
    }
    
    if choice == "3":
        model_path = input("Enter model path (e.g., meta-llama/Llama-3.2-3B-Instruct): ").strip()
        if not model_path:
            model_path = "meta-llama/Llama-3.2-3B-Instruct"
            print(f"Using default: {model_path}")
    else:
        model_path = model_map.get(choice, "meta-llama/Llama-3.2-3B-Instruct")
    
    # Set up AirLLM provider with Llama model
    # Note: apiKey can be used as model path, or we can put model in defaults
    config["providers"]["airllm"] = {
        "apiKey": "",  # Will be set to model path
        "apiBase": None,
        "extraHeaders": {}
    }
    
    # Set default model
    config["agents"]["defaults"]["model"] = model_path
    
    # Ask for Hugging Face token
    print("\n" + "="*70)
    print("HUGGING FACE TOKEN SETUP")
    print("="*70)
    print("\nDo you have a Hugging Face token? (Required for Llama models)")
    print("If not, we'll show you how to get one.\n")
    
    has_token = input("Do you have a Hugging Face token? (y/n): ").strip().lower()
    
    if has_token == 'y':
        hf_token = input("\nEnter your Hugging Face token (starts with 'hf_'): ").strip()
        if hf_token and hf_token.startswith('hf_'):
            # Store token in extraHeaders
            config["providers"]["airllm"]["extraHeaders"]["hf_token"] = hf_token
            # Also set apiKey to model path (AirLLM uses apiKey as model path if it contains '/')
            config["providers"]["airllm"]["apiKey"] = config["agents"]["defaults"]["model"]
            print("\n✓ Token configured!")
        else:
            print("⚠ Warning: Token doesn't look valid (should start with 'hf_')")
            print("You can add it later by editing the config file.")
            # Still set model path in apiKey
            config["providers"]["airllm"]["apiKey"] = config["agents"]["defaults"]["model"]
    else:
        get_hf_token_instructions()
        print("\nYou can add your token later by:")
        print(f"1. Editing: {CONFIG_PATH}")
        print("2. Adding your token to: providers.airllm.extraHeaders.hf_token")
        print("\nOr run this script again after getting your token.")
    
    return config

def save_config(config):
    """Save config to file."""
    CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
    with open(CONFIG_PATH, 'w') as f:
        json.dump(config, f, indent=2)
    
    # Set secure permissions
    os.chmod(CONFIG_PATH, 0o600)
    print(f"\n✓ Configuration saved to: {CONFIG_PATH}")
    print(f"✓ File permissions set to 600 (read/write for owner only)")

def main():
    """Main setup function."""
    print("\n" + "="*70)
    print("NANOBOT LLAMA + AIRLLM SETUP")
    print("="*70)
    print("\nThis script will configure nanobot to use Llama models with AirLLM.\n")
    
    if CONFIG_PATH.exists():
        print(f"Found existing config at: {CONFIG_PATH}")
        backup = input("\nCreate backup? (y/n): ").strip().lower()
        if backup == 'y':
            backup_path = CONFIG_PATH.with_suffix('.json.backup')
            import shutil
            shutil.copy(CONFIG_PATH, backup_path)
            print(f"✓ Backup created: {backup_path}")
    else:
        print(f"Creating new config at: {CONFIG_PATH}")
    
    config = create_llama_config()
    save_config(config)
    
    print("\n" + "="*70)
    print("SETUP COMPLETE!")
    print("="*70)
    print("\nConfiguration:")
    print(f"  Model: {config['agents']['defaults']['model']}")
    print(f"  Provider: airllm")
    if config["providers"]["airllm"].get("extraHeaders", {}).get("hf_token"):
        print(f"  HF Token: {'*' * 20} (configured)")
    else:
        print(f"  HF Token: Not configured (add it to use gated models)")
    
    print("\nNext steps:")
    print("  1. If you need a Hugging Face token, follow the instructions above")
    print("  2. Test it: nanobot agent -m 'Hello, what is 2+5?'")
    print("\n" + "="*70 + "\n")

if __name__ == "__main__":
    main()