package wav

import "core:fmt"
import "core:math"
import "core:strings"
import "core:strconv"
import "core:os"
import "xml"

/*
TODO: Support RF64
TODO: Support music metadata
*/

Wav :: struct {
	// Basic data
	path          : string,
	format        : Audio_Format,
	channels      : int,
	sample_rate   : int,
	bit_depth     : int,
	sample_size   : int, // Samples can be bigger than bit_depth for alignment
	reported_size : int,
	data_size     : int,
	data_start    : int,
	
	// The actual audio data
	audio : [][]f32,
	
	// Internals
	handle : os.Handle,
	load_head : int,
	
	// Metadata
	date : Date,
	channel_names : []string,
	channel_mask : Speaker_Set,
	samples_since_midnight : int, // Source of timecode
	tc_framerate : f32,
	tc_dropframe : bool,
	ubits : [8]u8,
	take : int,
	scene : string,
	note : string,
	project : string,
	tape : string,
	circled : bool,
}
Audio_Format :: enum {
	INT = 1,
	ADPCM = 2, // Compressed format. Not yet handled.
	FLOAT = 3,
}
Date :: struct {
	year, month, day : int,
}
Timecode :: struct {
	hour, minute, second, frame : int
}
Speakers :: enum u32 {
	Front_Left            = 0,
	Front_Right           = 1,
	Front_Center          = 2,
	LFE                   = 3,
	Back_Left             = 4,
	Back_Right            = 5,
	Front_Left_of_Center  = 6,
	Front_Right_of_Center = 7,
	Back_Center           = 8,
	Side_Left             = 9,
	Side_Right            = 10,
	Top_Center            = 11,
	Top_Front_Left        = 12,
	Top_Front_Center      = 13,
	Top_Front_Right       = 14,
	Top_Back_Left         = 15,
	Top_Back_Center       = 16,
	Top_Back_Right        = 17,
}
Speaker_Set :: bit_set[Speakers; u32]

VERBOSE :: false
BUFFER_SIZE :: 1<<18

main :: proc() {
	// Test
	/*enok, enok_ok := read("test/WAVs/ENOKS-BIRHTDAYT02.WAV", context.temp_allocator)
	when VERBOSE do fmt.printf("\n\nenok = %#v\n\n", enok)
	prins, prins_ok := read("test/WAVs/KRONPRINS01T01.wav", context.temp_allocator)
	when VERBOSE do fmt.printf("\n\nprins = %#v\n\n", prins)
	*/
	f8, f8_ok := read("test/WAVs/F8-SL098-T001.WAV", context.temp_allocator)
	when VERBOSE do fmt.printf("\n\nf8 = %#v\n\n", f8)
	load(&f8)
	wave_print(f8.audio[0])
	
	/*
	ski, ski_ok := read("test/WAVs/FOLEY, SKI, KLAEBO 01, LCR.wav", context.temp_allocator)
	when VERBOSE do fmt.printf("\n\nSKI = %#v\n\n", ski)
	load(&ski)
	wave_print(ski.audio[0])
	
	t, t_ok := read("test/WAVs/test_data.wav", context.temp_allocator)
	when VERBOSE do fmt.printf("\n\nTEST = %#v\n\n", t)
	load(&t, {1}, allocator=context.temp_allocator)
	wave_print(t.audio[1])
	*/
}

