import os
import sys
import json
import random
import math
import struct
import wave
import tempfile
import hashlib
from collections import defaultdict
from dataclasses import dataclass, field
from typing import List, Dict, Tuple, Optional
from enum import Enum

# =============================================================================
# MUSIC THEORY CONSTANTS
# =============================================================================

NOTES = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
NOTE_TO_MIDI = {note + str(octave): i + 12 * (octave - 1) 
                for octave in range(0, 9) for i, note in enumerate(NOTES)}
MIDI_TO_NOTE = {v: k for k, v in NOTE_TO_MIDI.items()}

MAJOR_SCALE = [0, 2, 4, 5, 7, 9, 11]
MINOR_SCALE = [0, 2, 3, 5, 7, 8, 10]

CHORD_TYPES = {
    'major': [0, 4, 7],
    'minor': [0, 3, 7],
    'diminished': [0, 3, 6],
    'augmented': [0, 4, 8],
    'major7': [0, 4, 7, 11],
    'minor7': [0, 3, 7, 10],
    'dominant7': [0, 4, 7, 10],
}

# =============================================================================
# DATA CLASSES
# =============================================================================

@dataclass
class Note:
    """Represents a musical note"""
    pitch: int  # MIDI note number
    duration: float  # in beats
    velocity: int = 64  # 0-127
    start_time: float = 0.0
    
    def get_frequency(self) -> float:
        """Convert MIDI note to frequency in Hz"""
        return 440.0 * (2 ** ((self.pitch - 69) / 12))
    
    def get_note_name(self) -> str:
        """Get note name (e.g., C4)"""
        return MIDI_TO_NOTE.get(self.pitch, f"Note{self.pitch}")

@dataclass
class Chord:
    """Represents a chord"""
    root: int  # MIDI note number of root
    chord_type: str  # e.g., 'major', 'minor'
    duration: float
    inversion: int = 0  # 0 = root position
    
    def get_notes(self) -> List[int]:
        """Get all notes in the chord"""
        intervals = CHORD_TYPES[self.chord_type]
        notes = [(self.root + interval) % 12 for interval in intervals]
        return notes

@dataclass
class Track:
    """Represents a musical track with notes"""
    notes: List[Note] = field(default_factory=list)
    instrument: str = "piano"
    tempo: float = 120.0  # BPM
    
    def add_note(self, note: Note):
        self.notes.append(note)
        self.notes.sort(key=lambda n: n.start_time)
    
    def get_duration(self) -> float:
        if not self.notes:
            return 0
        return max(n.start_time + n.duration for n in self.notes)

@dataclass
class Composition:
    """Complete musical composition"""
    tracks: Dict[str, Track] = field(default_factory=dict)
    title: str = "Untitled"
    tempo: float = 120.0
    
    def add_track(self, name: str, track: Track):
        self.tracks[name] = track
    
    def total_duration(self) -> float:
        return max(t.get_duration() for t in self.tracks.values()) if self.tracks else 0

# =============================================================================
# MARKOV CHAIN MELODY GENERATOR
# =============================================================================

class MarkovMelodyGenerator:
    """Generates melodies using Markov chains"""
    
    def __init__(self, scale: List[int] = MAJOR_SCALE, 
                 order: int = 2):
        self.scale = scale
        self.order = order
        self.transitions = defaultdict(lambda: defaultdict(float))
        
    def train(self, sequences: List[List[int]]):
        """Train on sequences of notes (MIDI numbers)"""
        for seq in sequences:
            for i in range(len(seq) - self.order):
                context = tuple(seq[i:i + self.order])
                next_note = seq[i + self.order]
                self.transitions[context][next_note] += 1
        
        # Normalize to probabilities
        for context in self.transitions:
            total = sum(self.transitions[context].values())
            for note in self.transitions[context]:
                self.transitions[context][note] /= total
    
    def generate(self, length: int = 16, 
                 seed: Optional[List[int]] = None) -> List[int]:
        """Generate a new melody sequence"""
        if seed is None:
            seed = [random.choice(list(self.transitions.keys())[0][0])]
        
        sequence = list(seed)
        
        for _ in range(length - len(seed)):
            context = tuple(sequence[-self.order:])
            
            if context in self.transitions:
                # Weighted random choice
                notes = list(self.transitions[context].keys())
                probs = list(self.transitions[context].values())
                next_note = random.choices(notes, weights=probs, k=1)[0]
            else:
                # Fallback: random note from scale
                next_note = random.choice(self.scale)
            
            sequence.append(next_note)
        
        return sequence

