#!/usr/bin/env python3
"""
MIDI Test Validator
Compares MIDI events from USB MIDI device with firmware event counters
"""

import serial
import time
import rtmidi
import random

# Serial port configuration - adjust COM port as needed
FIRMWARE_PORT = "COM9"  # Firmware USB serial port
BAUD_RATE = 115200

# MIDI device name to look for
MIDI_DEVICE_NAME = "Midilink"  # Will search for devices containing this string

# MIDI Octave to filter (0-10, or None for all octaves)
MIDI_OCTAVE = 1  # Only count notes in octave 4 (middle C is C4)

# MIDI Channel (0-15, where 0 = channel 1)
MIDI_CHANNEL = 0

# Test data generation settings
TEST_NUM_EVENTS = 200000    # Total number of Note ON events to generate
TEST_MIN_DELAY_MS = 5       # Minimum delay between events (ms)
TEST_MAX_DELAY_MS = 30      # Maximum delay between events (ms)
TEST_MAX_OVERLAPPING = 300  # Maximum number of overlapping notes

# Debug output
DEBUG_ENABLED = False  # Set to False to disable debug output

# MIDI definitions
MIDI_NOTE_ON = 0x90
MIDI_NOTE_OFF = 0x80

class MidiEventCounter:
    def __init__(self):
        self.red_on = 0
        self.red_off = 0
        self.green_on = 0
        self.green_off = 0
        self.blue_on = 0
        self.blue_off = 0
        
    def add_event(self, note, velocity, is_note_on, debug=False):
        # Calculate octave (MIDI note 0-127, octave = note // 12, where C4 = MIDI note 60)
        # MIDI octaves: C-1=0-11, C0=12-23, C1=24-35, C2=36-47, C3=48-59, C4=60-71, etc.
        octave = note // 12 - 1  # Adjust so C4 = octave 4
        note_in_octave = note % 12
        
        if debug:
            note_names = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
            note_name = f"{note_names[note_in_octave]}{octave}"
            print(f"[DEBUG] Note {note} = {note_name} (Octave: {octave}, Note in octave: {note_in_octave})")
        
        # Check if octave matches (if MIDI_OCTAVE is set)
        if MIDI_OCTAVE is not None and octave != MIDI_OCTAVE:
            if debug:
                print(f"[DEBUG]   -> Ignoring: Wrong octave (expected {MIDI_OCTAVE}, got {octave})")
            return
        
        # Determine color based on note (adjust to your config)
        # Assuming C=Red, D=Green, E=Blue
        color_mapped = None
        
        if note_in_octave in [0, 1]:  # C, C# = Red
            if is_note_on:
                self.red_on += 1
            else:
                self.red_off += 1
            color_mapped = "RED"
        elif note_in_octave in [2, 3]:  # D, D# = Green
            if is_note_on:
                self.green_on += 1
            else:
                self.green_off += 1
            color_mapped = "GREEN"
        elif note_in_octave in [4, 5]:  # E, F = Blue
            if is_note_on:
                self.blue_on += 1
            else:
                self.blue_off += 1
            color_mapped = "BLUE"
        else:
            color_mapped = "NONE (ignored)"
        
        if debug:
            event_type = "ON " if is_note_on else "OFF"
            print(f"[DEBUG] Mapped Note {note} -> {color_mapped} {event_type}")
    
    def print_summary(self):
        print("\nMIDI Device Event Counts:")
        print(f"  Red   - ON: {self.red_on:6d}, OFF: {self.red_off:6d}")
        print(f"  Green - ON: {self.green_on:6d}, OFF: {self.green_off:6d}")
        print(f"  Blue  - ON: {self.blue_on:6d}, OFF: {self.blue_off:6d}")

