import argparse
import numpy as np
import wave


def read_wav_int16(path: str):
    with wave.open(path, "rb") as wav:
        sample_width = wav.getsampwidth()
        if sample_width != 2:
            raise ValueError(f"{path} is {sample_width * 8}-bit PCM; expected 16-bit")
        sample_rate = wav.getframerate()
        channels = wav.getnchannels()
        frame_count = wav.getnframes()
        raw = wav.readframes(frame_count)
    samples = np.frombuffer(raw, dtype="<i2").astype(np.int16, copy=False)
    return samples, sample_rate, channels


def write_wav_int16(path: str, samples: np.ndarray, sample_rate: int, channels: int):
    with wave.open(path, "wb") as wav:
        wav.setnchannels(channels)
        wav.setsampwidth(2)
        wav.setframerate(sample_rate)
        wav.writeframes(np.asarray(samples, dtype="<i2").tobytes())


def main():
    parser = argparse.ArgumentParser(description="Amplify a WAV file.")
    parser.add_argument("input", type=str, help="Input WAV path")
    parser.add_argument("output", type=str, help="Output WAV path")
    parser.add_argument("--gain", type=float, default=5.0, help="Gain factor")
    args = parser.parse_args()

    data, sr, channels = read_wav_int16(args.input)
    print(f"Original max abs: {np.max(np.abs(data))}")

    amplified = np.clip(data.astype(np.float32) * args.gain, -32768, 32767).astype(np.int16)
    print(f"Amplified max abs: {np.max(np.abs(amplified))}")

    write_wav_int16(args.output, amplified, sr, channels)
    print(f"Saved: {args.output}")


if __name__ == "__main__":
    main()