"""
Script to automatically update all files with token from .env.
"""

import os
import re
from pathlib import Path
from typing import List, Set
from dotenv import load_dotenv

def load_token() -> str:
    """Load token from .env file."""
    root_dir = Path(__file__).parent.parent
    env_path = root_dir / '.env'
    load_dotenv(dotenv_path=env_path)
    
    token = os.getenv('HF_TOKEN')
    if not token:
        raise ValueError("HF_TOKEN not set in .env file")
    return token

def find_python_files(start_path: str) -> List[Path]:
    """Find all Python files in the directory tree."""
    return [
        Path(os.path.join(root, file))
        for root, _, files in os.walk(start_path)
        for file in files
        if file.endswith('.py') and 'venv' not in root and '__pycache__' not in root
    ]

def has_get_token_cached(content: str) -> bool:
    """Check if file uses get_hf_token_cached."""
    return (
        'get_hf_token_cached' in content and
        'from config import' in content
    )

def update_file(file_path: Path, token: str) -> bool:
    """Update a single file with token."""
    print(f"Checking {file_path}")
    
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Skip files that don't use the config token
    if not has_get_token_cached(content):
        print(f"Skipping {file_path} - no token config usage found")
        return False
    
    print(f"Updating {file_path} with token")
    
    # Add token initialization right after config import
    lines = content.split('\n')
    modified_lines = []
    
    for line in lines:
        modified_lines.append(line)
        if line.strip() == 'from config import get_hf_token_cached':
            # Add token setup after import
            modified_lines.append(f'\n# Initialize token from .env')
            modified_lines.append(f'os.environ["HF_TOKEN"] = "{token}"\n')
    
    content = '\n'.join(modified_lines)
    
    # Write changes
    with open(file_path, 'w', encoding='utf-8') as f:
        f.write(content)
    
    print(f"Updated {file_path}")
    return True

def main():
    """Main function to update all files."""
    try:
        # Get token from .env
        token = load_token()
        
        # Get project root
        root_dir = Path(__file__).parent.parent
        
        # Find and update all Python files
        updated = 0
        python_files = find_python_files(root_dir)
        
        for file_path in python_files:
            try:
                if update_file(file_path, token):
                    updated += 1
            except Exception as e:
                print(f"Error updating {file_path}: {e}")
        
        print(f"\nUpdated {updated} files with token from .env")
        
    except Exception as e:
        print(f"Error: {e}")
        print("Please make sure HF_TOKEN is set in your .env file")
        return 1
    
    return 0

if __name__ == '__main__':
    main()