package main

import "base:runtime"
import "base:intrinsics"

import "core:fmt"
import "core:os/os2"
import "core:c"
import "core:log"
import "core:encoding/endian"
import "core:strings"
import "core:time"
import "core:sys/posix"

import "shared:pulse"

BITS_PER_SAMPLE :: 16
CHANNELS :: 1
SAMPLE_RATE :: 44100
SAMPLE_FORMAT: pulse.pa_sample_format : .S16LE
WAV_HEADER_SIZE :: 44

SHOULD_EXIT := false
SHOULD_DELETE_FILES := false
WE_ARE_READY := false
logger: log.Logger

InputDevice :: struct {
    name: string,
    description: string,
    index: u32,
    sample_spec: pulse.pa_sample_spec,
}

State :: enum {
    Initial,
    GettingSourceInfo,
    StartCapture,
    Capturing,
}

ReadCallbackData :: struct {
    file: ^os2.File,
    audio_data_size: uint,
}

state_cb :: proc "c" (ctx: ^pulse.pa_context, userdata: rawptr) {
    context = runtime.default_context()
    state := pulse.context_get_state(ctx)

    #partial switch state {
    case .READY:
        WE_ARE_READY = true
    case .FAILED:
        os2.exit(1)
    case .TERMINATED:
        SHOULD_EXIT = true
    }
}

source_info_cb :: proc "c" (ctx: ^pulse.pa_context, source_info: ^pulse.pa_source_info, end_of_list: b32, userdata: rawptr) {
    context = runtime.default_context()
    context.logger = logger
    sources := cast(^[dynamic]InputDevice)(userdata)

    if end_of_list {
        return
    }

    // Skip all monitor devices
    if source_info.monitor_of_sink != 0xFFFFFFFF {
        return
    }

    log.info("Found input device:", source_info.description)

    if source_info.mute == 1 {
        log.warnf("Attempting to unmute input device \"%s\". Make sure it did unmute, we don't log any errors if it failed.", source_info.description)
        pulse.context_set_source_mute_by_name(ctx, source_info.name, false, nil, nil)
    }

    append(sources, InputDevice{
        name = strings.clone_from_cstring(source_info.name),
        description = strings.clone_from_cstring(source_info.description),
        index = source_info.index,
        sample_spec = source_info.sample_spec,
    })
}

read_callback :: proc "c" (stream: ^pulse.pa_stream, nbytes: c.size_t, userdata: rawptr) {
    context = runtime.default_context()
    context.logger = logger
    cb_data := transmute(^ReadCallbackData)userdata

    log.assert(cb_data.file != nil, "File is nil")
    log.assert(stream != nil, "Stream is nil")
    log.assert(nbytes > 0, "size is 0")

    size: c.size_t
    data: rawptr

    err := pulse.stream_peek(stream, &data, &size)
    log.assert(err == 0, fmt.tprintfln("Err non-zero: %v", pulse.strerror(err)))

    writ, writ_err := os2.write_ptr(cb_data.file, data, auto_cast size)
    log.assert(writ_err == nil, fmt.tprint("Failed to write to file:", writ_err))
    cb_data.audio_data_size += auto_cast writ

    err = pulse.stream_drop(stream)
    log.assert(err == 0, fmt.tprintfln("Err non-zero: %v", pulse.strerror(err)))
}

