#!/usr/bin/env python3
"""
DOOM WAD Flasher for ESP32-S3
Flashes WAD files to the correct flash addresses
"""

import subprocess
import sys
import os
import argparse

# WAD addresses from partitions.csv
WAD_ADDRESSES = {
    'doom2.wad': 0x120000,
    'doom.wad': 0x120000,        # Alternative IWAD
    'prboom-plus.wad': 0xFA0000,
}

def flash_wad(port, wad_file, address=None):
    """Flash a WAD file to ESP32 flash"""
    filename = os.path.basename(wad_file).lower()

    if address is None:
        if filename in WAD_ADDRESSES:
            address = WAD_ADDRESSES[filename]
        else:
            print(f"Unknown WAD file: {filename}")
            print(f"Known WADs: {list(WAD_ADDRESSES.keys())}")
            return False

    if not os.path.exists(wad_file):
        print(f"WAD file not found: {wad_file}")
        return False

    file_size = os.path.getsize(wad_file)
    print(f"Flashing {wad_file} ({file_size} bytes) to address 0x{address:X}")

    cmd = [
        'esptool.py',
        '--chip', 'esp32s3',
        '--port', port,
        'write_flash',
        f'0x{address:X}',
        wad_file
    ]

    try:
        subprocess.run(cmd, check=True)
        print(f"Successfully flashed {wad_file}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"Error flashing {wad_file}: {e}")
        return False
    except FileNotFoundError:
        print("esptool.py not found. Install it with: pip install esptool")
        return False

def main():
    parser = argparse.ArgumentParser(description='Flash DOOM WAD files to ESP32-S3')
    parser.add_argument('--port', '-p', default='COM3', help='Serial port (default: COM3)')
    parser.add_argument('--iwad', '-i', help='IWAD file (doom2.wad or doom.wad)')
    parser.add_argument('--pwad', help='PWAD file (prboom-plus.wad)')
    parser.add_argument('--address', '-a', type=lambda x: int(x, 0),
                       help='Custom flash address (hex, e.g., 0x120000)')
    parser.add_argument('files', nargs='*', help='WAD files to flash')

    args = parser.parse_args()

    success = True

    if args.iwad:
        if not flash_wad(args.port, args.iwad, 0x120000):
            success = False

    if args.pwad:
        if not flash_wad(args.port, args.pwad, 0xFA0000):
            success = False

    for wad_file in args.files:
        if not flash_wad(args.port, wad_file, args.address):
            success = False

    if not args.iwad and not args.pwad and not args.files:
        print("Usage: python flash_wads.py [--port PORT] --iwad doom2.wad --pwad prboom-plus.wad")
        print("   or: python flash_wads.py [--port PORT] wad_file1 wad_file2 ...")

    return 0 if success else 1

if __name__ == '__main__':
    sys.exit(main())
