from __future__ import annotations
import json
import math
from pathlib import Path
from typing import Optional

KNOWLEDGE_DIR = Path(__file__).parent.parent.parent / "knowledge"

NOTE_NAMES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]

SCALE_INTERVALS = {
    "major": [0, 2, 4, 5, 7, 9, 11],
    "minor": [0, 2, 3, 5, 7, 8, 10],
    "harmonic_minor": [0, 2, 3, 5, 7, 8, 11],
    "melodic_minor": [0, 2, 3, 5, 7, 9, 11],
    "dorian": [0, 2, 3, 5, 7, 9, 10],
    "phrygian": [0, 1, 3, 5, 7, 8, 10],
}

CHORD_TYPES = {
    "maj": [0, 4, 7],
    "min": [0, 3, 7],
    "dim": [0, 3, 6],
    "aug": [0, 4, 8],
    "7": [0, 4, 7, 10],
    "m7": [0, 3, 7, 10],
    "sus2": [0, 2, 7],
    "sus4": [0, 5, 7],
}


def note_to_midi(note_str: str) -> int:
    note_str = note_str.strip()
    if len(note_str) == 1:
        name = note_str[0]
        octave = 4
    elif len(note_str) == 2:
        if note_str[1] == "#":
            name = note_str[:2]
            octave = 4
        else:
            name = note_str[0]
            octave = int(note_str[1])
    else:
        name = note_str[:2] if note_str[1] == "#" else note_str[0]
        octave = int(note_str[-1])

    base = NOTE_NAMES.index(name)
    return (octave + 1) * 12 + base


def parse_chord_name(chord_str: str) -> tuple[int, list[int]]:
    chord_str = chord_str.strip()
    root = chord_str[0]
    idx = 1
    if len(chord_str) > 1 and chord_str[1] == "#":
        root += "#"
        idx = 2

    suffix = chord_str[idx:]
    if suffix == "m" or suffix == "min":
        chord_type = "min"
    elif suffix == "dim":
        chord_type = "dim"
    elif suffix == "7":
        chord_type = "7"
    elif suffix == "m7":
        chord_type = "m7"
    elif suffix == "sus2":
        chord_type = "sus2"
    elif suffix == "sus4":
        chord_type = "sus4"
    elif suffix == "":
        chord_type = "maj"
    else:
        chord_type = "maj"

    root_midi = note_to_midi(root + "4")
    intervals = CHORD_TYPES.get(chord_type, [0, 4, 7])
    return root_midi, intervals


def generate_dembow(bars: int = 8, ppq: int = 96) -> list[dict]:
    notes = []
    for bar in range(bars):
        offset = bar * 4.0
        kick_positions = [0.0, 2.5]
        snare_positions = [2.0, 4.0]
        hihat_positions = [i * 0.5 for i in range(8)]

        for p in kick_positions:
            notes.append({"position": offset + p, "length": 0.25, "key": 36, "velocity": 110})
        for p in snare_positions:
            notes.append({"position": offset + p, "length": 0.15, "key": 38, "velocity": 105})
        for p in hihat_positions:
            notes.append({"position": offset + p, "length": 0.1, "key": 42, "velocity": 75})
        for i in [1, 3, 5, 7]:
            notes.append({
                "position": offset + hihat_positions[i],
                "length": 0.1,
                "key": 46,
                "velocity": 60,
            })
    return notes


def generate_bass_808(
    chord_progression: list[str],
    beats_per_chord: int = 4,
    octave: int = 2,
    bars: int = 8,
) -> list[dict]:
    notes = []
    pos = 0.0
    total_beats = bars * 4
    while pos < total_beats:
        for chord_name in chord_progression:
            root_midi, _ = parse_chord_name(chord_name)
            bass_note = root_midi - (4 - octave) * 12
            notes.append({
                "position": pos,
                "length": min(beats_per_chord - 0.1, total_beats - pos),
                "key": bass_note,
                "velocity": 100,
            })
            pos += beats_per_chord
            if pos >= total_beats:
                break
    return notes


def generate_piano_stabs(
    chord_progression: list[str],
    beats_per_chord: int = 4,
    bars: int = 8,
) -> list[dict]:
    notes = []
    pos = 0.0
    total_beats = bars * 4
    while pos < total_beats:
        for chord_name in chord_progression:
            root_midi, intervals = parse_chord_name(chord_name)
            chord_notes = [root_midi + iv for iv in intervals]
            for stab_pos in [0.5, 1.5, 2.5, 3.5]:
                actual_pos = pos + stab_pos
                if actual_pos >= total_beats:
                    break
                for cn in chord_notes:
                    notes.append({
                        "position": actual_pos,
                        "length": 0.2,
                        "key": cn,
                        "velocity": 70,
                    })
            pos += beats_per_chord
            if pos >= total_beats:
                break
    return notes