/*
Reads in the wav file metadata, without loading the sound data into ram.
*/
read :: proc(path : string, allocator:=context.allocator) -> (Wav, bool) #optional_ok {
	file : Wav
	file.path = path
	file.take = -1
	file.tc_framerate = -1
	file.take = -1
	
	load_err : os.Error
	file.handle, load_err = os.open(path)
	defer os.close(file.handle)
	defer file.handle = 0
	if load_err != os.General_Error.None {
		fmt.eprintfln("ERROR %v: Unable to load file \"%v\"", load_err, path)
		return {}, false
	}
	
	
	temp_buf := new([BUFFER_SIZE]u8)[:]
	temp_bext : []u8
	temp_ixml : string
	defer delete(temp_buf)
	
	os.read(file.handle, temp_buf)
	
	head : int = 0
	
	// RIFF header
	when VERBOSE do fmt.println(string(temp_buf[0:4]))
	if string(temp_buf[0:4]) != "RIFF" do return {}, false
	head += 4
	
	// Size
	file.reported_size = int(little_endian_u32(temp_buf[head:head+4]))
	when VERBOSE do fmt.println("Reported size:", file.reported_size)
	head += 4
	
	// Confirming again that this is a wave file
	when VERBOSE do fmt.println(string(temp_buf[head:head+4]))
	if string(temp_buf[head:head+4]) != "WAVE" do return {}, false
	head += 4
	
	when VERBOSE do fmt.println("\nChunks:\n")
	
	// Looping through chunks
	null_chunks := 0
	for _ in 0..<BUFFER_SIZE {
		chunk_id_raw := temp_buf[head:head+4]
		chunk_id := string(chunk_id_raw)
		head += 4
		chunk_size := int(little_endian_u32(temp_buf[head:head+4]))
		head += 4
		when VERBOSE do fmt.println(chunk_id, chunk_size,"\n-------------------------------------")
		current_chunk_start := head
		next_chunk_start := head + chunk_size
		
		data_reached := false
		if chunk_id == "data" {
			file.data_size = chunk_size
			file.data_start = head
			data_reached = true
		}
		if next_chunk_start < BUFFER_SIZE && !data_reached {
			print_data : string
			data_end := head+chunk_size
			read_end := min(head+15000, data_end)
			switch chunk_id {
				case "iXML":
					temp_ixml = string(temp_buf[head:data_end])
					//print_data = temp_ixml
					null_chunks = 0
				case "bext":
					temp_bext = temp_buf[head:data_end]
					//print_data = string(temp_bext)
					null_chunks = 0
				case "fmt ":
					raw_format_field := little_endian_u16(temp_buf[head:])
					wave_format_extensible := false
					switch raw_format_field {
						case 1..=3:
							file.format = Audio_Format(raw_format_field)
							when VERBOSE do fmt.println("Format:", file.format)
						case 65534:
							wave_format_extensible = true
							when VERBOSE do fmt.println("Format: wave_format_extensible")
					}
					head += 2
					file.channels = int(little_endian_u16(temp_buf[head:]))
					when VERBOSE do fmt.println("Channels:", file.channels)
					head += 2
					file.sample_rate = int(little_endian_u32(temp_buf[head:]))
					when VERBOSE do fmt.println("Sample rate:", file.sample_rate)
					head += 4
					
					// Skipping byte rate and block align.
					// These two legacy fields are only ever redundant or wrong,
					// and are implied by the other fields.
					head += 4 + 2
					
					if !wave_format_extensible {
						file.bit_depth = int(little_endian_u16(temp_buf[head:]))
						file.sample_size = file.bit_depth
						when VERBOSE do fmt.println("Bit depth:", file.bit_depth)
						head += 2
					} else if chunk_size>=40 {
						
						// This is what would normally be the bit depth field
						// But in the wave_format_extensible case, it essentailly tells you the
						// stride between each sample, aka "sample size" in this lib's API.
						bits_pr_sample := int(little_endian_u16(temp_buf[head:]))
						head += 2
						
						when VERBOSE do fmt.println(" - wave_format_extensible -")
						when VERBOSE do fmt.println("Size of extension:", little_endian_u16(temp_buf[head:]))
						head += 2
						
						// Valid bits pr sample (aka is every 24-bit sample padded to be 32 bit aligned?)
						// OR this field can be samples pr block, if this is a compressed format...
						// ...sigh... yes, you can cram a compressed format into a wav file. I know.
						valid_bits_pr_sample := int(little_endian_u16(temp_buf[head:]))
						head += 2
						
						// Channel mask (Which speaker channels are in the file)
						file.channel_mask = transmute(Speaker_Set)little_endian_u32(temp_buf[head:])
						when VERBOSE do fmt.println("Channel Mask (RAW):", temp_buf[head:head+4])
						when VERBOSE do fmt.println("Channel Mask (Speakers):", file.channel_mask)
						head += 4
						
						// 16 bytes of SubFormat GUID, with the first two bytes being the format
						when VERBOSE do fmt.println("GUID:", temp_buf[head:head+16])
						file.format = Audio_Format(little_endian_u16(temp_buf[head:]))
						when VERBOSE do fmt.println("Format:", file.format)
						head += 2
						// Other GUID stuff
						head += 14
						
						
						if file.format == .ADPCM {
							// TODO: Support ADPCM (Very low priority)
							when VERBOSE do fmt.println("Samples pr chunk (In bit-depth field because it's unsupported):", file.bit_depth)
						} else {
							if valid_bits_pr_sample < bits_pr_sample {
								when VERBOSE do fmt.printfln("Warning: Bad metadata. valid_bits_pr_sample ({}) < bits_pr_sample ({}). Ignoring valid_bits_pr_sample. Using bits_pr_sample for bit_depth and sample_size.", valid_bits_pr_sample, bits_pr_sample)
								file.bit_depth   = bits_pr_sample
								file.sample_size = bits_pr_sample
								when VERBOSE do fmt.println("Bit depth & sample size:", bits_pr_sample)
							} else {
								// This is what is supposed to happen with wave_format_extensible
								file.bit_depth   = valid_bits_pr_sample
								file.sample_size = bits_pr_sample
								when VERBOSE do fmt.println("Bit depth:", file.bit_depth)
								when VERBOSE do fmt.println("Sample size/stride:", file.sample_size)
							}
						}
					}
					
					if file.format == .ADPCM {
						fmt.eprintln("WARNING! ADPCM file submitted. This is NOT SUPPORTED.")
					}
					
					head = data_end
					null_chunks = 0
				case "\x00\x00\x00\x00":
					if chunk_size == 0 {
						//null_chunks += 1
					}
			}
			when VERBOSE do fmt.println(print_data, "\n")
		} else {
			when VERBOSE do fmt.println("End of buffer reached.")
			break
		}
		
		
		head = next_chunk_start
		
		if null_chunks > 3 {
			when VERBOSE do fmt.println("Got more than 3 null chunks in a row. Quitting parse.")
			break
		}
		if data_reached {
			when VERBOSE do fmt.println("Data reached.")
		} 
	}
	
	file.channel_names = make([]string, file.channels, allocator=allocator)
	file.audio = make([][]f32, file.channels, allocator=allocator)
	
	
	// iXML Parsing
	if temp_ixml != "" {
		
		// Stripping null padding
		end := len(temp_ixml) - 1
		for end >= 0 && temp_ixml[end] == 0 {
			end -= 1
		}
		temp_ixml = temp_ixml[:end]
		
		
		naming_channel := 0
		/*
		interleave_set is here because we don't want to overwrite naming_channel
		with a number from CHANNEL_INDEX, if we've already set it with INTERLEAVE_INDEX.
		INTERLEAVE_INDEX is the number that actually tells you which channel the name
		belongs to. CHANNEL_INDEX is to be treated as a backup.
		*/
		interleave_set := false
		
		xml_recurse :: proc(doc: ^xml.Document, element_id: xml.Element_ID, file: ^Wav, naming_channel: ^int, interleave_set: ^bool, allocator:=context.allocator, indent := 0) {
			
			naming_channel := naming_channel
			interleave_set := interleave_set
			
			
			tab :: proc(indent: int) {
				for _ in 0..=indent {
					fmt.printf("\t")
				}
			}
		
			when VERBOSE do fmt.printf("\n")
			when VERBOSE do tab(indent)
		
			element := doc.elements[element_id]
			
			if element.kind == .Element {
				when VERBOSE do fmt.printf("<%v>", element.ident)
				
				if len(element.value) > 0 {
					value := element.value[0]
					switch element.ident {
						case "TRACK":
							interleave_set^ = false
						case "CHANNEL_INDEX":
							if !interleave_set^ {
								switch v in value {
									case string:
										naming_channel^, _ = strconv.parse_int(v)
									case u32:
										naming_channel^ = int(v)
								}
							}
							naming_channel^ -= 1
						case "INTERLEAVE_INDEX":
							interleave_set^ = true
							switch v in value {
								case string:
									naming_channel^, _ = strconv.parse_int(v)
								case u32:
									naming_channel^ = int(v)
							}
							naming_channel^ -= 1
						case "NAME":
							#partial switch v in value {
								case string:
									file.channel_names[naming_channel^] = strings.clone(v, allocator=allocator)
							}
						case "UBITS":
							#partial switch v in value {
								case string:
									for r, i in v {
										file.ubits[i] = u8(r) - u8('0')
									}
							}
						case "TAKE":
							#partial switch v in value {
								case string:
									file.take, _ = strconv.parse_int(v)
							}
						case "SCENE":
							#partial switch v in value {
								case string:
									file.scene = strings.clone(v, allocator=allocator)
							}
						case "PROJECT":
							#partial switch v in value {
								case string:
									file.project = strings.clone(v, allocator=allocator)
							}
						case "TAPE":
							#partial switch v in value {
								case string:
									file.tape = strings.clone(v, allocator=allocator)
							}
						case "CIRCLED":
							#partial switch v in value {
								case string:
									file.circled = v == "TRUE"
							}
						case "TIMECODE_FLAG":
							#partial switch v in value {
								case string:
									file.tc_dropframe = v != "NDF"
							}
						case "TIMECODE_RATE":
							#partial switch v in value {
								case string:
									end : int
									for r, i in v {
										if r == '/' {
											end = i
											break
										}
									}
									file.tc_framerate, _ = strconv.parse_f32(v[:end])
							}
						case "NOTE":
							#partial switch v in value {
								case string:
									if v != "" {
										file.note = strings.clone(v, allocator=allocator)
									}
							}
					}
				}
				
				for value in element.value {
					switch v in value {
						case string:
							when VERBOSE do fmt.printf(": %v", v)
						case xml.Element_ID:
							xml_recurse(doc, v, file, naming_channel, interleave_set, allocator, indent + 1)
					}
				}
		
				for attr in element.attribs {
					when VERBOSE do tab(indent + 1)
					when VERBOSE do fmt.printf("[Attr] %v: %v\n", attr.key, attr.val)
				}
			} else if element.kind == .Comment {
				when VERBOSE do fmt.printf("[COMMENT] %v\n", element.value)
			}
		
			return
		}
		
		
		parsed_ixml : ^xml.Document
		parsed_ixml, _ = xml.parse(temp_ixml, xml.Options{
			flags={.Ignore_Unsupported},
			expected_doctype = "",
		})
		xml_recurse(parsed_ixml, 0, &file, &naming_channel, &interleave_set, allocator)
	}
	
	// bext parsing
	if len(temp_bext) > 0 {
		
		// Stripping null padding
		/*end := len(temp_bext) - 1
		for end >= 0 && temp_bext[end] == 0 {
			end -= 1
		}
		temp_bext = temp_bext[:end]*/
		
		naming_channel := 0
		
		description := string(temp_bext[:256])
		for line in strings.split_lines(description) {
			
			if file.channel_names[naming_channel] == "" &&
			  (strings.starts_with(line, "sTRK") || strings.starts_with(line, "zTRK")) {
				eq_index := strings.index(line, "=")
				file.channel_names[naming_channel] = strings.clone(line[eq_index+1:], allocator=allocator)
				naming_channel += 1
			}
			if strings.starts_with(line[1:], "TAKE=")      {
				file.take, _ = strconv.parse_int(line[6:])
			}
			if strings.starts_with(line[1:], "CIRCLED=")   {
				file.circled = line[9:] == "TRUE"
			}
			if strings.starts_with(line[1:], "SPEED=")     {
				value := line[7:]
				num_end : int
				type_start : int
				for r, i in value {
					if !strings.contains_rune("1234567890.", r) && num_end == 0 {
						num_end = i
					}
					if r=='N' || r=='D' {
						type_start = i
						break
					}
				}
				file.tc_framerate, _ = strconv.parse_f32(value[:num_end])
				file.tc_dropframe = value[type_start:] != "ND"
				
			}
			// Only if ixml doesn't exist, so we don't allocate the note string twice.
			if file.note == "" {
				if strings.starts_with(line[1:], "NOTE=")  {
					v := line[6:]
					if v != "" {
						file.note = strings.clone(v, allocator=allocator)
					}
				}
			}
		}
		head := 0
		when VERBOSE do fmt.printf("Description: \n%v\n", string(temp_bext[head:256]))
		head += 256
		when VERBOSE do fmt.printf("Originator: %v\n", string(temp_bext[head:head+32]))
		head += 32
		when VERBOSE do fmt.printf("Originator Reference: %v\n", string(temp_bext[head:head+32]))
		head += 32
		date := string(temp_bext[head:head+10])
		when VERBOSE do fmt.printf("Origination Date: %v\n", date)
		date_splits := strings.split(date, "-")
		file.date.year,  _ = strconv.parse_int(date_splits[0])
		file.date.month, _ = strconv.parse_int(date_splits[1])
		file.date.day,   _ = strconv.parse_int(date_splits[2])
		delete(date_splits)
		head += 10
		when VERBOSE do fmt.printf("Origination Time: %v\n", string(temp_bext[head:head+8]))
		head += 8
		
		file.samples_since_midnight = int(little_endian_u64(temp_bext[head:head+8]))
		when VERBOSE do fmt.printf("Time Reference: %v (Samples since midnight, source of timecode)\n", file.samples_since_midnight)
		head += 8
		
		when VERBOSE do fmt.printf("Version: %v\n", little_endian_u16(temp_bext[head:head+2]))
		head += 2
		when VERBOSE do fmt.printf("UMID Skipped.\n")
		head += 64
		when VERBOSE do fmt.printf("Skipped reserved nothingness.\n")
		head += 190
		when VERBOSE do fmt.printf("Coding history:\n%v\n", string(temp_bext[head:]))
	}
	when VERBOSE do fmt.printfln("%#v", file)
	when VERBOSE do fmt.println()
	
	// just here to make some printing prettier
	temp_bext = nil
	
	temp_buf = nil
	return file, true
}

