import msgpack
import enum
import numpy as np
import msgpack_numpy
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, position: float = 0.0, name: str = "Chunk", chunk_type: ChunkType = ChunkType.CHUNK):
        self.position = position # position is how many bars into the song the chunk is
        self.name = name
        self.chunk_type = chunk_type
        
    def from_json(json: dict) -> ChannelChunk:
        return ChannelChunk(
            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, audio_data: np.ndarray, sample_rate: int, position: float = 0.0, name: str = "Sample"):
        super().__init__(position, name, chunk_type=ChunkType.AUDIO)
        self.audio_data = audio_data
        self.sample_rate = sample_rate
        
    def from_json(json: dict) -> ChannelChunk:
        return AudioChannelChunk(
            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, name: str = "", volume: int = 0, pan: int = 0, mute: bool = False, solo: bool = False, chunks: list[ChannelChunk] = []):
        self.name = name
        self.volume = volume
        self.pan = pan
        self.mute = mute
        self.solo = solo
        self.chunks = chunks
        
    def from_json(json: dict) -> ProjectChannel:
        return ProjectChannel(
            name = json["name"],
            volume = json["volume"],
            pan = json["pan"],
            mute = json["mute"],
            solo = json["solo"],
            chunks = [chunk_type_associations[ChunkType(chunk["type"])].from_json(chunk) for chunk in json["chunks"]]
        )
        
    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)):
        self.version = version
        self.bpm = bpm
        self.time_signature = time_signature
        self.channels = channels
        
    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:
        return Project(
            version = json["version"],
            time_signature = TimeSignature(json["time_signature"]["beats_per_measure"], json["time_signature"]["note_value"]),
            bpm = json["bpm"],
            channels = [ProjectChannel.from_json(channel) for channel in json["channels"]]
        )
        
    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]
        }
        
    def write_to_file(self, file_path: str):
        with open(file_path, "wb") as f:
            f.write(msgpack.packb(self.to_json()))