def midi_callback(message, data):
    """Callback for MIDI input - called when MIDI messages are received"""
    event_counter = data
    midi_message = message[0]
    
    if len(midi_message) >= 3:
        status = midi_message[0]
        note = midi_message[1]
        velocity = midi_message[2]
        
        # Check for Note On or Note Off
        if (status & 0xF0) == MIDI_NOTE_ON or (status & 0xF0) == MIDI_NOTE_OFF:
            channel = status & 0x0F
            is_note_on = ((status & 0xF0) == MIDI_NOTE_ON and velocity > 0)
            event_type = "Note ON" if is_note_on else "Note OFF"
            
            if DEBUG_ENABLED:
                print(f"[MIDI] {event_type}, Channel: {channel}, Note: {note}, Velocity: {velocity}")
            
            event_counter.add_event(note, velocity, is_note_on, debug=DEBUG_ENABLED)

def find_midi_device(device_name):
    """Find MIDI input device by name"""
    midiin = rtmidi.MidiIn()
    ports = midiin.get_ports()
    
    print(f"\nAvailable MIDI Input devices:")
    for i, port in enumerate(ports):
        print(f"  [{i}] {port}")
    
    # Search for device containing the name
    for i, port in enumerate(ports):
        if device_name.lower() in port.lower():
            print(f"\n✓ Found '{device_name}' input at port {i}: {port}")
            return i
    
    print(f"\n✗ Could not find MIDI input device containing '{device_name}'")
    return None

def find_midi_output_device(device_name):
    """Find MIDI output device by name"""
    midiout = rtmidi.MidiOut()
    ports = midiout.get_ports()
    
    print(f"\nAvailable MIDI Output devices:")
    for i, port in enumerate(ports):
        print(f"  [{i}] {port}")
    
    # Search for device containing the name
    for i, port in enumerate(ports):
        if device_name.lower() in port.lower():
            print(f"\n✓ Found '{device_name}' output at port {i}: {port}")
            return i
    
    print(f"\n✗ Could not find MIDI output device containing '{device_name}'")
    return None

def generate_test_data(midiout):
    """Generate random MIDI test data with overlapping notes"""
    # Notes to use: C, C#, D, D#, E, F (Red, Green, Blue colors with alternates)
    # Calculate MIDI note numbers for the configured octave
    base_note = (MIDI_OCTAVE + 1) * 12  # Base note for the octave
    
    notes = {
        'C': base_note + 0,   # Red
        'C#': base_note + 1,  # Red alternate
        'D': base_note + 2,   # Green
        'D#': base_note + 3,  # Green alternate
        'E': base_note + 4,   # Blue
        'F': base_note + 5    # Blue alternate
    }
    
    note_names = list(notes.keys())
    note_values = list(notes.values())
    
    print(f"\nGenerating {TEST_NUM_EVENTS} random MIDI events...")
    print(f"Octave: {MIDI_OCTAVE}, Channel: {MIDI_CHANNEL + 1}")
    print(f"Using notes: {', '.join([f'{name}({notes[name]})' for name in note_names])}")
    print("="*50 + "\n")
    
    active_notes = []  # Track currently active (on but not yet off) notes
    event_queue = []   # Queue of all events to send
    
    # Generate NOTE ON events
    for i in range(TEST_NUM_EVENTS):
        note_idx = random.randint(0, len(note_values) - 1)
        note = note_values[note_idx]
        velocity = random.randint(64, 127)  # Random velocity
        
        event_queue.append({
            'type': 'on',
            'note': note,
            'note_name': note_names[note_idx],
            'velocity': velocity
        })
        active_notes.append(note)
        
        # Randomly send some NOTE OFF events if we have overlapping notes
        while len(active_notes) >= TEST_MAX_OVERLAPPING:
            off_note = active_notes.pop(0)
            off_note_name = note_names[note_values.index(off_note)]
            event_queue.append({
                'type': 'off',
                'note': off_note,
                'note_name': off_note_name,
                'velocity': 0
            })
        
        # Random chance to send an OFF event even if under max
        if len(active_notes) > 0 and random.random() < 0.4:
            off_note = active_notes.pop(random.randint(0, len(active_notes) - 1))
            off_note_name = note_names[note_values.index(off_note)]
            event_queue.append({
                'type': 'off',
                'note': off_note,
                'note_name': off_note_name,
                'velocity': 0
            })
    
    # Send remaining NOTE OFF events
    while len(active_notes) > 0:
        off_note = active_notes.pop(0)
        off_note_name = note_names[note_values.index(off_note)]
        event_queue.append({
            'type': 'off',
            'note': off_note,
            'note_name': off_note_name,
            'velocity': 0
        })
    
    # Send all events with random delays
    total_events = len(event_queue)
    for idx, event in enumerate(event_queue):
        if event['type'] == 'on':
            message = [MIDI_NOTE_ON | MIDI_CHANNEL, event['note'], event['velocity']]
            if DEBUG_ENABLED:
                print(f"→ Note ON:  {event['note_name']:3s} (MIDI {event['note']:3d}), Velocity: {event['velocity']:3d}")
        else:
            message = [MIDI_NOTE_OFF | MIDI_CHANNEL, event['note'], event['velocity']]
            if DEBUG_ENABLED:
                print(f"← Note OFF: {event['note_name']:3s} (MIDI {event['note']:3d})")
        
        midiout.send_message(message)
        
        # Show progress bar (only if debug is disabled to avoid cluttering output)
        if not DEBUG_ENABLED:
            progress = (idx + 1) / total_events
            bar_length = 40
            filled_length = int(bar_length * progress)
            bar = '█' * filled_length + '░' * (bar_length - filled_length)
            print(f'\rProgress: [{bar}] {idx + 1}/{total_events} events ({progress*100:.1f}%)', end='', flush=True)
        
        # Random delay before next event
        delay_ms = random.randint(TEST_MIN_DELAY_MS, TEST_MAX_DELAY_MS)
        time.sleep(delay_ms / 1000.0)
    
    if not DEBUG_ENABLED:
        print()  # New line after progress bar
    
    print("\n" + "="*50)
    print("✓ Test data generation complete")
    print("="*50)

