#!/usr/bin/env python3
"""AI ATLAS: inspect and scale short mono16-bit PCM WAVs, offline.
Python3.9+ standard library only. No uploads, audio playback, or overwrites.
Example:
  python audio_lab.py measure jingle-abrupt.wav jingle-faded.wav
  python audio_lab.py normalize jingle-faded.wav peak.wav --mode peak --target-dbfs -3
  python audio_lab.py normalize jingle-faded.wav rms.wav --mode rms --target-dbfs -18
Metrics are not LUFS, true peak, hearing safety, or a subjective quality score.
"""
from array import array
import argparse
import json
import math
from pathlib import Path
import sys
import wave

FULL_SCALE = 32767.0  # Positive16-bit full-scale reference, also used by the guide's waveform.
MAX_FRAMES = 5_000_000  # This small educational tool is not a streaming editor.

def load_pcm(path):
    with wave.open(str(path), 'rb') as source:
        channels, width, rate, frames, compression, _ = source.getparams()
        if (channels,width,compression) != (1,2,'NONE'):
            raise ValueError('Expected mono16-bit uncompressed PCM WAV; convert a copy first.')
        if not 0 < frames <= MAX_FRAMES or rate <= 0:
            raise ValueError('Expected1 to5,000,000 samples and a positive sample rate.')
        raw = source.readframes(frames)
        if len(raw) != frames*2:
            raise ValueError('Truncated WAV sample data.')
    samples = array('h')
    samples.frombytes(raw)
    if sys.byteorder != 'little':
        samples.byteswap()
    return samples,rate

def db(value):
    return None if value == 0 else 20*math.log10(value)

def measure(samples,rate,boundary_seconds=.5):
    n=len(samples)
    if not n or rate <= 0:
        raise ValueError('Nonempty samples and positive rate are required.')
    if not math.isfinite(boundary_seconds) or boundary_seconds <= 0:
        raise ValueError('Boundary spacing must be finite and positive.')
    step_float=rate*boundary_seconds
    step=round(step_float)
    if step < 1 or abs(step_float-step) > 1e-6:
        raise ValueError('Boundary spacing must correspond to a whole number of samples.')
    peak=max(abs(v) for v in samples)/FULL_SCALE
    rms=math.sqrt(sum(v*v for v in samples)/n)/FULL_SCALE
    boundaries=list(range(step,n,step))
    return {'sampleRate':rate,'channels':1,'bits':16,'samples':n,'seconds':n/rate,
            'samplePeakDbfs':db(peak),'rmsDbfs':db(rms),'meanAmplitude':sum(samples)/n/FULL_SCALE,
            'railSamples':sum(v in (-32768,32767) for v in samples),
            'boundarySpacingSeconds':boundary_seconds,'boundaryCount':len(boundaries),
            'maxBoundaryJump':max((abs(samples[i]-samples[i-1])/FULL_SCALE for i in boundaries),default=0),
            'firstSample':samples[0],'lastSample':samples[-1],
            'amplitudeReference':32767,
            'limits':'RMS is unweighted and includes silence/DC. Rail hits are a warning, not proof of clipping; zero hits do not prove undistorted audio. No LUFS or true-peak calculation.'}

def normalized(samples,rate,mode,target_dbfs):
    if mode not in ('peak','rms'):
        raise ValueError('Mode must be peak or rms.')
    if not math.isfinite(target_dbfs) or not -60 <= target_dbfs <= -0.1:
        raise ValueError('Educational target must be between -60 and -0.1dBFS.')
    before=measure(samples,rate)
    current=before['samplePeakDbfs' if mode=='peak' else 'rmsDbfs']
    if current is None:
        raise ValueError('Silence has no finite gain that reaches this target.')
    gain=10**((target_dbfs-current)/20)
    lo=min(samples)*gain;hi=max(samples)*gain
    if lo < -32768 or hi > 32767:
        raise ValueError('Target would exceed16-bit range. Lower the target; this tool will not clip or compress.')
    result=array('h',(round(value*gain) for value in samples))
    return result,{'mode':mode,'targetDbfs':target_dbfs,'gainDb':20*math.log10(gain),
                   'before':before,'after':measure(result,rate),
                   'processing':'One fixed gain for the full file, then16-bit rounding. No limiter, compressor, DC removal, noise reduction, or dither.'}

def write_pcm_exclusive(path,samples,rate):
    # All calculations/validation happen before creating a new output.
    raw=array('h',samples)
    if sys.byteorder != 'little':
        raw.byteswap()
    with Path(path).open('xb') as destination:
        with wave.open(destination,'wb') as output:
            output.setnchannels(1);output.setsampwidth(2);output.setframerate(rate)
            output.writeframes(raw.tobytes())

def main(argv=None):
    parser=argparse.ArgumentParser(description=__doc__)
    sub=parser.add_subparsers(dest='command',required=True)
    inspect=sub.add_parser('measure',help='Read files and print JSON without modifying them')
    inspect.add_argument('inputs',nargs='+',type=Path)
    inspect.add_argument('--boundary-seconds',type=float,default=.5,help='Known note spacing in this exercise, not automatic note detection')
    scale=sub.add_parser('normalize',help='Create a new scaled WAV; existing paths are never overwritten')
    scale.add_argument('input',type=Path);scale.add_argument('output',type=Path)
    scale.add_argument('--mode',choices=['peak','rms'],required=True)
    scale.add_argument('--target-dbfs',type=float,required=True)
    args=parser.parse_args(argv)
    if args.command=='measure':
        result=[]
        for path in args.inputs:
            samples,rate=load_pcm(path)
            result.append({'file':path.name,**measure(samples,rate,args.boundary_seconds)})
    else:
        if args.input.resolve()==args.output.resolve() or args.output.exists():
            raise ValueError('Choose a new output filename; the source or existing file will not be overwritten.')
        samples,rate=load_pcm(args.input)
        output,result=normalized(samples,rate,args.mode,args.target_dbfs)
        write_pcm_exclusive(args.output,output,rate)
        result={'source':args.input.name,'output':args.output.name,**result}
    print(json.dumps(result,ensure_ascii=False,indent=2,allow_nan=False))

if __name__=='__main__':
    try:main()
    except (ValueError,OSError,wave.Error) as error:
        print('ERROR: '+str(error),file=sys.stderr);sys.exit(1)