import mido
import sys
import os
import argparse
import json

# General MIDI Program Numbers mapped to Instrument Names
GENERAL_MIDI_PROGRAMS = [
    "Acoustic Grand Piano", "Bright Acoustic Piano", "Electric Grand Piano", "Honky-tonk Piano",
    "Electric Piano 1", "Electric Piano 2", "Harpsichord", "Clavinet",
    "Celesta", "Glockenspiel", "Music Box", "Vibraphone",
    "Marimba", "Xylophone", "Tubular Bells", "Dulcimer",
    "Drawbar Organ", "Percussive Organ", "Rock Organ", "Church Organ",
    "Reed Organ", "Accordion", "Harmonica", "Tango Accordion",
    "Acoustic Guitar (nylon)", "Acoustic Guitar (steel)", "Electric Guitar (jazz)", "Electric Guitar (clean)",
    "Electric Guitar (muted)", "Overdriven Guitar", "Distortion Guitar", "Guitar harmonics",
    "Acoustic Bass", "Electric Bass (finger)", "Electric Bass (pick)", "Fretless Bass",
    "Slap Bass 1", "Slap Bass 2", "Synth Bass 1", "Synth Bass 2",
    "Violin", "Viola", "Cello", "Contrabass",
    "Tremolo Strings", "Pizzicato Strings", "Orchestral Harp", "Timpani",
    "String Ensemble 1", "String Ensemble 2", "SynthStrings 1", "SynthStrings 2",
    "Choir Aahs", "Voice Oohs", "Synth Choir", "Orchestra Hit",
    "Trumpet", "Trombone", "Tuba", "Muted Trumpet",
    "French Horn", "Brass Section", "SynthBrass 1", "SynthBrass 2",
    "Soprano Sax", "Alto Sax", "Tenor Sax", "Baritone Sax",
    "Oboe", "English Horn", "Bassoon", "Clarinet",
    "Piccolo", "Flute", "Recorder", "Pan Flute",
    "Blown Bottle", "Shakuhachi", "Whistle", "Ocarina",
    "Lead 1 (square)", "Lead 2 (sawtooth)", "Lead 3 (calliope)", "Lead 4 (chiff)",
    "Lead 5 (charang)", "Lead 6 (voice)", "Lead 7 (fifths)", "Lead 8 (bass + lead)",
    "Pad 1 (new age)", "Pad 2 (warm)", "Pad 3 (polysynth)", "Pad 4 (choir)",
    "Pad 5 (bowed)", "Pad 6 (metallic)", "Pad 7 (halo)", "Pad 8 (sweep)",
    "FX 1 (rain)", "FX 2 (soundtrack)", "FX 3 (crystal)", "FX 4 (atmosphere)",
    "FX 5 (brightness)", "FX 6 (goblins)", "FX 7 (echoes)", "FX 8 (sci-fi)",
    "Sitar", "Banjo", "Shamisen", "Koto",
    "Kalimba", "Bag pipe", "Fiddle", "Shanai",
    "Tinkle Bell", "Agogo", "Steel Drums", "Woodblock",
    "Taiko Drum", "Melodic Tom", "Synth Drum", "Reverse Cymbal",
    "Guitar Fret Noise", "Breath Noise", "Seashore", "Bird Tweet",
    "Telephone Ring", "Helicopter", "Applause", "Gunshot"
]

def get_instrument_name(program_number):
    """
    Maps a MIDI program number to its instrument name.
    MIDI program numbers are 0-based.
    """
    if 0 <= program_number < len(GENERAL_MIDI_PROGRAMS):
        return GENERAL_MIDI_PROGRAMS[program_number]
    else:
        return f"Unknown Program ({program_number})"

def has_musical_messages(track):
    """
    Determines if a MIDI track contains any musical messages.
    Returns True if the track has at least one musical message, False otherwise.
    """
    musical_types = {
        'note_on', 'note_off', 'program_change', 'control_change',
        'pitchwheel', 'aftertouch', 'polyphonic_key_pressure'
    }
    for msg in track:
        if msg.type in musical_types:
            return True
    return False