prepare_output_files :: proc(file_name1, file_name2: string, four_byte_slice, two_byte_slice: []byte) -> (^os2.File, ^os2.File) {
    output_file, open_err := os2.open(file_name1, {.Write, .Create, .Trunc}, 0o644)
    if open_err != nil {
        log.panic("open_err:", open_err)
    }

    filtered_file, filtered_err := os2.open(file_name2, {.Write, .Create, .Trunc}, 0o644)
    if filtered_err != nil {
        log.panic("filtered_open:", filtered_err)
    }

    os2.write(output_file, []byte{'R', 'I', 'F', 'F'}) // Riff file
    os2.write(output_file, []byte{0,0,0,0})            // File size - 8 bytes. To be filled in later
    os2.write(output_file, []byte{'W', 'A', 'V', 'E'}) // WAVE data
    os2.write(output_file, []byte{'f', 'm', 't', ' '}) // Format chunk marker with NUL byte
    endian.put_i32(four_byte_slice, .Little, 16)
    os2.write(output_file, four_byte_slice)            // Chunk size: 16 for PCM
    endian.put_i16(two_byte_slice, .Little, 1)
    os2.write(output_file, two_byte_slice)             // Audio format: 1 for PCM
    os2.write(output_file, two_byte_slice)             // Number of channels
    endian.put_i32(four_byte_slice, .Little, SAMPLE_RATE)
    os2.write(output_file, four_byte_slice)            // Sample rate
    endian.put_i32(four_byte_slice, .Little, SAMPLE_RATE * CHANNELS * BITS_PER_SAMPLE / 8)
    os2.write(output_file, four_byte_slice)            // Byte rate
    endian.put_i16(two_byte_slice, .Little, CHANNELS * BITS_PER_SAMPLE / 8)
    os2.write(output_file, two_byte_slice)             // Block align
    endian.put_i16(two_byte_slice, .Little, BITS_PER_SAMPLE)
    os2.write(output_file, two_byte_slice)             // Bits per sample
    os2.write(output_file, []byte{'d', 'a', 't', 'a'}) // Data chunk header
    os2.write(output_file, []byte{0,0,0,0}) // Length of data. To be filled in later

    os2.write(filtered_file, []byte{'R', 'I', 'F', 'F'}) // Riff file
    os2.write(filtered_file, []byte{0,0,0,0})            // File size - 8 bytes. To be filled in later
    os2.write(filtered_file, []byte{'W', 'A', 'V', 'E'}) // WAVE data
    os2.write(filtered_file, []byte{'f', 'm', 't', ' '}) // Format chunk marker with NUL byte
    endian.put_i32(four_byte_slice, .Little, 16)
    os2.write(filtered_file, four_byte_slice)            // Chunk size: 16 for PCM
    endian.put_i16(two_byte_slice, .Little, 1)
    os2.write(filtered_file, two_byte_slice)             // Audio format: 1 for PCM
    os2.write(filtered_file, two_byte_slice)             // Number of channels
    endian.put_i32(four_byte_slice, .Little, SAMPLE_RATE)
    os2.write(filtered_file, four_byte_slice)            // Sample rate
    endian.put_i32(four_byte_slice, .Little, SAMPLE_RATE * CHANNELS * BITS_PER_SAMPLE / 8)
    os2.write(filtered_file, four_byte_slice)            // Byte rate
    endian.put_i16(two_byte_slice, .Little, CHANNELS * BITS_PER_SAMPLE / 8)
    os2.write(filtered_file, two_byte_slice)             // Block align
    endian.put_i16(two_byte_slice, .Little, BITS_PER_SAMPLE)
    os2.write(filtered_file, two_byte_slice)             // Bits per sample
    os2.write(filtered_file, []byte{'d', 'a', 't', 'a'}) // Data chunk header
    os2.write(filtered_file, []byte{0,0,0,0}) // Length of data. To be filled in later

    return output_file, filtered_file
}

