#!/usr/bin/env python3
"""
Minimal script to automatically fix YAML config signatures by adding missing arguments.
"""

import argparse
import re
import sys
from pathlib import Path
from typing import Dict, List, Set, Tuple

import yaml


def extract_signature_params(signature: str) -> Set[str]:
    """Extract parameter names from a tool signature string."""
    required_pattern = r'<(\w+)>'
    optional_pattern = r'\[<([^>]+)>\]'
    
    required_params = set(re.findall(required_pattern, signature))
    optional_params = set(re.findall(optional_pattern, signature))
    
    # Clean up optional params that might have -- prefix
    cleaned_optional = set()
    for param in optional_params:
        if param.startswith('--'):
            cleaned_optional.add(param[2:].replace('-', '_'))
        else:
            cleaned_optional.add(param)
    
    return required_params | cleaned_optional


def extract_defined_arguments(arguments: List[Dict]) -> Tuple[Set[str], Set[str]]:
    """Extract required and optional argument names from the arguments definition."""
    required_args = set()
    optional_args = set()
    
    for arg in arguments:
        name = arg['name']
        if arg.get('required', True):
            required_args.add(name)
        else:
            optional_args.add(name)
    
    return required_args, optional_args


def fix_signature(signature: str, missing_required: Set[str], missing_optional: Set[str]) -> str:
    """Fix the signature by adding missing arguments."""
    # Split signature into command and existing params
    parts = signature.split()
    command = parts[0]
    
    # Extract existing params to preserve order
    existing_params = ' '.join(parts[1:])
    
    # Add missing required params after existing required ones
    for param in sorted(missing_required):
        existing_params += f' <{param}>'
    
    # Add missing optional params at the end
    for param in sorted(missing_optional):
        existing_params += f' [<{param}>]'
    
    return f"{command} {existing_params}".strip()


def fix_config_file(config_path: Path, dry_run: bool = False) -> bool:
    """Fix a config file by adding missing arguments to the signature."""
    try:
        with open(config_path) as f:
            config = yaml.safe_load(f)
        
        # Get the first (and likely only) tool
        tool_name = next(iter(config['tools']))
        tool_config = config['tools'][tool_name]
        
        signature = tool_config['signature']
        arguments = tool_config['arguments']
        
        signature_params = extract_signature_params(signature)
        required_args, optional_args = extract_defined_arguments(arguments)
        defined_args = required_args | optional_args
        
        missing_from_signature = defined_args - signature_params
        missing_required = missing_from_signature & required_args
        missing_optional = missing_from_signature & optional_args
        
        print(f"Processing {config_path}")
        print(f"Tool: {tool_name}")
        print(f"Current signature: {signature}")
        
        if not missing_from_signature:
            print("✅ No missing arguments - signature is correct")
            return True
        
        print(f"Missing arguments: {sorted(missing_from_signature)}")
        
        if missing_required:
            print(f"  Required: {sorted(missing_required)}")
        if missing_optional:
            print(f"  Optional: {sorted(missing_optional)}")
        
        new_signature = fix_signature(signature, missing_required, missing_optional)
        print(f"Fixed signature: {new_signature}")
        
        if dry_run:
            print("🔍 DRY RUN - no changes made")
            return True
        
        # Update the config
        config['tools'][tool_name]['signature'] = new_signature
        
        # Write back to file
        with open(config_path, 'w') as f:
            yaml.dump(config, f, default_flow_style=False, sort_keys=False)
        
        print("✅ Config file fixed!")
        return True
        
    except Exception as e:
        print(f"❌ Error fixing {config_path}: {e}")
        return False


def main():
    parser = argparse.ArgumentParser(description="Fix YAML config signatures automatically")
    parser.add_argument("config_files", nargs="+", help="YAML config files to fix")
    parser.add_argument("--dry-run", action="store_true", 
                       help="Show what would be changed without making changes")
    args = parser.parse_args()
    
    all_successful = True
    for config_file in args.config_files:
        config_path = Path(config_file)
        if not config_path.exists():
            print(f"❌ File not found: {config_path}")
            all_successful = False
            continue
            
        success = fix_config_file(config_path, dry_run=args.dry_run)
        all_successful = all_successful and success
        print()
    
    sys.exit(0 if all_successful else 1)


if __name__ == "__main__":
    main()
