#!/usr/bin/env python3
"""
Script to move model files from the 'econ' subdirectory to the root directory
for the Hugging Face repository: Gaston895/aegisconduct
"""

import os
import shutil
from pathlib import Path
import sys

def move_files_to_root(repo_path="."):
    """
    Moves all files from the 'econ' subdirectory to the repository root.
    
    Args:
        repo_path (str): Path to the local clone of the repository.
    """
    
    # Define paths
    repo_dir = Path(repo_path).resolve()
    econ_dir = repo_dir / "econ"
    root_dir = repo_dir
    
    print(f"Repository root: {root_dir}")
    print(f"Econ subdirectory: {econ_dir}")
    
    # Check if 'econ' directory exists
    if not econ_dir.exists() or not econ_dir.is_dir():
        print(f"❌ Error: 'econ' subdirectory not found at {econ_dir}")
        print("Please ensure you are in the correct directory and the 'econ' folder exists.")
        return False
    
    # List files in the 'econ' directory
    files_to_move = list(econ_dir.iterdir())
    
    if not files_to_move:
        print("ℹ️  No files found in the 'econ' directory.")
        return True
    
    print(f"📁 Found {len(files_to_move)} files/directories in 'econ':")
    for item in files_to_move:
        print(f"   - {item.name}")
    
    # Move files
    moved_count = 0
    for item in files_to_move:
        source_path = item
        dest_path = root_dir / item.name
        
        # Check if a file with the same name already exists in root
        if dest_path.exists():
            print(f"⚠️  Warning: {item.name} already exists in root. Skipping...")
            continue
        
        try:
            # Move the file/directory
            shutil.move(str(source_path), str(dest_path))
            moved_count += 1
            print(f"✅ Moved: {item.name}")
        except Exception as e:
            print(f"❌ Failed to move {item.name}: {e}")
    
    # Check if 'econ' directory is now empty and remove it
    try:
        if not any(econ_dir.iterdir()):
            econ_dir.rmdir()
            print(f"🗑️  Removed empty 'econ' directory")
    except Exception as e:
        print(f"⚠️  Could not remove 'econ' directory: {e}")
    
    print(f"\n🎉 Successfully moved {moved_count} out of {len(files_to_move)} items to the root directory.")
    
    if moved_count < len(files_to_move):
        print("💡 Some files may not have been moved due to conflicts. Please review manually.")
    
    return True

def update_config_json(repo_path="."):
    """
    Updates the config.json file if it references the old 'econ' path.
    """
    
    config_path = Path(repo_path) / "config.json"
    
    if config_path.exists():
        try:
            import json
            with open(config_path, 'r', encoding='utf-8') as f:
                config = json.load(f)
            
            # Check if config needs updating
            needs_update = False
            # You can add specific checks here based on your config structure
            
            if needs_update:
                with open(config_path, 'w', encoding='utf-8') as f:
                    json.dump(config, f, indent=2)
                print("✅ Updated config.json")
            else:
                print("ℹ️  config.json doesn't appear to need updates")
                
        except Exception as e:
            print(f"⚠️  Could not check/update config.json: {e}")

def main():
    """Main function to orchestrate the file movement."""
    
    print("=" * 60)
    print("AEGISCONDUCT MODEL FILES REORGANIZATION")
    print("=" * 60)
    print("This script moves files from 'econ' subdirectory to repository root.")
    print(f"Repository: Gaston895/aegisconduct")
    print("=" * 60)
    
    # Ask for confirmation
    response = input("\n⚠️  WARNING: This will modify your local repository structure.\nDo you want to continue? (yes/no): ").strip().lower()
    
    if response not in ['yes', 'y']:
        print("Operation cancelled.")
        return
    
    # Step 1: Move files
    print("\n" + "=" * 60)
    print("STEP 1: Moving files from 'econ' to root directory")
    print("=" * 60)
    
    success = move_files_to_root()
    
    if not success:
        print("❌ File movement failed. Exiting.")
        return
    
    # Step 2: Update config if needed
    print("\n" + "=" * 60)
    print("STEP 2: Checking configuration files")
    print("=" * 60)
    
    update_config_json()
    
    # Step 3: Instructions for next steps
    print("\n" + "=" * 60)
    print("NEXT STEPS")
    print("=" * 60)
    print("1. Review the moved files in your repository root")
    print("2. Update your application code to load from root (not 'econ' subfolder)")
    print("3. Test that the model loads correctly:")
    print("   - From Python: model = AutoModelForCausalLM.from_pretrained('Gaston895/aegisconduct')")
    print("   - No 'subfolder' or 'revision' parameter needed")
    print("4. Commit and push changes to Hugging Face Hub:")
    print("   git add .")
    print("   git commit -m 'Move model files from econ subdir to root'")
    print("   git push origin main")
    print("=" * 60)
    
    print("\n✅ Script completed successfully!")

if __name__ == "__main__":
    main()