/*
Loads the full-length audio of the wav file into memory.
Pass in a slice of ints to select only specific channels to load.
*/
load :: proc(file : ^Wav, channels : []int = {}, allocator:=context.allocator) {
	channels := channels
	all_channels := false
	if len(channels) == 0 {
		all_channels = true
		channels = make([]int, file.channels)
		for c, i in channels {
			channels[i] = i
		}
	}
	err_quit := false
	for c in channels {
		if c >= file.channels {
			fmt.eprintfln("ERROR: Requested higher channel ({}) to be loaded than exists in file ({})!", c, file.channels)
		}
	}
	if err_quit do return
	
	
	// The main work part
	length := get_sample_count(file^)
	
	for c in channels {
		file.audio[c] = make([]f32, length, allocator=allocator)
	}
	
	samples_in_buffer :: 1024
	bytes_pr_sample := file.sample_size/8
	buffer_bytes : i64 = i64(bytes_pr_sample*file.channels*samples_in_buffer)
	buffer := make([]u8, buffer_bytes)
	
	decode := decode_24
	switch file.format {
		case .INT:
			switch file.bit_depth {
				case 16:
					decode = decode_16
				case 24:
					decode = decode_24
				case 32:
					decode = decode_32
			}
		case .FLOAT:
			switch file.bit_depth {
				case 32:
					decode = decode_f32
			}
		case .ADPCM:
			fmt.eprintln("ERROR: ADPCM is not supported!")
			return
	}
	
	head : i64 = i64(file.data_start)
	absolute_sample : u64 = 0
	file.handle, _ = os.open(file.path)
	for {
		os.read_at(file.handle, buffer, head)
		//fmt.println(buffer)
		for local_sample in 0..<samples_in_buffer {
			for c in channels {
				offset := (local_sample*file.channels+c)*bytes_pr_sample
				file.audio[c][absolute_sample] = decode(buffer[offset:])
			}
			absolute_sample += 1
			if absolute_sample >= u64(length) {
				break
			}
		}
		if absolute_sample >= u64(length) {
			break
		}
		head += buffer_bytes
	}
	os.close(file.handle)
	delete(buffer)
	
	
	if all_channels do delete(channels)
	return
}
tprint_timecode :: proc(file : Wav) -> string {
	return print_timecode(file, allocator=context.temp_allocator)
}
print_timecode :: proc(file : Wav, allocator := context.allocator) -> string {
	tc := get_timecode(file)
	using tc
	return fmt.aprintf("%02d:%02d:%02d:%02d",hour,minute,second,frame,
	                    allocator=allocator)
}
get_timecode :: proc(file : Wav) -> (output:Timecode) {
	seconds_since_midnight := file.samples_since_midnight / file.sample_rate
	output.hour   = int( seconds_since_midnight / 3600)
	output.minute = int((seconds_since_midnight % 3600) / 60)
	output.second = int( seconds_since_midnight % 60)
	output.frame  = int( math.round(f64(file.samples_since_midnight % file.sample_rate ) * f64(file.tc_framerate) / f64(file.sample_rate)))
	return
}