def collect_tempo_changes(midi):
    """
    Collects all tempo changes from all tracks in the MIDI file.
    Returns a list of tuples: (absolute_tick_time, bpm)
    """
    tempo_changes = []
    ticks_per_beat = midi.ticks_per_beat
    DEFAULT_TEMPO = 500000  # Default tempo (microseconds per beat) is 120 BPM

    for track in midi.tracks:
        absolute_time = 0
        for msg in track:
            absolute_time += msg.time
            if msg.type == 'set_tempo':
                bpm = mido.tempo2bpm(msg.tempo)
                tempo_changes.append((absolute_time, bpm))

    # If no tempo changes found, assume default tempo
    if not tempo_changes:
        tempo_changes.append((0, mido.tempo2bpm(DEFAULT_TEMPO)))

    # Sort tempo changes by tick time
    tempo_changes.sort(key=lambda x: x[0])

    return tempo_changes

def analyze_midi(file_path):
    try:
        midi = mido.MidiFile(file_path)
    except IOError:
        print(f"Error: Could not open MIDI file: {file_path}")
        return False
    except mido.KeySignatureError:
        print(f"Error: Invalid MIDI file: {file_path}")
        return False

    analysis = {
        "song_title": os.path.splitext(os.path.basename(file_path))[0],
        "tempo": {
            "min_bpm": None,
            "max_bpm": None
        },
        "pitch_bend": {
            "min_semitones": None,
            "max_semitones": None
        },
        "tracks": [],
        "notes": "",
        "song_offset": 0
    }

    # Collect tempo changes
    tempo_changes = collect_tempo_changes(midi)
    tempos = [bpm for (_, bpm) in tempo_changes]
    analysis["tempo"]["min_bpm"] = min(tempos)
    analysis["tempo"]["max_bpm"] = max(tempos)

    musical_track_count = 0  # To number only musical tracks
    # Initialize per channel RPN state and pitch bend range
    channel_rpn_state = {channel: {'selected_rpn_msb': None, 'selected_rpn_lsb': None, 'rpn_selected': None} for channel in range(16)}
    channel_pitch_bend_range = {channel: 2 for channel in range(16)}  # Default to ±2 semitones

    global_pitch_bends = []

    for track in midi.tracks:
        if not has_musical_messages(track):
            continue  # Skip non-musical tracks

        musical_track_count += 1
        track_info = {
            "track_name": f"Track {musical_track_count}",
            "Channel Assignment": [],
            "Pitch Bend Sensitivity": {},
            "Max Note Velocity": "N/A",
            "Min Note Velocity": "N/A",
            "Uses Program Change": False
        }

        program_changes = []
        pitch_bends_semitones = []

        absolute_time = 0  # To keep track of the absolute time in ticks

        for msg in track:
            absolute_time += msg.time  # Accumulate delta times to get absolute time
            if msg.type == 'track_name':
                track_info["track_name"] = msg.name
            elif hasattr(msg, 'channel'):
                channel = msg.channel  # 0-15
                if (channel + 1) not in track_info["Channel Assignment"]:
                    track_info["Channel Assignment"].append(channel + 1)  # Channels represented as 1-16

                if msg.type == 'note_on' and msg.velocity > 0:
                    velocity = msg.velocity
                    if track_info["Max Note Velocity"] == "N/A" or velocity > track_info["Max Note Velocity"]:
                        track_info["Max Note Velocity"] = velocity
                    if track_info["Min Note Velocity"] == "N/A" or velocity < track_info["Min Note Velocity"]:
                        track_info["Min Note Velocity"] = velocity
                elif msg.type == 'program_change':
                    track_info["Uses Program Change"] = True
                    program_changes.append({
                        "program_number": msg.program,
                        "instrument_name": get_instrument_name(msg.program),
                        "tick": absolute_time
                    })
                elif msg.type == 'control_change':
                    # Handle RPN messages for Pitch Bend Sensitivity
                    if msg.control == 101:
                        # Set RPN MSB
                        channel_rpn_state[channel]['selected_rpn_msb'] = msg.value
                        if msg.value != 0:
                            channel_rpn_state[channel]['rpn_selected'] = None
                    elif msg.control == 100:
                        # Set RPN LSB
                        channel_rpn_state[channel]['selected_rpn_lsb'] = msg.value
                        # Check if RPN 0 (Pitch Bend Sensitivity) is selected
                        if (channel_rpn_state[channel]['selected_rpn_msb'] == 0 and
                            channel_rpn_state[channel]['selected_rpn_lsb'] == 0):
                            channel_rpn_state[channel]['rpn_selected'] = 'pitch_bend_range'
                        else:
                            channel_rpn_state[channel]['rpn_selected'] = None
                    elif msg.control == 6:
                        # Data Entry MSB
                        if channel_rpn_state[channel].get('rpn_selected') == 'pitch_bend_range':
                            # Set pitch bend range in semitones (integer part)
                            channel_pitch_bend_range[channel] = msg.value
                    elif msg.control == 38:
                        # Data Entry LSB
                        # Currently not handling fractional semitones
                        pass
                elif msg.type == 'pitchwheel':
                    # Calculate semitone shift based on current pitch bend range
                    current_range = channel_pitch_bend_range[channel]
                    semitones = (msg.pitch / 8192) * current_range
                    pitch_bends_semitones.append(semitones)
                    global_pitch_bends.append(semitones)

        # Update Pitch Bend Sensitivity for channels used in this track
        for channel in track_info["Channel Assignment"]:
            sensitivity = channel_pitch_bend_range[channel - 1]
            track_info["Pitch Bend Sensitivity"][f"Channel {channel}"] = round(sensitivity, 2)

        # Analyze pitch bends
        if pitch_bends_semitones:
            track_info["pitch_bend"] = {
                "min_semitones": round(min(pitch_bends_semitones), 2),
                "max_semitones": round(max(pitch_bends_semitones), 2)
            }
        else:
            track_info["pitch_bend"] = {
                "min_semitones": 0.0,
                "max_semitones": 0.0
            }

        # Add program changes if any
        if program_changes:
            track_info["Program Changes"] = program_changes

        analysis["tracks"].append(track_info)

    # Analyze global pitch bends
    if global_pitch_bends:
        analysis["pitch_bend"]["min_semitones"] = round(min(global_pitch_bends), 2)
        analysis["pitch_bend"]["max_semitones"] = round(max(global_pitch_bends), 2)
    else:
        analysis["pitch_bend"]["min_semitones"] = 0.0
        analysis["pitch_bend"]["max_semitones"] = 0.0

    if musical_track_count == 0:
        print(f"Warning: No musical tracks found in MIDI file: {file_path}")
        return False

    # Write the analysis to a JSON file
    output_path = os.path.splitext(file_path)[0] + '.json'
    try:
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(analysis, f, indent=4, ensure_ascii=False)
        print(f"Analysis complete. Results saved to '{output_path}'")
        return True
    except IOError:
        print(f"Error: Could not write to output file: {output_path}")
        return False

def main():
    parser = argparse.ArgumentParser(
        description="Analyze a single MIDI file and generate a JSON report."
    )
    parser.add_argument(
        'input_file',
        help="Path to the MIDI file to analyze."
    )

    args = parser.parse_args()

    input_file = args.input_file

    if not os.path.isfile(input_file):
        print(f"Error: The provided path is not a file or does not exist: {input_file}")
        sys.exit(1)

    if not input_file.lower().endswith(('.mid', '.midi')):
        print(f"Error: The provided file is not a MIDI file: {input_file}")
        sys.exit(1)

    success = analyze_midi(input_file)
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()