def read_firmware_counters(ser_fw):
    """Send 'a' command and read firmware event counters"""
    # Clear input buffer
    ser_fw.reset_input_buffer()
    
    # Send command
    ser_fw.write(b'a\r')
    time.sleep(0.2)
    
    # Read all available data
    response = ser_fw.read(ser_fw.in_waiting).decode('ascii')
    
    # Split by \r to get individual lines
    lines = response.split('\r')
    
    counters = {}
    colors = ['red', 'green', 'blue']
    
    line_idx = 0
    for color in colors:
        if line_idx < len(lines):
            line = lines[line_idx].strip()
            if ';' in line:
                parts = line.split(';')
                on_count = int(parts[0].strip())
                off_count = int(parts[1].strip())
                counters[color] = {'on': on_count, 'off': off_count}
            else:
                print(f"Warning: Could not parse line for {color}: {line}")
                counters[color] = {'on': 0, 'off': 0}
            line_idx += 1
        else:
            print(f"Warning: No data for {color}")
            counters[color] = {'on': 0, 'off': 0}
    
    return counters

def compare_results(midi_counts, firmware_counts):
    """Compare MIDI device and firmware event counts"""
    print("\nFirmware Event Counts:")
    print(f"  Red   - ON: {firmware_counts['red']['on']:6d}, OFF: {firmware_counts['red']['off']:6d}")
    print(f"  Green - ON: {firmware_counts['green']['on']:6d}, OFF: {firmware_counts['green']['off']:6d}")
    print(f"  Blue  - ON: {firmware_counts['blue']['on']:6d}, OFF: {firmware_counts['blue']['off']:6d}")
    
    print("\n" + "="*50)
    print("COMPARISON RESULTS:")
    print("="*50)
    
    all_match = True
    
    # Compare Red
    if midi_counts.red_on == firmware_counts['red']['on'] and \
       midi_counts.red_off == firmware_counts['red']['off']:
        print("Red:   PASS ✓")
    else:
        print(f"Red:   FAIL ✗ (MIDI: {midi_counts.red_on}/{midi_counts.red_off}, "
              f"Firmware: {firmware_counts['red']['on']}/{firmware_counts['red']['off']})")
        all_match = False
    
    # Compare Green
    if midi_counts.green_on == firmware_counts['green']['on'] and \
       midi_counts.green_off == firmware_counts['green']['off']:
        print("Green: PASS ✓")
    else:
        print(f"Green: FAIL ✗ (MIDI: {midi_counts.green_on}/{midi_counts.green_off}, "
              f"Firmware: {firmware_counts['green']['on']}/{firmware_counts['green']['off']})")
        all_match = False
    
    # Compare Blue
    if midi_counts.blue_on == firmware_counts['blue']['on'] and \
       midi_counts.blue_off == firmware_counts['blue']['off']:
        print("Blue:  PASS ✓")
    else:
        print(f"Blue:  FAIL ✗ (MIDI: {midi_counts.blue_on}/{midi_counts.blue_off}, "
              f"Firmware: {firmware_counts['blue']['on']}/{firmware_counts['blue']['off']})")
        all_match = False
    
    print("="*50)
    
    if all_match:
        print("✓ ALL TESTS PASSED")
    else:
        print("✗ SOME TESTS FAILED")
    
    return all_match

