"""Command-line interface for model-exporter."""

import argparse
import sys
from pathlib import Path

from model_exporter.exporters import get_exporter, list_supported_models
from model_exporter.formats import list_formats


def export_command(args):
    """Handle export command."""
    try:
        # Get appropriate exporter
        exporter = get_exporter(
            args.model,
            cache_dir=args.cache_dir
        )
        
        # Load model
        exporter.load_model()
        
        # Print model info if requested
        if args.info:
            exporter.print_model_info()
            print()
        
        # Export to specified format
        print(f"Exporting to {args.format} format...")
        exporter.export(
            output_dir=args.output,
            format_name=args.format,
            save_vocab=not args.no_vocab
        )
        
        print(f"\n✓ Export complete!")
        
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


def list_command(args):
    """Handle list command."""
    if args.type == 'models':
        print("Supported model families:\n")
        families = list_supported_models()
        for family, models in families.items():
            print(f"{family}:")
            for model in models:
                print(f"  - {model}")
            print()
    
    elif args.type == 'formats':
        print("Supported export formats:\n")
        formats = list_formats()
        for fmt in formats:
            print(f"  - {fmt}")
        print()


def info_command(args):
    """Handle info command for inspecting weight files."""
    weights_file = Path(args.file)
    if not weights_file.exists():
        print(f"Error: File not found: {weights_file}", file=sys.stderr)
        sys.exit(1)
    
    print(f"Weights file: {weights_file}\n")
    
    # Detect format based on file extension
    if weights_file.suffix == '.bin':
        from model_exporter.formats.binary import BinaryFormatExporter
        BinaryFormatExporter.load_info(weights_file, max_tensors=args.max_tensors)
    elif weights_file.suffix == '.safetensors':
        from model_exporter.formats.safetensors import SafetensorsFormatExporter
        SafetensorsFormatExporter.load_info(weights_file, max_tensors=args.max_tensors)
    else:
        print(f"Error: Unsupported file format: {weights_file.suffix}", file=sys.stderr)
        print("Supported formats: .bin (binary), .safetensors", file=sys.stderr)
        sys.exit(1)


def main():
    """Main CLI entry point."""
    parser = argparse.ArgumentParser(
        prog='model-exporter',
        description='Universal model weight exporter for multiple formats'
    )
    
    subparsers = parser.add_subparsers(dest='command', help='Commands')
    
    # Export command
    export_parser = subparsers.add_parser(
        'export',
        help='Export model weights to specified format'
    )
    export_parser.add_argument(
        'model',
        type=str,
        help='HuggingFace model name (e.g., gpt2, gpt2-medium)'
    )
    export_parser.add_argument(
        '--format',
        type=str,
        default='binary',
        choices=list_formats(),
        help='Export format (default: binary)'
    )
    export_parser.add_argument(
        '--output',
        type=str,
        default='./weights',
        help='Output directory (default: ./weights)'
    )
    export_parser.add_argument(
        '--cache-dir',
        type=str,
        default=None,
        help='Cache directory for downloaded models'
    )
    export_parser.add_argument(
        '--no-vocab',
        action='store_true',
        help='Skip vocabulary export'
    )
    export_parser.add_argument(
        '--info',
        action='store_true',
        help='Print model information before export'
    )
    export_parser.set_defaults(func=export_command)
    
    # List command
    list_parser = subparsers.add_parser(
        'list',
        help='List supported models or formats'
    )
    list_parser.add_argument(
        'type',
        choices=['models', 'formats'],
        help='What to list'
    )
    list_parser.set_defaults(func=list_command)
    
    # Info command
    info_parser = subparsers.add_parser(
        'info',
        help='Inspect binary weights file'
    )
    info_parser.add_argument(
        'file',
        type=str,
        help='Path to binary weights file'
    )
    info_parser.add_argument(
        '--max-tensors',
        type=int,
        default=5,
        help='Maximum number of tensors to display (default: 5)'
    )
    info_parser.set_defaults(func=info_command)
    
    # Parse and execute
    args = parser.parse_args()
    
    if not args.command:
        parser.print_help()
        sys.exit(1)
    
    args.func(args)


if __name__ == '__main__':
    main()