# =============================================================================
# CHORD PROGRESSION GENERATOR
# =============================================================================

class ChordProgressionGenerator:
    """Generates chord progressions based on music theory rules"""
    
    PROGRESSIONS = {
        'pop': ['I', 'V', 'vi', 'IV'],
        'jazz': ['ii', 'V', 'I', 'vi'],
        'blues': ['I', 'IV', 'I', 'V', 'IV', 'I'],
        'emotional': ['vi', 'IV', 'I', 'V'],
        'epic': ['i', 'VI', 'III', 'VII'],
        'ambient': ['I', 'vi', 'IV', 'V'],
    }
    
    def __init__(self, key: str = 'C', mode: str = 'major'):
        self.key = key
        self.mode = mode
        self.key_index = NOTES.index(key)
        self.scale = MAJOR_SCALE if mode == 'major' else MINOR_SCALE
        
    def get_chord_root(self, roman_numeral: str) -> int:
        """Convert Roman numeral to MIDI root note"""
        roman = roman_numeral.upper()
        if roman == 'I': degree = 0
        elif roman == 'II': degree = 1
        elif roman == 'III': degree = 2
        elif roman == 'IV': degree = 3
        elif roman == 'V': degree = 4
        elif roman == 'VI': degree = 5
        elif roman == 'VII': degree = 6
        else: degree = 0
        
        root_note = self.scale[degree % len(self.scale)]
        return 60 + self.key_index + root_note
    
    def generate_progression(self, style: str = 'pop', 
                             bars: int = 4) -> List[Chord]:
        """Generate a chord progression"""
        progression = self.PROGRESSIONS.get(style, self.PROGRESSIONS['pop'])
        chords = []
        
        for i, roman in enumerate(progression * (bars // len(progression) + 1)):
            if i >= bars:
                break
            root = self.get_chord_root(roman)
            chord_type = 'minor' if roman.islower() else 'major'
            chords.append(Chord(root=root, chord_type=chord_type, duration=1.0))
        
        return chords[:bars]

# =============================================================================
# AUDIO SYNTHESIZER
# =============================================================================

class SimpleSynth:
    """Simple synthesizer for generating audio from MIDI data"""
    
    def __init__(self, sample_rate: int = 44100):
        self.sample_rate = sample_rate
        
    def generate_sine_wave(self, frequency: float, duration: float, 
                          amplitude: float = 0.3) -> bytes:
        """Generate a sine wave"""
        samples = int(duration * self.sample_rate)
        wave_data = bytearray()
        
        for i in range(samples):
            t = i / self.sample_rate
            value = amplitude * math.sin(2 * math.pi * frequency * t)
            pcm_value = int(value * 32767)
            wave_data.extend(struct.pack('<h', pcm_value))
        
        return bytes(wave_data)
    
    def generate_saw_wave(self, frequency: float, duration: float,
                         amplitude: float = 0.3) -> bytes:
        """Generate a saw wave"""
        samples = int(duration * self.sample_rate)
        wave_data = bytearray()
        
        for i in range(samples):
            t = i / self.sample_rate
            value = amplitude * (2 * (t * frequency % 1) - 1)
            pcm_value = int(value * 32767)
            wave_data.extend(struct.pack('<h', pcm_value))
        
        return bytes(wave_data)
    
    def apply_envelope(self, audio_data: bytes, attack: float = 0.01,
                      decay: float = 0.1, sustain: float = 0.7,
                      release: float = 0.3) -> bytes:
        """Apply ADSR envelope to audio"""
        samples = len(audio_data) // 2
        envelope = bytearray()
        
        attack_samples = int(attack * self.sample_rate)
        decay_samples = int(decay * self.sample_rate)
        release_samples = int(release * self.sample_rate)
        
        for i in range(samples):
            pcm_value = struct.unpack('<h', audio_data[i*2:(i+1)*2])[0]
            
            if i < attack_samples:
                env = i / attack_samples
            elif i < attack_samples + decay_samples:
                env = 1.0 - (1.0 - sustain) * ((i - attack_samples) / decay_samples)
            elif i < samples - release_samples:
                env = sustain
            else:
                env = sustain * (1.0 - (i - (samples - release_samples)) / release_samples)
            
            new_value = int(pcm_value * env)
            envelope.extend(struct.pack('<h', max(-32768, min(32767, new_value))))
        
        return bytes(envelope)
    
    def save_wav(self, audio_data: bytes, filename: str):
        """Save audio data to WAV file"""
        with wave.open(filename, 'wb') as wav_file:
            wav_file.setnchannels(1)
            wav_file.setsampwidth(2)
            wav_file.setframerate(self.sample_rate)
            wav_file.writeframes(audio_data)

# =============================================================================
# MUSIC GENERATOR
# =============================================================================

class MusicGenerator:
    """Main music generation system"""
    
    def __init__(self, tempo: float = 120.0):
        self.tempo = tempo
        self.markov_gen = MarkovMelodyGenerator()
        self.chord_gen = ChordProgressionGenerator()
        self.synth = SimpleSynth()
        
        # Pre-built training data for melodies
        self.training_sequences = [
            [60, 62, 64, 65, 67, 69, 71, 72],  # C major scale
            [72, 71, 69, 67, 65, 64, 62, 60],  # Descending
            [60, 64, 67, 72, 67, 64, 60, 57],  # Arpeggio
            [60, 60, 62, 64, 65, 65, 64, 62],  # Simple melody
            [67, 69, 71, 72, 71, 69, 67, 65],  # G major
            [65, 67, 69, 70, 72, 74, 76, 77],  # A minor
        ]
        
        # Train the Markov generator
        self.markov_gen.train(self.training_sequences)
    
    def generate_melody_track(self, duration_bars: int = 8,
                             style: str = 'melodic') -> Track:
        """Generate a melody track"""
        track = Track(instrument='piano', tempo=self.tempo)
        
        melody = self.markov_gen.generate(length=duration_bars * 4)
        current_time = 0.0
        
        for i, pitch in enumerate(melody):
            if random.random() < 0.3:
                continue  # Skip some notes for rests
            
            duration = random.choice([0.25, 0.5, 0.5, 1.0])
            velocity = random.randint(60, 100)
            
            track.add_note(Note(
                pitch=pitch,
                duration=duration,
                velocity=velocity,
                start_time=current_time
            ))
            
            current_time += duration
        
        return track
    
    def generate_chord_track(self, duration_bars: int = 8,
                            style: str = 'pop') -> Track:
        """Generate a chord accompaniment track"""
        track = Track(instrument='strings', tempo=self.tempo)
        
        progressions = self.chord_gen.generate_progression(
            style=style, bars=duration_bars)
        
        current_time = 0.0
        for chord in progressions:
            # Play chord as arpeggio
            notes = chord.get_notes()
            for j, note_pitch in enumerate(notes):
                track.add_note(Note(
                    pitch=note_pitch + 48,
                    duration=0.5,
                    velocity=random.randint(40, 70),
                    start_time=current_time + j * 0.25
                ))
            current_time += chord.duration
        
        return track
    
    def generate_bass_track(self, duration_bars: int = 8) -> Track:
        """Generate a bass line"""
        track = Track(instrument='bass', tempo=self.tempo)
        
        progressions = self.chord_gen.generate_progression(
            style='pop', bars=duration_bars)
        
        current_time = 0.0
        for chord in progressions:
            track.add_note(Note(
                pitch=chord.root - 12,
                duration=chord.duration,
                velocity=random.randint(70, 90),
                start_time=current_time
            ))
            current_time += chord.duration
        
        return track
    
    def generate_composition(self, 
                            melody_style: str = 'melodic',
                            chord_style: str = 'pop',
                            duration_bars: int = 8) -> Composition:
        """Generate a complete composition"""
        composition = Composition(
            title=f"AI Generated - {chord_style}",
            tempo=self.tempo
        )
        
        melody_track = self.generate_melody_track(duration_bars, melody_style)
        chord_track = self.generate_chord_track(duration_bars, chord_style)
        bass_track = self.generate_bass_track(duration_bars)
        
        composition.add_track('melody', melody_track)
        composition.add_track('chords', chord_track)
        composition.add_track('bass', bass_track)
        
        return composition
    
    def render_to_audio(self, composition: Composition,
                       output_path: str) -> str:
        """Render composition to audio file"""
        total_duration = composition.total_duration()
        samples = int(total_duration * 60 / self.tempo * self.synth.sample_rate)
        
        master_audio = bytearray(samples * 2)
        for i in range(samples):
            struct.pack_into('<h', master_audio, i * 2, 0)
        
        for track_name, track in composition.tracks.items():
            track_audio = bytearray(samples * 2)
            
            for note in track.notes:
                freq = note.get_frequency()
                duration_beats = note.duration
                duration_sec = duration_beats * 60 / self.tempo
                
                if track.instrument == 'bass':
                    wave_data = self.synth.generate_saw_wave(freq, duration_sec)
                elif track.instrument == 'strings':
                    wave_data = self.synth.generate_saw_wave(freq, duration_sec)
                else:
                    wave_data = self.synth.generate_sine_wave(freq, duration_sec)
                
                wave_data = self.synth.apply_envelope(wave_data)
                
                start_sample = int(note.start_time * self.synth.sample_rate)
                
                for i in range(min(len(wave_data) // 2, samples - start_sample)):
                    pos = (start_sample + i) * 2
                    track_val = struct.unpack('<h', track_audio[pos:pos+2])[0]
                    wave_val = struct.unpack('<h', wave_data[i*2:i*2+2])[0]
                    new_val = int((track_val + wave_val) / 2)
                    struct.pack_into('<h', track_audio, pos, 
                                    max(-32768, min(32767, new_val)))
            
            for i in range(samples):
                pos = i * 2
                master_val = struct.unpack('<h', master_audio[pos:pos+2])[0]
                track_val = struct.unpack('<h', track_audio[pos:pos+2])[0]
                new_val = int((master_val + track_val) / 2)
                struct.pack_into('<h', master_audio, pos,
                                max(-32768, min(32767, new_val)))
        
        with wave.open(output_path, 'wb') as wav_file:
            wav_file.setnchannels(1)
            wav_file.setsampwidth(2)
            wav_file.setframerate(self.synth.sample_rate)
            wav_file.writeframes(bytes(master_audio))
        
        return output_path

# =============================================================================
# MAIN APP INTERFACE
# =============================================================================

class MusicGenerationApp:
    """Main application class for AI music generation"""
    
    def __init__(self):
        self.generator = MusicGenerator()
        self.compositions = {}
        self.output_dir = tempfile.mkdtemp(prefix='music_gen_')
        print(f"Music generation output directory: {self.output_dir}")
    
    def generate_song(self, style: str = 'pop', 
                     duration_bars: int = 16,
                     name: str = None) -> Composition:
        """Generate a complete song"""
        if name is None:
            name = f"Song_{len(self.compositions) + 1}"
        
        composition = self.generator.generate_composition(
            melody_style=style,
            chord_style=style,
            duration_bars=duration_bars
        )
        composition.title = name
        self.compositions[name] = composition
        
        return composition
    
    def export_to_midi(self, composition: Composition, filename: str) -> str:
        """Export composition to MIDI file"""
        filepath = os.path.join(self.output_dir, f"{filename}.mid")
        
        with open(filepath, 'wb') as f:
            # Header chunk
            f.write(b'MThd')
            f.write(struct.pack('>I', 6))
            f.write(struct.pack('>H', 1))
            f.write(struct.pack('>H', len(composition.tracks)))
            f.write(struct.pack('>H', 480))
            
            # Track chunks
            for track_name, track in composition.tracks.items():
                events = []
                
                # Tempo event
                tempo_microsec = int(60000000 / track.tempo)
                events.append((0, 0xFF, 0x51, struct.pack('>I', tempo_microsec)[1:]))
                
                # Note events
                for note in track.notes:
                    events.append((int(note.start_time * 480), 0x90, note.pitch, note.velocity))
                    events.append((int((note.start_time + note.duration) * 480), 0x80, note.pitch, 0))
                
                events.sort(key=lambda x: x[0])
                
                track_data = bytearray()
                for delta, status, data1, data2 in events:
                    while delta >= 0x80:
                        track_data.append((delta & 0x7F) | 0x80)
                        delta >>= 7
                    track_data.append(delta)
                    track_data.append(status)
                    track_data.append(data1)
                    if status != 0xFF:
                        track_data.append(data2)
                
                track_data.extend([0x00, 0xFF, 0x2F, 0x00])
                
                f.write(b'MTrk')
                f.write(struct.pack('>I', len(track_data)))
                f.write(track_data)
        
        return filepath
    
    def export_to_wav(self, composition: Composition, filename: str) -> str:
        """Export composition to WAV audio file"""
        filepath = os.path.join(self.output_dir, f"{filename}.wav")
        return self.generator.render_to_audio(composition, filepath)
    
    def print_composition(self, composition: Composition):
        """Print composition details"""
        print(f"\n{'='*60}")
        print(f"Composition: {composition.title}")
        print(f"Tempo: {composition.tempo} BPM")
        print(f"Duration: {composition.total_duration():.2f} beats")
        print(f"{'='*60}")
        
        for track_name, track in composition.tracks.items():
            print(f"\nTrack: {track_name} ({track.instrument})")
            print(f"  Notes: {len(track.notes)}")
            if track.notes:
                print(f"  First note: {track.notes[0].get_note_name()}")
                print(f"  Last note: {track.notes[-1].get_note_name()}")

# =============================================================================
# DEMONSTRATION
# =============================================================================

print("=" * 60)
print("AI MUSIC GENERATION APP")
print("=" * 60)

app = MusicGenerationApp()

styles = ['pop', 'jazz', 'emotional', 'ambient']

for style in styles:
    print(f"\nGenerating {style} style composition...")
    composition = app.generate_song(style=style, duration_bars=8, name=style)
    app.print_composition(composition)

first_composition = app.compositions['pop']

print("\n" + "=" * 60)
print("EXPORTING FILES")
print("=" * 60)

midi_file = app.export_to_midi(first_composition, 'ai_generated_song')
print(f"MIDI file saved: {midi_file}")

wav_file = app.export_to_wav(first_composition, 'ai_generated_song')
print(f"WAV file saved: {wav_file}")

print(f"\nMIDI file size: {os.path.getsize(midi_file)} bytes")
print(f"WAV file size: {os.path.getsize(wav_file)} bytes")

print("\n" + "=" * 60)
print("MUSIC GENERATION COMPLETE!")
print("=" * 60)

print(f"\nOutput directory: {app.output_dir}")
print("Files created:")
for f in os.listdir(app.output_dir):
    filepath = os.path.join(app.output_dir, f)
    print(f"  - {f} ({os.path.getsize(filepath)} bytes)")