tprint_duration :: proc(file : Wav) -> string {
	return print_duration(file, allocator=context.temp_allocator)
}
print_duration :: proc(file : Wav, allocator := context.allocator) -> string {
	total_seconds := int(math.round(get_duration(file)))
	hours   := int( total_seconds / 3600)
	minutes := int((total_seconds % 3600) / 60)
	seconds := int( total_seconds % 60)
	
	return fmt.aprintf("%02d:%02d:%02d", hours, minutes, seconds, allocator=allocator)
}
/* Returns duration in seconds */
get_duration :: proc(file : Wav) -> f64 {
	return f64(get_sample_count(file))/f64(file.sample_rate)
}
get_sample_count :: proc(file : Wav) -> int {
	return file.data_size/file.channels/(file.sample_size/8)
}


wave_print :: proc(wave : []f32) {
	for sample, i in wave {
		fmt.printf("[%012d] ", i)
		bar_print(sample)
		fmt.printf(" %+01.04f\n", sample)
	}
}
bar_print :: proc(x : f32, one_side_width : int = 64, sign := '#', square := true) {
	x := x
	positive := x>0
	if square {
		x = math.sqrt(abs(x))
		if !positive do x *= -1
	}
	bar_length := int(math.round(abs(min(x, 1))*f32(one_side_width)))
	spaces := one_side_width-bar_length
	if positive {
		for _ in 0..<one_side_width {
			fmt.print(" ")
		}
		fmt.print("|")
		for _ in 0..<bar_length {
			fmt.print(sign)
		}
		for _ in 0..<spaces {
			fmt.print(" ")
		}
	} else {
		for _ in 0..<spaces {
			fmt.print(" ")
		}
		for _ in 0..<bar_length {
			fmt.print(sign)
		}
		fmt.print("|")
		for _ in 0..<one_side_width {
			fmt.print(" ")
		}
	}
}