main :: proc() {
    logger = log.create_console_logger()
    context.logger = logger

    ml := pulse.mainloop_new()
    log.assert(ml != nil, "mainloop is a nil pointer")
    defer pulse.mainloop_free(ml)

    api := pulse.mainloop_get_api(ml)
    log.assert(api != nil, "mainloop_api is a nil pointer")

    ctx := pulse.context_new(api, nil)
    log.assert(ctx != nil, "pa_context is a nil pointer")

    err := pulse.context_connect(ctx, nil, .NOFLAGS, nil)
    log.assert(err == 0, fmt.tprintfln("Err non-zero: %v", pulse.strerror(err)))
    defer pulse.context_unref(ctx)
    defer pulse.context_disconnect(ctx)

    pulse.context_set_state_callback(ctx, state_cb, nil)

    state := State.Initial
    op: ^pulse.pa_operation

    sources := make([dynamic]InputDevice, 0)
    stream: ^pulse.pa_stream
    filtered_stream: ^pulse.pa_stream
    defer if stream != nil {
        pulse.stream_disconnect(stream)
        pulse.stream_unref(stream)
    }
    defer if filtered_stream != nil {
        pulse.stream_disconnect(filtered_stream)
        pulse.stream_unref(filtered_stream)
    }

    action: posix.sigaction_t
    action.sa_handler = proc "c" (signal: posix.Signal) {
        context = runtime.default_context()
        fmt.println("\nCaught SIGINT")
        intrinsics.atomic_store(&SHOULD_EXIT, true)
    }
    posix.sigaction(.SIGINT, &action, nil)

    four_byte_slice := make([]byte, 4)
    two_byte_slice := make([]byte, 2)
    defer delete(four_byte_slice)
    defer delete(two_byte_slice)

    now := time.now()
    year, month, day := time.date(now)
    hour, minute, second := time.clock_from_time(now)

    file_name1 := fmt.tprintf("capture-%v-%02d-%02dT%02d:%02d:%02d.wav", year, int(month), day, hour, minute, second)
    file_name2 := fmt.tprintf("capture-filtered-%v-%02d-%02dT%02d:%02d:%02d.wav", year, int(month), day, hour, minute, second)

    output_file, filtered_file := prepare_output_files(file_name1, file_name2, four_byte_slice, two_byte_slice)
    defer {
        os2.close(output_file)
        os2.close(filtered_file)
    }

    read_cb_data := ReadCallbackData{
        file = output_file,
        audio_data_size = 0,
    }
    filtered_cb_data := ReadCallbackData{
        file = filtered_file,
        audio_data_size = 0,
    }

    loop: for {
        if !WE_ARE_READY {
            pulse.mainloop_iterate(ml, true, nil)
            continue
        }

        if SHOULD_EXIT {
            break loop
        }

        switch state {
        case .Initial:
            op = pulse.context_get_source_info_list(ctx, source_info_cb, &sources)
            state = .GettingSourceInfo
        case .GettingSourceInfo:
            if pulse.operation_get_state(op) == .DONE {
                pulse.operation_unref(op)
                if len(sources) < 2 {
                    SHOULD_DELETE_FILES = true
                    break loop
                }
                log.debug(sources)
                state = .StartCapture
            }
        case .StartCapture:
            log.info("Starting Capture")
            stream = pulse.stream_new(ctx, "Default Capture", &{ rate = SAMPLE_RATE, format = SAMPLE_FORMAT, channels = 1 }, nil)
            filtered_stream = pulse.stream_new(ctx, "Filtered Capture", &{ rate = SAMPLE_RATE, format = SAMPLE_FORMAT, channels = 1 }, nil)
            pulse.stream_set_read_callback(stream, read_callback, &read_cb_data)
            pulse.stream_set_read_callback(filtered_stream, read_callback, &filtered_cb_data)
            err = pulse.stream_connect_record(stream, strings.clone_to_cstring(sources[0].name), nil, {})
            log.assert(err == 0, fmt.tprintfln("Err non-zero: %v", pulse.strerror(err)))
            if !strings.contains(sources[1].description, "Filtered") {
                log.warn("The device chosen as the \"Filtered\" device doesn't contain \"Filtered\"in its description")
            }

            err = pulse.stream_connect_record(filtered_stream, strings.clone_to_cstring(sources[1].name), nil, {})
            log.assert(err == 0, fmt.tprintfln("Err non-zero: %v", pulse.strerror(err)))

            state = .Capturing
        case .Capturing:
            // Placeholder state just so we leave the StartCapture state.
        }

        pulse.mainloop_iterate(ml, state != .StartCapture, nil)
    }

    endian.put_i32(four_byte_slice, .Little, i32(read_cb_data.audio_data_size + WAV_HEADER_SIZE - 8))
    os2.write_at(output_file, four_byte_slice, 4) // Overall size - 8 bytes

    endian.put_i32(four_byte_slice, .Little, i32(read_cb_data.audio_data_size))
    os2.write_at(output_file, four_byte_slice, 40) // Data size

    endian.put_i32(four_byte_slice, .Little, i32(filtered_cb_data.audio_data_size + WAV_HEADER_SIZE - 8))
    os2.write_at(filtered_file, four_byte_slice, 4) // Overall size - 8 bytes

    endian.put_i32(four_byte_slice, .Little, i32(filtered_cb_data.audio_data_size))
    os2.write_at(filtered_file, four_byte_slice, 40) // Data size

    if SHOULD_DELETE_FILES {
        log.error("NoiseTorch might not be running or there's no microphone plugged in")
        os2.remove(file_name1)
        os2.remove(file_name2)
    }

    log.info("Exiting...")
}