# 
# Simple FLAC encoder (Python)
# 
# Copyright (c) 2026 Project Nayuki. (MIT License)
# https://www.nayuki.io/page/simple-flac-implementation
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
# the Software, and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
# - The above copyright notice and this permission notice shall be included in
#   all copies or substantial portions of the Software.
# - The Software is provided "as is", without warranty of any kind, express or
#   implied, including but not limited to the warranties of merchantability,
#   fitness for a particular purpose and noninfringement. In no event shall the
#   authors or copyright holders be liable for any claim, damages or other
#   liability, whether in an action of contract, tort or otherwise, arising from,
#   out of or in connection with the Software or the use or other dealings in the
#   Software.
# 

import pathlib, sys
from types import TracebackType
from typing import BinaryIO


def main(argv: list[str]) -> None:
	if len(argv) != 3:
		sys.exit(f"Usage: python {argv[0]} InFile.wav OutFile.flac")
	with pathlib.Path(argv[1]).open("rb") as inp, \
			BitOutputStream(pathlib.Path(argv[2]).open("wb")) as out:
		encode_file(inp, out)


def encode_file(inp: BinaryIO, out: BitOutputStream) -> None:
	# Read and parse WAV file headers
	def fail_if(cond: bool, msg: str) -> None:
		if cond:
			raise ValueError(msg)
	fail_if(read_fully(inp, 4) != b"RIFF", "Invalid RIFF file header")
	read_little_int(inp, 4)
	fail_if(read_fully(inp, 4) != b"WAVE", "Invalid WAV file header")
	fail_if(read_fully(inp, 4) != b"fmt ", "Unrecognized WAV file chunk")
	fail_if(read_little_int(inp, 4) != 16, "Unsupported WAV file type")
	fail_if(read_little_int(inp, 2) != 0x0001, "Unsupported WAV file codec")
	numchannels: int = read_little_int(inp, 2)
	fail_if(not (1 <= numchannels <= 8), "Too many (or few) audio channels")
	samplerate: int = read_little_int(inp, 4)
	fail_if(not (1 <= samplerate < (1 << 20)), "Sample rate too large or invalid")
	read_little_int(inp, 4)
	read_little_int(inp, 2)
	sampledepth: int = read_little_int(inp, 2)
	fail_if(sampledepth not in (8,16,24,32), "Unsupported sample depth")
	fail_if(read_fully(inp, 4) != b"data", "Unrecognized WAV file chunk")
	sampledatalen: int = read_little_int(inp, 4)
	fail_if((sampledatalen <= 0) or (sampledatalen % (numchannels * (sampledepth // 8)) != 0), "Invalid length of audio sample data")
	
	# Start writing FLAC file header and stream info metadata block
	out.write_int(32, 0x664C6143)
	out.write_int(1, 1)
	out.write_int(7, 0)
	out.write_int(24, 34)
	out.write_int(16, BLOCK_SIZE)
	out.write_int(16, BLOCK_SIZE)
	out.write_int(24, 0)
	out.write_int(24, 0)
	out.write_int(20, samplerate)
	out.write_int(3, numchannels - 1)
	out.write_int(5, sampledepth - 1)
	numsamples: int = sampledatalen // (numchannels * (sampledepth // 8))
	out.write_int(36, numsamples)
	for _ in range(16):
		out.write_int(8, 0)
	
	# Read raw samples and encode FLAC audio frames
	i: int = 0
	while numsamples > 0:
		blocksize: int = min(numsamples, BLOCK_SIZE)
		encode_frame(inp, i, numchannels, sampledepth, samplerate, blocksize, out)
		numsamples -= blocksize
		i += 1

BLOCK_SIZE: int = 4096


def read_fully(inp: BinaryIO, n: int) -> bytes:
	result: bytes = inp.read(n)
	if len(result) < n:
		raise EOFError()
	return result


def read_little_int(inp: BinaryIO, n: int) -> int:
	result: int = 0
	for (i, b) in enumerate(read_fully(inp, n)):
		result |= b << (i * 8)
	return result


def encode_frame(inp: BinaryIO, frameindex: int, numchannels: int, sampledepth: int, samplerate: int, blocksize: int, out: BitOutputStream) -> None:
	bytespersample: int = sampledepth // 8
	samples: list[list[int]] = [[] for _ in range(numchannels)]
	for _ in range(blocksize):
		for chansamples in samples:
			val: int = read_little_int(inp, bytespersample)
			if sampledepth == 8:
				val -= 128
			else:
				val -= (val >> (sampledepth - 1)) << sampledepth
			chansamples.append(val)
	
	out.reset_crcs()
	out.write_int(14, 0x3FFE)
	out.write_int(1, 0)
	out.write_int(1, 0)
	out.write_int(4, 7)
	out.write_int(4, (14 if (samplerate % 10 == 0) else 13))
	out.write_int(4, numchannels - 1)
	out.write_int(3, {8:1, 16:4, 24:6, 32:0}[sampledepth])
	out.write_int(1, 0)
	out.write_int(8, 0xFC | (frameindex >> 30))
	for i in range(24, -1, -6):
		out.write_int(8, 0x80 | ((frameindex >> i) & 0x3F))
	out.write_int(16, blocksize - 1)
	out.write_int(16, samplerate // (10 if (samplerate % 10 == 0) else 1))
	out.write_int(8, out.crc8)
	
	for chansamples in samples:
		encode_subframe(chansamples, sampledepth, out)
	out.align_to_byte()
	out.write_int(16, out.crc16)


def encode_subframe(samples: list[int], sampledepth: int, out: BitOutputStream) -> None:
	out.write_int(1, 0)
	out.write_int(6, 1)  # Verbatim coding
	out.write_int(1, 0)
	for x in samples:
		out.write_int(sampledepth, x)



class BitOutputStream:
	
	out: BinaryIO
	bitbuffer: int
	bitbufferlen: int
	crc8: int
	crc16: int
	
	
	def __init__(self, out: BinaryIO):
		self.out = out
		self.bitbuffer = 0
		self.bitbufferlen = 0
		self.reset_crcs()
	
	
	def reset_crcs(self) -> None:
		self.crc8 = 0
		self.crc16 = 0
	
	
	def align_to_byte(self) -> None:
		self.write_int((-self.bitbufferlen) % 8, 0)
	
	
	def write_int(self, n: int, val: int) -> None:
		self.bitbuffer <<= n
		self.bitbuffer |= val & ((1 << n) - 1)
		self.bitbufferlen += n
		while self.bitbufferlen >= 8:
			self.bitbufferlen -= 8
			b: int = (self.bitbuffer >> self.bitbufferlen) & 0xFF
			self.out.write(bytes((b,)))
			self.crc8 = CRC8_TABLE[self.crc8 ^ b]
			self.crc16 = CRC16_TABLE[(self.crc16 >> 8) ^ b] ^ ((self.crc16 & 0xFF) << 8)
		self.bitbuffer &= (1 << self.bitbufferlen) - 1
	
	
	def close(self) -> None:
		self.out.close()
	
	
	def __enter__(self) -> BitOutputStream:
		return self
	
	
	def __exit__(self, type: type[BaseException]|None, value: BaseException|None, traceback: TracebackType|None) -> None:
		self.close()


CRC8_TABLE: list[int] = []
CRC16_TABLE: list[int] = []
for i in range(256):
	temp8: int = i
	temp16: int = i << 8
	for _ in range(8):
		temp8 = (temp8 << 1) ^ ((temp8 >> 7) * 0x107)
		temp16 = (temp16 << 1) ^ ((temp16 >> 15) * 0x18005)
	CRC8_TABLE.append(temp8)
	CRC16_TABLE.append(temp16)



if __name__ == "__main__":
	main(sys.argv)