// Package watermark implements STFT-domain spread-spectrum audio watermarking
// for fm-rds-tx, based on Kirovski & Malvar (IEEE TSP 2003).
//
// The watermark is embedded in STFT magnitude (dB scale) — a multiplicative
// modification that naturally scales with audio content (psychoacoustic masking).
// Block repetition coding provides drift tolerance. Reed-Solomon RS(16,8)
// provides error correction for the 128-bit payload.
//
// Encoder: internal/watermark/stft_watermark.go (STFTEmbedder)
// Decoder: cmd/wmdecode (STFTDetector)
package watermark

import (
	"crypto/sha256"
	"fmt"
)

// RS code parameters.
const (
	rsDataBytes  = 8                         // payload bytes
	rsCheckBytes = 8                         // parity bytes
	rsTotalBytes = rsDataBytes + rsCheckBytes // 16
	payloadBits  = rsTotalBytes * 8          // 128
)

// Exported constants for tools.
const (
	PayloadBits  = payloadBits
	RsDataBytes  = rsDataBytes
	RsTotalBytes = rsTotalBytes
	RsCheckBytes = rsCheckBytes
)

// --- GF(2^8) arithmetic — primitive polynomial 0x11d ---

func gfMul(a, b byte) byte {
	if a == 0 || b == 0 {
		return 0
	}
	return gfExp[(int(gfLog[a])+int(gfLog[b]))%255]
}

func gfInv(a byte) byte {
	if a == 0 {
		return 0
	}
	return gfExp[255-int(gfLog[a])]
}

func gfPow(a byte, n int) byte {
	if n == 0 {
		return 1
	}
	r := a
	for i := 1; i < n; i++ {
		r = gfMul(r, a)
	}
	return r
}

// --- RS(16,8) encoding ---

func rsEncode(data [rsDataBytes]byte) [rsTotalBytes]byte {
	work := make([]byte, rsTotalBytes)
	copy(work, data[:])
	for i := 0; i < rsDataBytes; i++ {
		fb := work[i]
		if fb != 0 {
			for j := 1; j <= rsCheckBytes; j++ {
				work[i+j] ^= gfMul(rsGen[j], fb)
			}
		}
	}
	var out [rsTotalBytes]byte
	copy(out[:rsDataBytes], data[:])
	copy(out[rsDataBytes:], work[rsDataBytes:])
	return out
}

// RSEncode encodes 8 data bytes into a 16-byte RS codeword (exported).
func RSEncode(data [rsDataBytes]byte) [rsTotalBytes]byte {
	return rsEncode(data)
}

// --- RS(16,8) decoding (Vandermonde solver) ---

// RSDecode recovers 8 data bytes from a (possibly corrupted) 16-byte codeword.
// erasurePositions lists the byte indices (0-15) of known-erased bytes.
// Returns the 8 payload bytes and true if decoding succeeded.
//
// The polynomial convention is C(x) = c[0]x^15 + c[1]x^14 + … + c[15],
// so byte position j maps to polynomial power (15-j). The locator for
// position j is α^(15-j). Decoding uses Vandermonde/Gaussian elimination
// instead of Forney's formula to avoid position-mapping bugs.
func RSDecode(recv [rsTotalBytes]byte, erasurePositions []int) ([rsDataBytes]byte, bool) {
	// Syndromes S[i] = C(α^i) via Horner's method.
	var S [rsCheckBytes]byte
	for i := 0; i < rsCheckBytes; i++ {
		var acc byte
		for _, c := range recv {
			acc = gfMul(acc, gfPow(2, i)) ^ c
		}
		S[i] = acc
	}

	// All syndromes zero → valid codeword.
	hasErr := false
	for _, s := range S {
		if s != 0 {
			hasErr = true
			break
		}
	}
	if !hasErr {
		var out [rsDataBytes]byte
		copy(out[:], recv[:rsDataBytes])
		return out, true
	}
	ne := len(erasurePositions)
	if ne == 0 || ne > rsCheckBytes {
		return [rsDataBytes]byte{}, false
	}

	// Locators: X[k] = α^(15-pos) for byte position pos.
	X := make([]byte, ne)
	for k, pos := range erasurePositions {
		X[k] = gfPow(2, rsTotalBytes-1-pos)
	}

	// Vandermonde system: S[i] = Σ_k e_k · X[k]^i
	// Augmented matrix [V | S], ne × (ne+1)
	mat := make([][]byte, ne)
	for i := range mat {
		mat[i] = make([]byte, ne+1)
		for k := 0; k < ne; k++ {
			mat[i][k] = gfPow(X[k], i)
		}
		mat[i][ne] = S[i]
	}

	// Gaussian elimination
	for col := 0; col < ne; col++ {
		pivot := -1
		for row := col; row < ne; row++ {
			if mat[row][col] != 0 {
				pivot = row
				break
			}
		}
		if pivot < 0 {
			return [rsDataBytes]byte{}, false
		}
		mat[col], mat[pivot] = mat[pivot], mat[col]
		inv := gfInv(mat[col][col])
		for j := 0; j <= ne; j++ {
			mat[col][j] = gfMul(mat[col][j], inv)
		}
		for row := 0; row < ne; row++ {
			if row == col {
				continue
			}
			f := mat[row][col]
			if f == 0 {
				continue
			}
			for j := 0; j <= ne; j++ {
				mat[row][j] ^= gfMul(f, mat[col][j])
			}
		}
	}

	// Apply corrections
	result := recv
	for k := 0; k < ne; k++ {
		result[erasurePositions[k]] ^= mat[k][ne]
	}

	// Verify
	for i := 0; i < rsCheckBytes; i++ {
		var acc byte
		for _, c := range result {
			acc = gfMul(acc, gfPow(2, i)) ^ c
		}
		if acc != 0 {
			return [rsDataBytes]byte{}, false
		}
	}

	var out [rsDataBytes]byte
	copy(out[:], result[:rsDataBytes])
	return out, true
}

