import pygame
import os
from tja_parser import TJA, Course, Note

class EditorState:
    def __init__(self):
        self.tja = None
        self.current_course_name = 'Oni'
        self.current_time = 0.0 # ms
        self.is_playing = False
        self.zoom_x = 0.2 # pixels per ms
        self.scroll_y = 0
        self.selected_notes = [] # List of Note objects
        self.file_path = None
        self.audio_path = None
        self.snap_grid = 4 # 1/4 beat (16th note)

    def load_file(self, path):
        if path.lower().endswith('.tja'):
            self.file_path = path
            self.tja = TJA(path)
            # Find audio
            wave = self.tja.headers.get('WAVE', '')
            base_dir = os.path.dirname(path)
            self.audio_path = os.path.join(base_dir, wave)
        elif path.lower().endswith('.ogg'):
            self.audio_path = path
            # Check for TJA
            base_name = os.path.splitext(path)[0]
            tja_path = base_name + '.tja'
            if os.path.exists(tja_path):
                self.file_path = tja_path
                self.tja = TJA(tja_path)
            else:
                self.tja = TJA()
                self.tja.headers['WAVE'] = os.path.basename(path)
                self.tja.headers['TITLE'] = os.path.basename(base_name)
        
        self.load_audio()

    def load_audio(self):
        if self.audio_path and os.path.exists(self.audio_path):
            try:
                pygame.mixer.music.load(self.audio_path)
            except Exception as e:
                print(f"Failed to load audio: {e}")

    def get_current_course(self):
        if not self.tja: return None
        if self.current_course_name not in self.tja.courses:
            self.tja.courses[self.current_course_name] = Course(self.current_course_name)
        return self.tja.courses[self.current_course_name]

    def toggle_play(self):
        if self.is_playing:
            pygame.mixer.music.pause()
            self.is_playing = False
        else:
            # Sync pygame music to current time
            try:
                # pygame.mixer.music.play(start=self.current_time / 1000.0) 
                # Note: 'start' in play() is usually working, but set_pos might be needed depending on implementation
                if self.current_time < 0:
                     # Start from 0 if negative (pre-song offset)
                     pygame.mixer.music.play(start=0)
                else:
                    pygame.mixer.music.play(start=self.current_time / 1000.0)
                self.is_playing = True
            except:
                pass

    def update(self, dt):
        if self.is_playing:
            self.current_time += dt
            # Optional: Sync with mixer position to avoid drift
            # mixer_pos = pygame.mixer.music.get_pos() # Returns ms played since start of 'play'
            # This is tricky because get_pos resets on play().
            # Simple dt addition is often smoother for short edits.

    def save(self):
        if self.tja and self.file_path:
            self.tja.save(self.file_path)
        elif self.tja and self.audio_path:
             # Save as .tja next to ogg
             base = os.path.splitext(self.audio_path)[0]
             self.file_path = base + ".tja"
             self.tja.save(self.file_path)

    def record_note(self, note_type):
        """Adds a note at current_time, snapped to grid"""
        course = self.get_current_course()
        if not course: return

        # Calculate snapped time
        bpm = self.tja.headers.get('BPM', 120)
        beat_ms = 60000 / bpm
        snap_ms = beat_ms / self.snap_grid # e.g. 1/4 beat = 16th note
        
        # Quantize current time
        raw_time = self.current_time
        snapped_time = round(raw_time / snap_ms) * snap_ms
        
        # Calculate Measure/Beat
        beat_total = snapped_time / beat_ms
        measure = int(beat_total / 4) # Assuming 4/4
        beat = beat_total % 4
        
        # Avoid duplicate at same time?
        # For now, just add. Or remove existing note at this spot?
        # Simple overwrite logic:
        existing = [n for n in course.notes if abs(n.time - snapped_time) < 5]
        for n in existing:
            course.notes.remove(n)
            
        self.add_note(note_type, snapped_time, measure, beat)

    def add_note(self, note_type, time, measure, beat):
        course = self.get_current_course()
        if course:
            new_note = Note(note_type, time, measure, beat)
            course.notes.append(new_note)

    def remove_selected(self):
        course = self.get_current_course()
        if course and self.selected_notes:
            for n in self.selected_notes:
                if n in course.notes:
                    course.notes.remove(n)
            self.selected_notes = []