package audio

import (
	"encoding/binary"
	"math"
	"os"
)

// Represents the header of a WAV file
type WAVHeader struct {
	RiffMark      [4]byte
	FileSize      uint32
	WaveMark      [4]byte
	FmtMark       [4]byte
	FmtSize       uint32
	FormatType    uint16
	NumChannels   uint16
	SampleRate    uint32
	ByteRate      uint32
	BlockAlign    uint16
	BitsPerSample uint16
	DataMark      [4]byte
	DataSize      uint32
}

// Writes the audio data to a WAV file
func (g *Generator) WriteWAV(filename string, samples []float64) error {
	file, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer file.Close()

	// Ensure the samples are normalized first
	Normalize(samples)

	header := WAVHeader{
		RiffMark:      [4]byte{'R', 'I', 'F', 'F'},
		WaveMark:      [4]byte{'W', 'A', 'V', 'E'},
		FmtMark:       [4]byte{'f', 'm', 't', ' '},
		FmtSize:       16,
		FormatType:    1, // PCM
		NumChannels:   Channels,
		SampleRate:    uint32(g.sampleRate),
		BitsPerSample: Bits,
		ByteRate:      uint32(g.sampleRate * float64(Channels) * float64(Bits) / 8),
		BlockAlign:    uint16(Channels * Bits / 8),
		DataMark:      [4]byte{'d', 'a', 't', 'a'},
	}

	// Calculate sizes - samples are interleaved LRLR...
	header.DataSize = uint32(len(samples) * int(Bits) / 8)
	header.FileSize = 36 + header.DataSize // 36 is the size of the header up to data

	// Write header
	if err := binary.Write(file, binary.LittleEndian, &header); err != nil {
		return err
	}

	// Write audio data - batch convert and write for better performance
	const batchSize = 8192 // Write in 8KB chunks
	int16Buffer := make([]int16, batchSize)

	for i := 0; i < len(samples); i += batchSize {
		end := i + batchSize
		if end > len(samples) {
			end = len(samples)
		}

		// Convert batch of samples to int16
		for j := i; j < end; j++ {
			sample := samples[j]
			// Clamp to prevent overflow
			if sample > 1.0 {
				sample = 1.0
			} else if sample < -1.0 {
				sample = -1.0
			}
			int16Buffer[j-i] = int16(sample * 32767.0)
		}

		// Write entire batch at once
		if err := binary.Write(file, binary.LittleEndian, int16Buffer[:end-i]); err != nil {
			return err
		}
	}

	return nil
}

// Normalizes the audio buffer to prevent clipping
func Normalize(buffer []float64) {
	if len(buffer) == 0 {
		return
	}

	// Find the maximum absolute value in the buffer
	maxVal := 0.0
	for _, sample := range buffer {
		absVal := math.Abs(sample)
		if absVal > maxVal {
			maxVal = absVal
		}
	}

	// If the maximum is very small, don't normalize to avoid amplifying noise
	if maxVal < 0.0001 {
		return
	}

	// Normalize to 90% of maximum to avoid clipping
	scale := 0.9 / maxVal
	for i := range buffer {
		buffer[i] *= scale
	}
}