def generate_lead_hook(
    chord_progression: list[str],
    beats_per_chord: int = 4,
    bars: int = 8,
    octave: int = 5,
) -> list[dict]:
    notes = []
    pos = 0.0
    total_beats = bars * 4

    hook_patterns = [
        [0, 0.5, 1.0, 2.0, 3.0],
        [0, 1.0, 1.5, 2.0, 3.5],
        [0, 0.25, 0.5, 2.0, 2.5, 3.0],
    ]

    pattern_idx = 0
    while pos < total_beats:
        for chord_name in chord_progression:
            root_midi, intervals = parse_chord_name(chord_name)
            scale_notes = [root_midi + iv for iv in [0, 2, 3, 5, 7, 8, 10]]
            target_octave_notes = [n + (octave - 4) * 12 for n in scale_notes]

            pattern = hook_patterns[pattern_idx % len(hook_patterns)]
            for i, p in enumerate(pattern):
                actual_pos = pos + p
                if actual_pos >= total_beats:
                    break
                note_idx = i % len(target_octave_notes)
                notes.append({
                    "position": actual_pos,
                    "length": 0.4 if i < len(pattern) - 1 else 0.8,
                    "key": target_octave_notes[note_idx],
                    "velocity": 90 if i % 2 == 0 else 75,
                })
            pos += beats_per_chord
            pattern_idx += 1
            if pos >= total_beats:
                break
    return notes


def generate_pad(
    chord_progression: list[str],
    beats_per_chord: int = 4,
    bars: int = 8,
    octave: int = 4,
) -> list[dict]:
    notes = []
    pos = 0.0
    total_beats = bars * 4
    while pos < total_beats:
        for chord_name in chord_progression:
            root_midi, intervals = parse_chord_name(chord_name)
            chord_notes = [root_midi + (octave - 4) * 12 + iv for iv in intervals]
            duration = min(beats_per_chord, total_beats - pos)
            for cn in chord_notes:
                notes.append({
                    "position": pos,
                    "length": duration,
                    "key": cn,
                    "velocity": 45,
                })
            pos += beats_per_chord
            if pos >= total_beats:
                break
    return notes


def generate_latin_perc(bars: int = 8) -> list[dict]:
    notes = []
    for bar in range(bars):
        offset = bar * 4.0
        shaker = [(i * 0.25) + 0.125 for i in range(16)]
        for p in shaker:
            notes.append({"position": offset + p, "length": 0.1, "key": 50, "velocity": 55})
        congas = [0.0, 1.0, 2.0, 3.0]
        for p in congas:
            notes.append({"position": offset + p, "length": 0.2, "key": 54, "velocity": 65})
        rim = [0.75, 2.75]
        for p in rim:
            notes.append({"position": offset + p, "length": 0.1, "key": 37, "velocity": 50})
    return notes


def compose_from_genre(
    genre_path: str | Path,
    custom_overrides: Optional[dict] = None,
) -> dict:
    with open(genre_path, "r", encoding="utf-8") as f:
        genre = json.load(f)

    if custom_overrides:
        genre.update(custom_overrides)

    bpm = genre["bpm"]["default"]
    ppq = genre.get("ppq", 96)
    key = genre["keys"][0]
    progression = genre["chord_progressions"][0]["chords"]
    beats_per_chord = genre["chord_progressions"][0].get("beats_per_chord", 4)
    bars = genre["structure"]["sections"][1]["bars"]

    composition = {
        "meta": {
            "genre": genre["genre"],
            "era": genre.get("era", ""),
            "bpm": bpm,
            "ppq": ppq,
            "key": key,
            "chord_progression": progression,
            "beats_per_chord": beats_per_chord,
            "bars": bars,
        },
        "tracks": [],
    }

    for role_name, role_config in genre["roles"].items():
        track = {
            "role": role_name,
            "description": role_config["description"],
            "preferred_plugins": role_config["preferred_plugins"],
            "mixer_slot": role_config.get("mixer_slot", 0),
        }

        if role_name == "drums":
            track["notes"] = generate_dembow(bars, ppq)
        elif role_name == "bass":
            track["notes"] = generate_bass_808(
                progression, beats_per_chord,
                octave=role_config.get("octave", 2),
                bars=bars,
            )
        elif role_name == "harmony":
            track["notes"] = generate_piano_stabs(progression, beats_per_chord, bars)
        elif role_name == "lead":
            track["notes"] = generate_lead_hook(
                progression, beats_per_chord, bars,
                octave=role_config.get("octave", 5),
            )
        elif role_name == "pad":
            track["notes"] = generate_pad(
                progression, beats_per_chord, bars,
                octave=role_config.get("octave", 4),
            )
        elif role_name == "perc":
            track["notes"] = generate_latin_perc(bars)

        composition["tracks"].append(track)

    return composition