// --- Key utilities ---

// KeyMatchesPayload returns true if SHA-256(key)[:8] matches payload.
func KeyMatchesPayload(key string, payload [rsDataBytes]byte) bool {
	h := sha256.Sum256([]byte(key))
	var expected [rsDataBytes]byte
	copy(expected[:], h[:rsDataBytes])
	return expected == payload
}

// KeyToPayload returns SHA-256(key)[:8].
func KeyToPayload(key string) [rsDataBytes]byte {
	var data [rsDataBytes]byte
	h := sha256.Sum256([]byte(key))
	copy(data[:], h[:rsDataBytes])
	return data
}

// PayloadHex returns the hex string of the RS-encoded payload for a key.
func PayloadHex(key string) string {
	data := KeyToPayload(key)
	cw := rsEncode(data)
	return fmt.Sprintf("%x", cw)
}

// --- STFT detector accessors ---

// PNChipAt returns the PN chip value at group g, bin b.
func (d *STFTDetector) PNChipAt(g, b int) int8 {
	return d.pnChips[g][b]
}

// GroupBit returns the data bit index for group g.
func (d *STFTDetector) GroupBit(g int) int {
	return d.groupToBit[g]
}

// --- GF tables ---

var gfExp = [512]byte{1, 2, 4, 8, 16, 32, 64, 128, 29, 58, 116, 232, 205, 135, 19, 38, 76, 152, 45, 90, 180, 117, 234, 201, 143, 3, 6, 12, 24, 48, 96, 192, 157, 39, 78, 156, 37, 74, 148, 53, 106, 212, 181, 119, 238, 193, 159, 35, 70, 140, 5, 10, 20, 40, 80, 160, 93, 186, 105, 210, 185, 111, 222, 161, 95, 190, 97, 194, 153, 47, 94, 188, 101, 202, 137, 15, 30, 60, 120, 240, 253, 231, 211, 187, 107, 214, 177, 127, 254, 225, 223, 163, 91, 182, 113, 226, 217, 175, 67, 134, 17, 34, 68, 136, 13, 26, 52, 104, 208, 189, 103, 206, 129, 31, 62, 124, 248, 237, 199, 147, 59, 118, 236, 197, 151, 51, 102, 204, 133, 23, 46, 92, 184, 109, 218, 169, 79, 158, 33, 66, 132, 21, 42, 84, 168, 77, 154, 41, 82, 164, 85, 170, 73, 146, 57, 114, 228, 213, 183, 115, 230, 209, 191, 99, 198, 145, 63, 126, 252, 229, 215, 179, 123, 246, 241, 255, 227, 219, 171, 75, 150, 49, 98, 196, 149, 55, 110, 220, 165, 87, 174, 65, 130, 25, 50, 100, 200, 141, 7, 14, 28, 56, 112, 224, 221, 167, 83, 166, 81, 162, 89, 178, 121, 242, 249, 239, 195, 155, 43, 86, 172, 69, 138, 9, 18, 36, 72, 144, 61, 122, 244, 245, 247, 243, 251, 235, 203, 139, 11, 22, 44, 88, 176, 125, 250, 233, 207, 131, 27, 54, 108, 216, 173, 71, 142, 1, 2, 4, 8, 16, 32, 64, 128, 29, 58, 116, 232, 205, 135, 19, 38, 76, 152, 45, 90, 180, 117, 234, 201, 143, 3, 6, 12, 24, 48, 96, 192, 157, 39, 78, 156, 37, 74, 148, 53, 106, 212, 181, 119, 238, 193, 159, 35, 70, 140, 5, 10, 20, 40, 80, 160, 93, 186, 105, 210, 185, 111, 222, 161, 95, 190, 97, 194, 153, 47, 94, 188, 101, 202, 137, 15, 30, 60, 120, 240, 253, 231, 211, 187, 107, 214, 177, 127, 254, 225, 223, 163, 91, 182, 113, 226, 217, 175, 67, 134, 17, 34, 68, 136, 13, 26, 52, 104, 208, 189, 103, 206, 129, 31, 62, 124, 248, 237, 199, 147, 59, 118, 236, 197, 151, 51, 102, 204, 133, 23, 46, 92, 184, 109, 218, 169, 79, 158, 33, 66, 132, 21, 42, 84, 168, 77, 154, 41, 82, 164, 85, 170, 73, 146, 57, 114, 228, 213, 183, 115, 230, 209, 191, 99, 198, 145, 63, 126, 252, 229, 215, 179, 123, 246, 241, 255, 227, 219, 171, 75, 150, 49, 98, 196, 149, 55, 110, 220, 165, 87, 174, 65, 130, 25, 50, 100, 200, 141, 7, 14, 28, 56, 112, 224, 221, 167, 83, 166, 81, 162, 89, 178, 121, 242, 249, 239, 195, 155, 43, 86, 172, 69, 138, 9, 18, 36, 72, 144, 61, 122, 244, 245, 247, 243, 251, 235, 203, 139, 11, 22, 44, 88, 176, 125, 250, 233, 207, 131, 27, 54, 108, 216, 173, 71, 142, 1, 2}
var gfLog = [256]byte{0, 0, 1, 25, 2, 50, 26, 198, 3, 223, 51, 238, 27, 104, 199, 75, 4, 100, 224, 14, 52, 141, 239, 129, 28, 193, 105, 248, 200, 8, 76, 113, 5, 138, 101, 47, 225, 36, 15, 33, 53, 147, 142, 218, 240, 18, 130, 69, 29, 181, 194, 125, 106, 39, 249, 185, 201, 154, 9, 120, 77, 228, 114, 166, 6, 191, 139, 98, 102, 221, 48, 253, 226, 152, 37, 179, 16, 145, 34, 136, 54, 208, 148, 206, 143, 150, 219, 189, 241, 210, 19, 92, 131, 56, 70, 64, 30, 66, 182, 163, 195, 72, 126, 110, 107, 58, 40, 84, 250, 133, 186, 61, 202, 94, 155, 159, 10, 21, 121, 43, 78, 212, 229, 172, 115, 243, 167, 87, 7, 112, 192, 247, 140, 128, 99, 13, 103, 74, 222, 237, 49, 197, 254, 24, 227, 165, 153, 119, 38, 184, 180, 124, 17, 68, 146, 217, 35, 32, 137, 46, 55, 63, 209, 91, 149, 188, 207, 205, 144, 135, 151, 178, 220, 252, 190, 97, 242, 86, 211, 171, 20, 42, 93, 158, 132, 60, 57, 83, 71, 109, 65, 162, 31, 45, 67, 216, 183, 123, 164, 118, 196, 23, 73, 236, 127, 12, 111, 246, 108, 161, 59, 82, 41, 157, 85, 170, 251, 96, 134, 177, 187, 204, 62, 90, 203, 89, 95, 176, 156, 169, 160, 81, 11, 245, 22, 235, 122, 117, 44, 215, 79, 174, 213, 233, 230, 231, 173, 232, 116, 214, 244, 234, 168, 80, 88, 175}
var rsGen = [9]byte{1, 255, 11, 81, 54, 239, 173, 200, 24}