#!/usr/bin/env python3
"""
Example: Create an MPQ archive from a directory
"""

from pystorm import MPQArchive, MPQ_CREATE_ARCHIVE_V2, MPQ_FILE_COMPRESS, MPQ_COMPRESSION_ZLIB
from pathlib import Path
import sys


def create_archive_from_directory(source_dir, archive_path, compress=True):
    """
    Create an MPQ archive from all files in a directory
    
    Args:
        source_dir: Directory containing files to archive
        archive_path: Path for the output MPQ archive
        compress: Whether to compress files (default: True)
    """
    source_dir = Path(source_dir)
    archive_path = Path(archive_path)
    
    if not source_dir.exists():
        print(f"Error: Directory not found: {source_dir}")
        return False
    
    if not source_dir.is_dir():
        print(f"Error: Not a directory: {source_dir}")
        return False
    
    print(f"Creating archive from: {source_dir}")
    print(f"Output archive: {archive_path}")
    print(f"Compression: {'enabled' if compress else 'disabled'}\n")
    
    # Collect all files
    files_to_add = []
    for file_path in source_dir.rglob("*"):
        if file_path.is_file():
            files_to_add.append(file_path)
    
    if not files_to_add:
        print("Error: No files found in directory")
        return False
    
    print(f"Found {len(files_to_add)} file(s) to add\n")
    
    try:
        # Create the archive
        with MPQArchive(str(archive_path), flags=MPQ_CREATE_ARCHIVE_V2) as archive:
            added = 0
            failed = 0
            
            for i, file_path in enumerate(files_to_add, 1):
                # Get relative path for archived name
                archived_name = str(file_path.relative_to(source_dir))
                # Replace backslashes with forward slashes (MPQ convention)
                archived_name = archived_name.replace("\\", "/")
                
                # Determine flags
                flags = MPQ_FILE_COMPRESS if compress else 0
                compression = MPQ_COMPRESSION_ZLIB if compress else 0
                
                try:
                    archive.add_file(str(file_path), archived_name, flags, compression)
                    added += 1
                    
                    # Get file size for display
                    file_size = file_path.stat().st_size
                    size_kb = file_size / 1024
                    print(f"[{i}/{len(files_to_add)}] ✓ {archived_name} ({size_kb:.2f} KB)")
                    
                except Exception as e:
                    failed += 1
                    print(f"[{i}/{len(files_to_add)}] ✗ {archived_name} - Error: {e}")
            
            # Flush changes
            print("\nFlushing changes to disk...")
            archive.flush()
            
            print(f"\n{'='*60}")
            print(f"Archive creation complete!")
            print(f"  Files added: {added}")
            print(f"  Failed: {failed}")
            print(f"  Output: {archive_path}")
            
            # Get archive size
            if archive_path.exists():
                archive_size = archive_path.stat().st_size
                size_mb = archive_size / (1024 * 1024)
                print(f"  Archive size: {size_mb:.2f} MB")
            
            print(f"{'='*60}")
            
            return failed == 0
            
    except Exception as e:
        print(f"Error creating archive: {e}")
        return False


def main():
    """Main entry point"""
    if len(sys.argv) < 2:
        print("Usage: python create_archive.py <source_dir> [output.mpq] [--no-compress]")
        print("\nExample:")
        print("  python create_archive.py my_files output.mpq")
        print("  python create_archive.py my_files output.mpq --no-compress")
        sys.exit(1)
    
    source_dir = sys.argv[1]
    archive_path = sys.argv[2] if len(sys.argv) > 2 and not sys.argv[2].startswith('--') else "output.mpq"
    compress = "--no-compress" not in sys.argv
    
    success = create_archive_from_directory(source_dir, archive_path, compress)
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()