def main():
    print("MIDI Test Validator")
    print("="*50)
    print("\nSelect mode:")
    print("  1. Manual mode (record MIDI events manually)")
    print("  2. Auto-test mode (send random test data)")
    
    mode = input("\nEnter mode (1 or 2): ").strip()
    
    auto_test = (mode == '2')
    
    # Find and open MIDI input device
    midi_in_port = find_midi_device(MIDI_DEVICE_NAME)
    if midi_in_port is None:
        print("\nPlease check:")
        print("  1. MIDI device is connected")
        print("  2. MIDI_DEVICE_NAME variable matches your device")
        return
    
    midiin = rtmidi.MidiIn()
    event_counter = MidiEventCounter()
    
    # Set callback and open input port
    midiin.set_callback(midi_callback, data=event_counter)
    midiin.open_port(midi_in_port)
    print(f"✓ Opened MIDI input port")
    
    # Find and open MIDI output device (for auto-test mode)
    midiout = None
    if auto_test:
        midi_out_port = find_midi_output_device(MIDI_DEVICE_NAME)
        if midi_out_port is None:
            print("\nCould not find MIDI output device for auto-test mode")
            midiin.close_port()
            return
        
        midiout = rtmidi.MidiOut()
        midiout.open_port(midi_out_port)
        print(f"✓ Opened MIDI output port")
    
    # Open firmware serial port
    try:
        ser_firmware = serial.Serial(FIRMWARE_PORT, BAUD_RATE, timeout=1)
        print(f"✓ Connected to firmware on {FIRMWARE_PORT}")
    except serial.SerialException as e:
        print(f"Error opening serial port: {e}")
        midiin.close_port()
        if midiout:
            midiout.close_port()
        return
    
    # Clear firmware counters at startup
    print("\nClearing firmware counters...")
    read_firmware_counters(ser_firmware)  # This reads and resets the counters
    print("✓ Firmware counters cleared")
    
    if MIDI_OCTAVE is not None:
        print(f"\nConfiguration: Filtering for MIDI Octave {MIDI_OCTAVE} only")
    else:
        print(f"\nConfiguration: Accepting all MIDI octaves")
    
    print(f"Debug output: {'ENABLED' if DEBUG_ENABLED else 'DISABLED'}")
    
    if auto_test:
        # Auto-test mode: send test data
        print("\n" + "="*50)
        print("AUTO-TEST MODE")
        print("="*50)
        time.sleep(0.5)  # Small delay to ensure everything is ready
        
        generate_test_data(midiout)
        
        # Give some time for last events to be processed
        time.sleep(0.5)
    else:
        # Manual mode: wait for user input
        print("\n" + "="*50)
        print("MANUAL MODE - Recording MIDI events...")
        print("Press ENTER to stop recording and compare results")
        print("="*50 + "\n")
        
        try:
            input()  # Wait for user to press enter
        except KeyboardInterrupt:
            print("\n\nStopped by user (Ctrl+C)")
    
    # Display results
    event_counter.print_summary()
    
    # Read firmware counters
    print("\nReading firmware counters...")
    firmware_counts = read_firmware_counters(ser_firmware)
    
    # Compare results
    compare_results(event_counter, firmware_counts)
    
    # Close ports
    midiin.close_port()
    if midiout:
        midiout.close_port()
    ser_firmware.close()
    print("\nPorts closed.")

if __name__ == "__main__":
    main()