import os
import re
import math

class Note:
    def __init__(self, note_type, time_ms, measure_idx, beat_pos, duration_ms=0):
        self.type = int(note_type)  # 1=Don, 2=Ka, 3=BigDon, 4=BigKa, etc.
        self.time = time_ms
        self.measure = measure_idx
        self.beat = beat_pos # Beat index within the measure (e.g. 0.0, 0.5, 1.0)
        self.duration = duration_ms # For rolls/balloons

    def __repr__(self):
        return f"Note(type={self.type}, time={self.time:.1f}, m={self.measure}, b={self.beat:.2f})"

class Event:
    def __init__(self, event_type, value, time_ms, measure_idx, beat_pos):
        self.type = event_type # 'bpm', 'scroll', 'measure', 'barline'
        self.value = value
        self.time = time_ms
        self.measure = measure_idx
        self.beat = beat_pos

class Course:
    def __init__(self, difficulty):
        self.difficulty = difficulty
        self.level = 0
        self.balloon = []
        self.score_init = 0
        self.score_diff = 0
        self.notes = []
        self.events = []
        # Raw lines for reference or partial parsing
        self.lines = [] 

class TJA:
    def __init__(self, file_path=None):
        self.headers = {
            'TITLE': 'New Song',
            'SUBTITLE': '',
            'BPM': 120.0,
            'WAVE': 'song.ogg',
            'OFFSET': 0.0,
            'DEMOSTART': 0.0,
            'SONGVOL': 100,
            'SEVOL': 100
        }
        self.courses = {} # 'Easy', 'Normal', 'Hard', 'Oni', 'Edit'
        self.common_lines = [] # Lines before any COURSE: command
        
        if file_path and os.path.exists(file_path):
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
            except UnicodeDecodeError:
                 with open(file_path, 'r', encoding='shift_jis', errors='replace') as f:
                    content = f.read()
            self.parse(content)
        else:
            # Initialize default courses
            for diff in ['Easy', 'Normal', 'Hard', 'Oni', 'Edit']:
                self.courses[diff] = Course(diff)

    def parse(self, text):
        lines = text.splitlines()
        current_course_obj = None
        current_lines = self.common_lines

        # First pass: Separate lines by COURSE
        for line in lines:
            line = line.strip()
            if not line: 
                continue
                
            if line.startswith('COURSE:'):
                course_name = line.split(':', 1)[1].strip()
                # Normalize course name
                c_map = {'0': 'Easy', 'Easy': 'Easy', '1': 'Normal', 'Normal': 'Normal',
                         '2': 'Hard', 'Hard': 'Hard', '3': 'Oni', 'Oni': 'Oni',
                         '4': 'Edit', 'Edit': 'Edit'}
                c_name = c_map.get(course_name, course_name)
                
                if c_name not in self.courses:
                    self.courses[c_name] = Course(c_name)
                current_course_obj = self.courses[c_name]
                current_lines = current_course_obj.lines
            elif line.startswith('#START'):
                # Start of chart data, but we just store lines for now
                if current_course_obj:
                    current_lines.append(line)
            elif ':' in line and not line.startswith('#'):
                # Header
                key, val = line.split(':', 1)
                key = key.strip().upper()
                val = val.strip()
                
                if current_course_obj is None:
                    # Global header
                    if key in ['TITLE', 'SUBTITLE', 'WAVE']:
                        self.headers[key] = val
                    elif key in ['BPM', 'OFFSET', 'DEMOSTART', 'SONGVOL', 'SEVOL']:
                        try:
                            self.headers[key] = float(val)
                        except:
                            pass
                else:
                    # Course specific header
                    if key == 'LEVEL':
                        try: current_course_obj.level = int(val)
                        except: pass
                    elif key == 'BALLOON':
                        current_course_obj.balloon = [int(x) for x in val.split(',') if x.strip()]
                    elif key == 'SCOREINIT':
                        current_course_obj.score_init = int(val.split(',')[0])
                    elif key == 'SCOREDIFF':
                        current_course_obj.score_diff = int(val)
            else:
                current_lines.append(line)

        # Second pass: Process notes for each course
        base_bpm = self.headers['BPM']
        base_offset = self.headers['OFFSET']
        
        for course_name, course in self.courses.items():
            self._parse_course_notes(course, base_bpm, base_offset)

    def _parse_course_notes(self, course, bpm, offset):
        course.notes = []
        course.events = []
        
        current_bpm = bpm
        current_measure_frac = 4.0 / 4.0 # 4/4 time
        
        # Calculate start time based on OFFSET
        # TJA OFFSET: "A negative value makes the song start earlier (notes appear later)"
        # Wait, TJA spec says: OFFSET is the time (in seconds) that the music starts relative to the first beat.
        # If OFFSET is -1.9, music starts 1.9s BEFORE the first beat.
        # So first beat is at +1.9s relative to music start.
        # Current Time Pointer
        current_time = -offset * 1000 # Convert to ms
        
        # We need to process lines sequentially
        # Notes are comma-separated.
        # Commands start with #
        
        measure_buffer = ""
        measure_idx = 0
        
        for line in course.lines:
            line = line.split('//')[0].strip() # Remove comments
            if not line: continue
            
            if line.startswith('#START'):
                current_time = -offset * 1000
                measure_idx = 0
                continue
            
            if line.startswith('#END'):
                break
                
            if line.startswith('#'):
                # Command
                if line.startswith('#BPMCHANGE'):
                    try:
                        val = float(line.split()[1])
                        course.events.append(Event('bpm', val, current_time, measure_idx, 0))
                        current_bpm = val
                    except: pass
                elif line.startswith('#MEASURE'):
                    try:
                        m_str = line.split()[1]
                        num, den = map(int, m_str.split('/'))
                        current_measure_frac = num / den
                        course.events.append(Event('measure', current_measure_frac, current_time, measure_idx, 0))
                    except: pass
                elif line.startswith('#SCROLL'):
                    try:
                        val = float(line.split()[1])
                        course.events.append(Event('scroll', val, current_time, measure_idx, 0))
                    except: pass
                elif line.startswith('#GOGOSTART'):
                    course.events.append(Event('gogo', True, current_time, measure_idx, 0))
                elif line.startswith('#GOGOEND'):
                    course.events.append(Event('gogo', False, current_time, measure_idx, 0))
                elif line.startswith('#BARLINEOFF'):
                    course.events.append(Event('barline', False, current_time, measure_idx, 0))
                elif line.startswith('#BARLINEON'):
                    course.events.append(Event('barline', True, current_time, measure_idx, 0))
                continue
            
            # Accumulate measure string
            measure_buffer += line
            if measure_buffer.endswith(','):
                # Process measure
                notes_str = measure_buffer[:-1] # Remove comma
                measure_buffer = ""
                
                # Calculate duration of this measure
                # beats per measure = 4 * current_measure_frac (e.g. 4/4 -> 4 beats, 3/4 -> 3 beats)
                beats_in_measure = 4 * current_measure_frac
                measure_duration = (60000 / current_bpm) * beats_in_measure
                
                num_notes = len(notes_str)
                if num_notes > 0:
                    time_per_char = measure_duration / num_notes
                    beat_per_char = beats_in_measure / num_notes
                    
                    for i, char in enumerate(notes_str):
                        if char == '0': continue
                        if char in '1234': # Don, Ka, BigDon, BigKa
                            note_time = current_time + (i * time_per_char)
                            note_beat = i * beat_per_char
                            course.notes.append(Note(char, note_time, measure_idx, note_beat))
                        # Note: Logic for rolls (5,6,7,8) is more complex (start/end), simplified here for MVP
                        elif char in '5679': # Rolls/Balloon Start
                             note_time = current_time + (i * time_per_char)
                             note_beat = i * beat_per_char
                             course.notes.append(Note(char, note_time, measure_idx, note_beat))
                        elif char == '8': # End roll
                             note_time = current_time + (i * time_per_char)
                             note_beat = i * beat_per_char
                             course.notes.append(Note(char, note_time, measure_idx, note_beat))

                current_time += measure_duration
                measure_idx += 1

    def save(self, file_path):
        # Basic reconstruction logic
        with open(file_path, 'w', encoding='utf-8') as f:
            # Write Headers
            for k, v in self.headers.items():
                if v:
                    f.write(f"{k}:{v}\n")
            f.write("\n")
            
            # Write Courses
            for name, course in self.courses.items():
                if not course.notes and not course.events:
                    continue # Skip empty courses
                
                f.write(f"COURSE:{name}\n")
                f.write(f"LEVEL:{course.level}\n")
                if course.balloon:
                    f.write(f"BALLOON:{','.join(map(str, course.balloon))}\n")
                f.write(f"SCOREINIT:{course.score_init}\n")
                f.write(f"SCOREDIFF:{course.score_diff}\n")
                f.write("\n#START\n")
                
                # Reconstruct measures from notes
                # This is the hard part: "Export" logic
                # For MVP, if we haven't edited the structure deeply, we could just dump the parsed events?
                # But the prompt implies editing.
                # Let's try a simple generation: 
                # Group notes by measure index.
                
                max_measure = 0
                if course.notes:
                    max_measure = max(n.measure for n in course.notes)
                if course.events:
                    max_measure = max(max_measure, max(e.measure for e in course.events))
                
                current_bpm = self.headers['BPM']
                
                for m in range(max_measure + 1):
                    # Find events in this measure
                    m_events = [e for e in course.events if e.measure == m]
                    for e in m_events:
                        if e.type == 'bpm': f.write(f"#BPMCHANGE {e.value}\n")
                        if e.type == 'measure': f.write(f"#MEASURE {int(e.value*4)}/4\n") # Simplified
                        if e.type == 'gogo': f.write("#GOGOSTART\n" if e.value else "#GOGOEND\n")
                    
                    # Construct note string
                    # Default to 16 divisions per measure (standard) or higher if needed
                    # Find notes in this measure
                    m_notes = sorted([n for n in course.notes if n.measure == m], key=lambda x: x.beat)
                    
                    # Determine required resolution
                    # Simple approach: Use 16 chars per measure (4 chars per beat)
                    grid_size = 16 
                    line = ['0'] * grid_size
                    
                    for n in m_notes:
                        # map beat (0..4) to index (0..16)
                        # beat is 0.0 to 4.0 (in 4/4)
                        # index = (beat / 4.0) * 16 = beat * 4
                        idx = int(round(n.beat * (grid_size / 4.0))) # Assuming 4/4
                        if 0 <= idx < grid_size:
                            line[idx] = str(n.type)
                    
                    f.write("".join(line) + ",\n")
                
                f.write("#END\n\n")