"""muse instrumentation — MIDI channel and note-range map for a track.

Shows which MIDI channels carry notes, the pitch range each channel spans,
velocity statistics per channel, and the approximate register (bass/mid/treble).
Agents handling multi-channel orchestration use this to verify that instrument
assignments are coherent before committing.

Usage::

    muse instrumentation tracks/full_score.mid
    muse instrumentation tracks/orchestra.mid --commit HEAD~3
    muse instrumentation tracks/ensemble.mid --json

Output::

    Instrumentation map: tracks/full_score.mid — working tree
    Channels: 4  ·  Total notes: 128

    Ch   Notes  Range        Register   Mean vel
    ───────────────────────────────────────────────
     0      32  C2–G2        bass         78.4
     1      40  C3–C5        mid          72.1
     2      28  G4–E6        treble       65.3
     3      28  F#3–D5       mid          80.0
"""

from __future__ import annotations

import argparse
import json
import logging
import pathlib
import sys
from collections import defaultdict
from typing import TypedDict

from muse.core.errors import ExitCode
from muse.core.repo import read_repo_id, require_repo
from muse.core.store import read_current_branch, resolve_commit_ref
from muse.plugins.midi._query import NoteInfo, load_track, load_track_from_workdir
from muse.plugins.midi.midi_diff import _pitch_name

logger = logging.getLogger(__name__)


class ChannelInfo(TypedDict):
    """Statistics for one MIDI channel."""

    channel: int
    note_count: int
    pitch_min: int
    pitch_max: int
    pitch_min_name: str
    pitch_max_name: str
    register: str
    mean_velocity: float


def _register(pitch_min: int, pitch_max: int) -> str:
    mid = (pitch_min + pitch_max) / 2
    if mid < 48:
        return "bass"
    if mid < 72:
        return "mid"
    return "treble"


def _channel_info(channel: int, notes: list[NoteInfo]) -> ChannelInfo:
    pitches = [n.pitch for n in notes]
    vels = [n.velocity for n in notes]
    lo, hi = min(pitches), max(pitches)
    return ChannelInfo(
        channel=channel,
        note_count=len(notes),
        pitch_min=lo,
        pitch_max=hi,
        pitch_min_name=_pitch_name(lo),
        pitch_max_name=_pitch_name(hi),
        register=_register(lo, hi),
        mean_velocity=round(sum(vels) / len(vels), 1),
    )



def _read_branch(root: pathlib.Path) -> str:
    return read_current_branch(root)


def register(subparsers: "argparse._SubParsersAction[argparse.ArgumentParser]") -> None:
    """Register the instrumentation subcommand."""
    parser = subparsers.add_parser("instrumentation", help="Show per-channel note distribution, pitch range, and register.", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("track", metavar="TRACK", help="Workspace-relative path to a .mid file.")
    parser.add_argument("--commit", "-c", metavar="REF", default=None, dest="ref", help="Analyse a historical snapshot instead of the working tree.")
    parser.add_argument("--json", action="store_true", dest="as_json", help="Emit results as JSON.")
    parser.set_defaults(func=run)


def run(args: argparse.Namespace) -> None:
    """Show per-channel note distribution, pitch range, and register.

    ``muse instrumentation`` groups notes by MIDI channel and reports:
    note count, lowest/highest pitch, register classification, and mean
    velocity.  Use it to verify that instrument roles are coherent — that
    the bass channel stays low, that the melody channel occupies the right
    register, and that no channel is accidentally silent.

    For agents coordinating multi-channel scores, this is the fast sanity
    check before every commit: ``muse instrumentation tracks/score.mid``.
    """
    track: str = args.track
    ref: str | None = args.ref
    as_json: bool = args.as_json

    root = require_repo()
    commit_label = "working tree"

    if ref is not None:
        repo_id = read_repo_id(root)
        branch = _read_branch(root)
        commit = resolve_commit_ref(root, repo_id, branch, ref)
        if commit is None:
            print(f"❌ Commit '{ref}' not found.", file=sys.stderr)
            raise SystemExit(ExitCode.USER_ERROR)
        result = load_track(root, commit.commit_id, track)
        commit_label = commit.commit_id[:8]
    else:
        result = load_track_from_workdir(root, track)

    if result is None:
        print(f"❌ Track '{track}' not found or not a valid MIDI file.", file=sys.stderr)
        raise SystemExit(ExitCode.USER_ERROR)

    notes, _tpb = result
    if not notes:
        print(f"  (no notes found in '{track}')")
        return

    by_channel: dict[int, list[NoteInfo]] = defaultdict(list)
    for n in notes:
        by_channel[n.channel].append(n)

    channels = [_channel_info(ch, ch_notes) for ch, ch_notes in sorted(by_channel.items())]

    if as_json:
        print(json.dumps(
            {"track": track, "commit": commit_label, "channels": list(channels)},
            indent=2,
        ))
        return

    print(f"\nInstrumentation map: {track} — {commit_label}")
    print(f"Channels: {len(channels)}  ·  Total notes: {len(notes)}\n")
    print(f"  {'Ch':>3}  {'Notes':>6}  {'Range':<14}  {'Register':<10}  {'Mean vel':>8}")
    print("  " + "─" * 50)
    for ch in channels:
        rng = f"{ch['pitch_min_name']}–{ch['pitch_max_name']}"
        print(
            f"  {ch['channel']:>3}  {ch['note_count']:>6}  {rng:<14}  "
            f"{ch['register']:<10}  {ch['mean_velocity']:>8.1f}"
        )