from __future__ import annotations

import argparse
import json
from pathlib import Path


BITRATES = {
    (1, 1): [0, 32, 40, 48, 56, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320, 0],
    (1, 2): [0, 32, 48, 56, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320, 384, 0],
    (1, 3): [0, 32, 40, 48, 56, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320, 0],
    (2, 1): [0, 32, 48, 56, 64, 80, 96, 112, 128, 144, 160, 176, 192, 224, 256, 0],
    (2, 2): [0, 8, 16, 24, 32, 40, 48, 56, 64, 80, 96, 112, 128, 144, 160, 0],
    (2, 3): [0, 8, 16, 24, 32, 40, 48, 56, 64, 80, 96, 112, 128, 144, 160, 0],
}
SAMPLE_RATES = {1: [44100, 48000, 32000], 2: [22050, 24000, 16000], 25: [11025, 12000, 8000]}


def synchsafe(b: bytes) -> int:
    return (b[0] << 21) | (b[1] << 14) | (b[2] << 7) | b[3]


def analyze(path: Path) -> dict:
    data = path.read_bytes()
    pos = 0
    if data[:3] == b"ID3" and len(data) >= 10:
        pos = 10 + synchsafe(data[6:10])
        if data[5] & 0x10:
            pos += 10

    frames = 0
    duration = 0.0
    bitrates = []
    sample_rates = []
    channel_modes = []
    first_frame = None
    while pos + 4 <= len(data):
        h = int.from_bytes(data[pos : pos + 4], "big")
        if (h >> 21) & 0x7FF != 0x7FF:
            pos += 1
            continue
        version_bits = (h >> 19) & 0x3
        layer_bits = (h >> 17) & 0x3
        bitrate_idx = (h >> 12) & 0xF
        sr_idx = (h >> 10) & 0x3
        padding = (h >> 9) & 0x1
        channel_mode = (h >> 6) & 0x3
        if version_bits == 1 or layer_bits == 0 or bitrate_idx in (0, 15) or sr_idx == 3:
            pos += 1
            continue
        version = {3: 1, 2: 2, 0: 25}[version_bits]
        layer = {3: 1, 2: 2, 1: 3}[layer_bits]
        br_version = 1 if version == 1 else 2
        bitrate = BITRATES[(br_version, layer)][bitrate_idx] * 1000
        sr = SAMPLE_RATES[version][sr_idx]
        if layer == 1:
            samples = 384
            frame_len = ((12 * bitrate // sr) + padding) * 4
        elif layer == 3 and version != 1:
            samples = 576
            frame_len = (72 * bitrate // sr) + padding
        else:
            samples = 1152
            frame_len = (144 * bitrate // sr) + padding
        if frame_len <= 4 or pos + frame_len > len(data):
            pos += 1
            continue
        if first_frame is None:
            first_frame = pos
        frames += 1
        duration += samples / sr
        bitrates.append(bitrate)
        sample_rates.append(sr)
        channel_modes.append(channel_mode)
        pos += frame_len

    return {
        "path": str(path),
        "bytes": len(data),
        "frames": frames,
        "duration_seconds": round(duration, 6),
        "duration_hms": f"{int(duration//60):02d}:{duration%60:05.2f}",
        "sample_rate_hz": max(set(sample_rates), key=sample_rates.count) if sample_rates else None,
        "channels": "mono" if channel_modes and all(x == 3 for x in channel_modes) else "stereo",
        "average_bitrate_kbps": round(sum(bitrates) / len(bitrates) / 1000, 2) if bitrates else None,
        "min_bitrate_kbps": min(bitrates) // 1000 if bitrates else None,
        "max_bitrate_kbps": max(bitrates) // 1000 if bitrates else None,
        "vbr": len(set(bitrates)) > 1,
        "first_audio_frame_offset": first_frame,
    }


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("path", type=Path)
    args = parser.parse_args()
    print(json.dumps(analyze(args.path), ensure_ascii=False, indent=2))