#!/usr/bin/env python3
"""
Local wrapper for cloud training script
"""

import os
import sys
import argparse
from pathlib import Path

def main():
    parser = argparse.ArgumentParser(description='Run cloud training locally')
    parser.add_argument('--model_name', type=str, default='local-morphological-transformer', help='Model name')
    parser.add_argument('--dataset', type=str, default='10L_90NL', choices=['10L_90NL', '50L_50NL', '90L_10NL'], help='Dataset name')
    parser.add_argument('--run', type=str, default='1', choices=['1', '2', '3'], help='Run number')
    parser.add_argument('--output_dir', type=str, default='./cloud_output', help='Output directory')
    parser.add_argument('--model_dir', type=str, default='./cloud_models', help='Model directory')
    parser.add_argument('--wandb_project', type=str, default='morphological-transformer-local', help='WandB project name')
    parser.add_argument('--hf_token', type=str, help='Hugging Face token (optional)')
    args = parser.parse_args()
    
    # Set environment variables
    os.environ['MODEL_NAME'] = args.model_name
    os.environ['DATASET_NAME'] = args.dataset
    os.environ['RUN_NUMBER'] = args.run
    os.environ['DATA_DIR'] = f'./{args.dataset}'
    os.environ['OUTPUT_DIR'] = args.output_dir
    os.environ['MODEL_DIR'] = args.model_dir
    os.environ['WANDB_PROJECT'] = args.wandb_project
    
    if args.hf_token:
        os.environ['HF_TOKEN'] = args.hf_token
    
    print(f"🚀 Starting cloud training with:")
    print(f"  - Model: {args.model_name}")
    print(f"  - Dataset: {args.dataset}")
    print(f"  - Run: {args.run}")
    print(f"  - Output: {args.output_dir}")
    print(f"  - Models: {args.model_dir}")
    print(f"  - WandB: {args.wandb_project}")
    
    # Import and run the cloud training script
    try:
        from scripts.hf_cloud_training import main as cloud_main
        cloud_main()
    except ImportError as e:
        print(f"❌ Failed to import cloud training script: {e}")
        sys.exit(1)
    except Exception as e:
        print(f"❌ Training failed: {e}")
        sys.exit(1)

if __name__ == '__main__':
    main()