decode_f32 :: proc(x : []u8) -> f32 {
	return (^f32)(&x[0])^
}
decode_32 :: proc(x : []u8) -> f32 {
	integer := i32(x[0])       |
	           i32(x[1]) <<  8 |
	           i32(x[2]) << 16 |
	           i32(x[3]) << 24
	return f32(integer) / f32(1 << 31)
}
decode_24 :: proc(x : []u8) -> f32 {
	integer := i32(x[0])       |
	           i32(x[1]) <<  8 |
	           i32(x[2]) << 16
	
	// Cheeky sign extension
	integer = (integer << 8) >> 8
	
	return f32(integer) / f32(1 << 23)
}
decode_16 :: proc(x : []u8) -> f32 {
	integer := i16(x[0]) |
	           i16(x[1]) <<  8
	return f32(integer) / f32(1 << 15)
}

little_endian_u64 :: proc(x : []u8) -> u64 {
	return u64(x[0])       |
	       u64(x[1]) <<  8 |
	       u64(x[2]) << 16 |
	       u64(x[3]) << 24 |
	       u64(x[4]) << 32 |
	       u64(x[5]) << 40 |
	       u64(x[6]) << 48 |
	       u64(x[7]) << 56
}

little_endian_u32 :: proc(x : []u8) -> u32 {
	return u32(x[0])       |
	       u32(x[1]) <<  8 |
	       u32(x[2]) << 16 |
	       u32(x[3]) << 24
}

little_endian_u16 :: proc(x : []u8) -> u16 {
	return u16(x[0])       |
	       u16(x[1]) <<  8
}