"""Add an Apple-encoder AAC column to the chart study, measured on macOS.

The main run happens in the container, where ffmpeg has libvorbis. Its AAC
column is ffmpeg's own encoder, which is not what Apple Music is served with,
and saying so in the copy is honest but weaker than measuring the real one.
macOS ships Apple's AAC encoder as afconvert, so this pass runs on the host and
merges a second AAC column into the same JSON.

    scripts/chart_codec_apple.py [--root DIR]

Requires afconvert (macOS), ffmpeg for decoding, and numpy/soundfile. It reads
docs/data/chart_codec_study.json, adds codecs['apple_native'] to every track and
a second toolchain entry, and writes the file back. Run it after the main study,
never before: it will not create rows it does not find.
"""

import argparse
import json
import os
import platform
import shutil
import subprocess
import sys
import tempfile

import numpy as np
import soundfile as sf

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, '..', 'src'))

from codec_null_test import align, band_residual, db, rms, to_wav   # noqa: E402
from chart_codec_study import (                                     # noqa: E402
    BAND_WINDOW_S, DEFAULT_ROOT, MANIFEST, band_window, our_true_peak,
    true_peak_ebur128,
)

DATA = os.path.join(HERE, '..', 'docs', 'data', 'chart_codec_study.json')

# Constrained VBR at 256 kbps with the codec's quality dial at maximum. Apple
# Music is served as 256 kbps VBR AAC; this is that bitrate through Apple's
# encoder rather than a guess at Apple's exact internal settings, and the
# command line goes into the data so the choice is visible.
AFCONVERT = ['-f', 'm4af', '-d', 'aac', '-b', '256000', '-q', '127', '-s', '2']


def run(cmd):
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode != 0:
        raise RuntimeError(' '.join(cmd[:4]) + ': ' + (r.stderr or '')[:300])


def apple_version():
    r = subprocess.run(['afconvert', '--help'], capture_output=True, text=True)
    first = [ln.strip() for ln in (r.stdout or '').splitlines() if ln.strip()]
    return {
        'tool': 'afconvert (Apple CoreAudio)',
        'version': next((ln for ln in first if ln.startswith('Version')), ''),
        'host': platform.platform(),
    }


def measure(src_path, tmp):
    ref_path = os.path.join(tmp, 'ref.wav')
    enc = os.path.join(tmp, 'a.m4a')
    dec = os.path.join(tmp, 'a.wav')
    for p in (ref_path, enc, dec):
        if os.path.exists(p):
            os.remove(p)

    to_wav(src_path, ref_path, 44100)
    cmd = ['afconvert'] + AFCONVERT + [ref_path, enc]
    run(cmd)
    to_wav(enc, dec, 44100)

    ref, _ = sf.read(ref_path, always_2d=True, dtype='float64')
    ref = ref[:, 0]
    test, _ = sf.read(dec, always_2d=True, dtype='float64')
    test = test[:, 0]

    lag = align(ref, test)
    if lag >= 0:
        a, b = ref[:len(ref) - lag], test[lag:lag + len(ref) - lag]
    else:
        a, b = ref[-lag:], test[:len(ref) + lag]
    n = min(len(a), len(b))
    a, b = a[:n], b[:n]

    g = float(np.dot(a, b) / np.dot(b, b)) if np.dot(b, b) else 1.0
    diff = a - b * g
    residual = db(rms(diff) / rms(a))
    lo, hi = band_window(a, 44100)
    tp, sample_peak = our_true_peak(dec)
    return {
        'label': 'AAC-LC 256 kbps (Apple encoder)',
        'lag_samples': int(lag),
        'gain_correction_db': db(abs(g)),
        'residual_db': None if residual == -np.inf else round(residual, 2),
        'bands': [{'lo': l_, 'hi': h_,
                   'rel_db': None if r == -np.inf else round(r, 1)}
                  for (l_, h_), r in band_residual(a[lo:hi], diff[lo:hi], 44100)],
        'true_peak_4x': tp,
        'sample_peak': sample_peak,
        'true_peak_ebur128': true_peak_ebur128(dec),
        'encoded_bytes': os.path.getsize(enc),
        'encoder_input_rate': 44100,
        'command': ' '.join(['afconvert'] + AFCONVERT + ['in.wav', 'out.m4a']),
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--root', default=DEFAULT_ROOT)
    args = ap.parse_args()
    for tool in ('afconvert', 'ffmpeg'):
        if not shutil.which(tool):
            sys.exit('%s not found; this pass only runs on macOS' % tool)

    d = json.load(open(DATA, encoding='utf-8'))
    by_key = {t['key']: t for t in d['tracks']}
    paths = {k: os.path.join(args.root, rel) for k, _a, _t, rel in MANIFEST}

    tmp = tempfile.mkdtemp(prefix='apple-aac-')
    for key, track in by_key.items():
        src = paths.get(key)
        if not src or not os.path.exists(src):
            sys.exit('missing audio for %s' % key)
        r = measure(src, tmp)
        track['codecs']['apple_native'] = r
        ffm = track['codecs']['apple']
        print('%-22s ffmpeg %6.1f dB / TP %+.2f    Apple %6.1f dB / TP %+.2f'
              % (track['artist'], ffm['residual_db'], ffm['true_peak_4x'],
                 r['residual_db'], r['true_peak_4x']))

    d.setdefault('toolchain', {})['apple'] = apple_version()
    with open(DATA, 'w', encoding='utf-8') as f:
        json.dump(d, f, indent=2, ensure_ascii=False)
    print('\nmerged apple_native into %s' % DATA)


if __name__ == '__main__':
    main()