// Package audio реализует минимальный набор для аудио-стеганографии в
// CTF: разбор WAV, быстрое преобразование Фурье и построение
// спектрограммы. Только WAV, не MP3/FLAC — и это осознанный выбор, не
// недоделка: сжатие с потерями убивает как раз ту мелкую структуру
// спектра, в которую прячут сообщения, так что задачи на эту тему
// в CTF практически всегда именно в несжатом PCM WAV.
package audio

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

// WAV — результат разбора файла: моно-сигнал, нормализованный в
// диапазон [-1, 1] независимо от исходной битности. Многоканальные
// файлы усредняются в моно при разборе — для спектрограммы порядок
// каналов не важен, а городить в UI выбор канала ради этой конкретной
// задачи не стали.
type WAV struct {
	SampleRate int
	Samples    []float64
}

// LoadWAV разбирает несжатый PCM WAV (RIFF/WAVE, formatTag=1).
func LoadWAV(path string) (*WAV, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("чтение %s: %w", path, err)
	}
	if len(data) < 44 || string(data[0:4]) != "RIFF" || string(data[8:12]) != "WAVE" {
		return nil, fmt.Errorf("не похоже на WAV-файл (нет заголовка RIFF/WAVE)")
	}

	var (
		sampleRate    int
		bitsPerSample int
		channels      int
		audioFormat   uint16
		pcmData       []byte
		foundFmt      bool
		foundData     bool
	)

	pos := 12
	for pos+8 <= len(data) {
		chunkID := string(data[pos : pos+4])
		chunkSize := int(binary.LittleEndian.Uint32(data[pos+4 : pos+8]))
		chunkStart := pos + 8
		if chunkStart+chunkSize > len(data) || chunkSize < 0 {
			chunkSize = len(data) - chunkStart // повреждённый/обрезанный файл — берём что есть, не падаем
		}
		if chunkSize < 0 {
			break
		}

		switch chunkID {
		case "fmt ":
			if chunkSize < 16 {
				return nil, fmt.Errorf("повреждённый чанк fmt")
			}
			audioFormat = binary.LittleEndian.Uint16(data[chunkStart : chunkStart+2])
			channels = int(binary.LittleEndian.Uint16(data[chunkStart+2 : chunkStart+4]))
			sampleRate = int(binary.LittleEndian.Uint32(data[chunkStart+4 : chunkStart+8]))
			bitsPerSample = int(binary.LittleEndian.Uint16(data[chunkStart+14 : chunkStart+16]))
			foundFmt = true
		case "data":
			pcmData = data[chunkStart : chunkStart+chunkSize]
			foundData = true
		}

		pos = chunkStart + chunkSize
		if pos%2 == 1 { // чанки RIFF выровнены по чётной границе
			pos++
		}
	}

	if !foundFmt || !foundData {
		return nil, fmt.Errorf("не найден чанк fmt или data — файл повреждён или это не WAV")
	}
	if audioFormat != 1 {
		return nil, fmt.Errorf("поддерживается только несжатый PCM (formatTag=1), получен %d — если файл сжат (MP3/ADPCM и т.п.), сначала перекодируй в PCM WAV", audioFormat)
	}
	if channels < 1 {
		return nil, fmt.Errorf("некорректное число каналов: %d", channels)
	}

	samples, err := decodePCM(pcmData, bitsPerSample, channels)
	if err != nil {
		return nil, err
	}
	if len(samples) == 0 {
		return nil, fmt.Errorf("в файле нет сэмплов")
	}
	return &WAV{SampleRate: sampleRate, Samples: samples}, nil
}

// decodePCM разбирает сырые PCM-байты в моно float64 [-1, 1], усредняя
// каналы, если их больше одного.
func decodePCM(data []byte, bitsPerSample, channels int) ([]float64, error) {
	bytesPerSample := bitsPerSample / 8
	if bytesPerSample == 0 {
		return nil, fmt.Errorf("неподдерживаемая битность: %d", bitsPerSample)
	}
	frameSize := bytesPerSample * channels
	if frameSize == 0 {
		return nil, fmt.Errorf("некорректный размер фрейма")
	}
	numFrames := len(data) / frameSize

	samples := make([]float64, numFrames)
	for i := 0; i < numFrames; i++ {
		var sum float64
		for c := 0; c < channels; c++ {
			off := i*frameSize + c*bytesPerSample
			sum += decodeSample(data[off:off+bytesPerSample], bitsPerSample)
		}
		samples[i] = sum / float64(channels)
	}
	return samples, nil
}

func decodeSample(b []byte, bitsPerSample int) float64 {
	switch bitsPerSample {
	case 8:
		// 8-битный PCM в WAV — беззнаковый (0..255, тишина = 128).
		return (float64(b[0]) - 128) / 128
	case 16:
		v := int16(binary.LittleEndian.Uint16(b))
		return float64(v) / 32768
	case 24:
		v := int32(b[0]) | int32(b[1])<<8 | int32(b[2])<<16
		if v&0x800000 != 0 {
			v -= 0x1000000 // старший бит 24-битного числа установлен — это отрицательное значение, поправляем через вычитание 2^24
		}
		return float64(v) / 8388608
	case 32:
		v := int32(binary.LittleEndian.Uint32(b))
		return float64(v) / 2147483648
	default:
		return 0
	}
}