package main

import zmq "../../"
import "core:fmt"
import "core:log"
import "core:os"

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

	major, minor, patch := zmq.version()
	log.infof(
		"zmq echo (libzmq %d.%d.%d) over inproc; DEALER/DEALER via zmq_poll",
		major,
		minor,
		patch,
	)

	ctx := zmq.ctx_new()
	defer zmq.ctx_term(ctx)

	// a is user-facing (stdin/stdout); b is the echoer. DEALER/DEALER is
	// async two-way: either end may send/recv freely, so the poll loop
	// needs no request/reply state tracking (unlike REQ/REP).
	a, serr := zmq.socket(ctx, .DEALER)
	if serr != nil {
		log.panic(zmq.strerror(serr))
	}
	defer zmq.close(a)
	if err := zmq.bind(a, "inproc://echo"); err != nil {
		log.panic(zmq.strerror(err))
	}

	b, cerr := zmq.socket(ctx, .DEALER)
	if cerr != nil {
		log.panic(zmq.strerror(cerr))
	}
	defer zmq.close(b)
	if err := zmq.connect(b, "inproc://echo"); err != nil {
		log.panic(zmq.strerror(err))
	}

	echo(a, b)
}

echo :: proc(a, b: zmq.Socket) {
	pending: [dynamic]u8 // stdin line assembly
	defer delete(pending)
	msg: [255]u8

	// Poll stdin + both DEALERs together. Each item polls its socket
	// when non-nil, otherwise its fd.
	items := [3]zmq.Poll_Item {
		{fd = 0, events = {.POLLIN}}, // stdin
		{socket = a, events = {.POLLIN}}, // reply arrived from b
		{socket = b, events = {.POLLIN}}, // request arrived from a
	}

	// in_flight tracks lines sent via a whose reply hasn't been printed yet.
	// On stdin EOF we stop reading stdin but keep looping until every
	// in-flight echo has been printed, so piped input isn't truncated.
	in_flight := 0
	eof := false
	run := true

	fmt.print("> ")
	for run {
		_, err := zmq.poll(items[:], -1) // block until something is ready
		if err != nil {
			fmt.eprintfln("poll: %s", zmq.strerror(err))
			return
		}

		// stdin ready: read what's available, assemble into lines.
		if !eof && .POLLIN in items[0].revents {
			tmp: [256]u8
			n, _ := os.read(os.stdin, tmp[:])
			if n == 0 {
				// EOF: flush any trailing partial line, then stop
				// watching stdin and drain remaining replies.
				eof = true
				items[0].events = {}
				if len(pending) > 0 {
					if _, e := zmq.send(a, pending[:]); e != nil {
						fmt.eprintfln("send a: %s", zmq.strerror(e))
					} else {
						in_flight += 1
					}
					resize(&pending, 0)
				}
			} else {
				append(&pending, ..tmp[:n])
			}
		}

		// a got a reply from b -> print it.
		if .POLLIN in items[1].revents {
			n, e := zmq.recv(a, msg[:])
			if e != nil {
				fmt.eprintfln("recv a: %s", zmq.strerror(e))
			} else {
				fmt.printf("< %s\n> ", string(msg[:n]))
				in_flight -= 1
			}
		}

		// b got a request from a -> echo it back.
		if .POLLIN in items[2].revents {
			n, e := zmq.recv(b, msg[:])
			if e != nil {
				fmt.eprintfln("recv b: %s", zmq.strerror(e))
			} else {
				if _, se := zmq.send(b, msg[:n]); se != nil {
					fmt.eprintfln("send b: %s", zmq.strerror(se))
				}
			}
		}

		// Drain complete lines from the stdin buffer -> send via a.
		if !eof {
			for {
				idx := -1
				for i in 0 ..< len(pending) {
					if pending[i] == '\n' {
						idx = i
						break
					}
				}
				if idx < 0 {
					break
				}
				line := pending[:idx]
				if len(line) > 0 {
					if _, e := zmq.send(a, line); e != nil {
						fmt.eprintfln("send a: %s", zmq.strerror(e))
					} else {
						in_flight += 1
					}
				}
				// consume the line + its newline
				consumed := idx + 1
				copy(pending[:], pending[consumed:])
				resize(&pending, len(pending) - consumed)
			}
		}

		if eof && in_flight == 0 {
			run = false
		}
	}
}