import msgpack
import enum
import numpy as np
import msgpack_numpy
import pedalboard
from dataclasses import dataclass, asdict


msgpack_numpy.patch()

@dataclass
class TimeSignature:
    beats_per_measure: float
    note_value: float

class ChunkType(enum.Enum):
    CHUNK = 1
    AUDIO = enum.auto()
    MIDI = enum.auto()

class ChannelChunk:
    def __init__(self, channel, position: float = 0.0, name: str = "Chunk", chunk_type: ChunkType = ChunkType.CHUNK):
        self.channel = channel
        self.position = position # position is how many bars into the song the chunk is
        self.name = name
        self.chunk_type = chunk_type
        
    def render(self):
        pass
        
    def from_json(json: dict, channel) -> ChannelChunk:
        return ChannelChunk(
            channel,
            chunk_type = ChunkType(json["type"]),
            name = json["name"],
            position = json["position"]
        )
        
    def to_json(self):
        return {
            "type": self.chunk_type.value,
            "name": self.name,
            "position": self.position
        }

class AudioChannelChunk(ChannelChunk):
    def __init__(self, channel, audio_data: np.ndarray, sample_rate: int, position: float = 0.0, name: str = "Sample"):
        super().__init__(channel, position, name, chunk_type=ChunkType.AUDIO)
        self.audio_data = audio_data
        self.sample_rate = sample_rate
        
    def render(self):
        start_sample = int(self.position * self.channel.project.samples_per_bar)
        
        audio = self.audio_data.T
        
        # ensure stereo
        if audio.ndim == 1:
            audio = np.stack([audio, audio], axis=1)
            
        end_sample = start_sample + len(audio)

        return start_sample, end_sample, audio
        
    def from_json(json: dict, channel) -> ChannelChunk:
        return AudioChannelChunk(
            channel,
            name = json["name"],
            position = json["position"],
            audio_data = json["audio_data"],
            sample_rate = json["sample_rate"]
        )
        
    def to_json(self):
        return {
            "type": self.chunk_type.value,
            "name": self.name,
            "position": self.position,
            "audio_data": self.audio_data,
            "sample_rate": self.sample_rate
        }

chunk_type_associations = {
    ChunkType.CHUNK: ChannelChunk,
    ChunkType.AUDIO: AudioChannelChunk
}

class ProjectChannel:
    def __init__(self, project, name: str = "", volume: int = 0, pan: int = 0, mute: bool = False, solo: bool = False, channel_chunks: list[ChannelChunk] = []):
        self.project = project
        self.name = name
        self.volume = volume
        self.pan = pan
        self.mute = mute
        self.solo = solo
        self.chunks = []
        
        self.board = pedalboard.Pedalboard()
        
    def pan_stereo(self, stereo, pan):
        pan = np.clip(pan, -1.0, 1.0)

        left_gain = np.cos((pan + 1) * np.pi / 4)
        right_gain = np.sin((pan + 1) * np.pi / 4)

        out = stereo.copy()
        out[:, 0] *= left_gain
        out[:, 1] *= right_gain

        return out
        
    def render(self):
        buffer = np.zeros((self.project.total_song_samples, 2), dtype=np.float32)
        
        # render each chunk
        for chunk in self.chunks:
            start, end, audio = chunk.render()
            buffer[start:end] += audio
            
        # apply effects
        #buffer = self.board(buffer, self.project.sample_rate)
            
        # apply volume
        gain = 10 ** (self.volume / 20)
        buffer *= gain
        
        # pan
        buffer = self.pan_stereo(buffer, self.pan/100)
        
        return buffer
        
    def from_json(json: dict, project) -> ProjectChannel:
        channel = ProjectChannel(
            project,
            name = json["name"],
            volume = json["volume"],
            pan = json["pan"],
            mute = json["mute"],
            solo = json["solo"],
            
        )
        
        channel.chunks = [chunk_type_associations[ChunkType(chunk["type"])].from_json(chunk, channel) for chunk in json["chunks"]]
        
        return channel
        
    def to_json(self):
        return {
            "name": self.name,
            "volume": self.volume,
            "pan": self.pan,
            "mute": self.mute,
            "solo": self.solo,
            "chunks": [chunk.to_json() for chunk in self.chunks]
        }

class Project:
    def __init__(self, channels: list[ProjectChannel] = [], version: float = 1.0, bpm: float = 120, time_signature: TimeSignature = TimeSignature(4, 4), song_length: float = 16, sample_rate: int = 44100):
        self.version = version
        self.bpm = bpm
        self.sample_rate = sample_rate
        self.time_signature = time_signature
        self.channels = channels
        self.song_length = song_length # length of the song in bars
        
        self.seconds_per_bar = (60.0 / self.bpm) * self.time_signature.beats_per_measure
        self.samples_per_bar = int(self.seconds_per_bar * self.sample_rate)
        self.total_song_samples = int(self.samples_per_bar * self.song_length)
        
    def render(self):
        buffer = np.zeros((self.total_song_samples, 2), dtype=np.float32)
        
        for channel in self.channels:
            buffer += channel.render()
        
        return buffer
        
    def from_file(file_path: str) -> Project:
        with open(file_path, "rb") as f:
            return Project.from_json(msgpack.unpackb(f.read()))
            
    def from_json(json: dict) -> Project:
        project = Project(
            version = json["version"],
            time_signature = TimeSignature(json["time_signature"]["beats_per_measure"], json["time_signature"]["note_value"]),
            bpm = json["bpm"],
            song_length = json["song_length"],
            sample_rate = json["sample_rate"]
        )
        
        project.channels = [ProjectChannel.from_json(channel, project) for channel in json["channels"]]
        
        return project
        
    def to_json(self):
        return {
            "version": self.version,
            "time_signature": asdict(self.time_signature),
            "bpm": self.bpm,
            "channels": [channel.to_json() for channel in self.channels],
            "song_length": self.song_length,
            "sample_rate": self.sample_rate
        }
        
    def write_to_file(self, file_path: str):
        with open(file_path, "wb") as f:
            f.write(msgpack.packb(self.to_json()))