package main

import (
	"fmt"
	"runtime"
	"sync"
	"time"
)

// QuickCheckRxa validates if a given rxa produces the program as output
func QuickCheckRxa(rxa int, prog []int) bool {
	rxb := 0
	rxc := 0
	ip := 0
	output := []int{}

	for ip < len(prog) {
		opcode := prog[ip]
		operand := prog[ip+1]
		switch opcode {
		case 0: // adv: rxa / 2^operand
			if operand >= 0 && operand <= 3 {
				rxa = rxa / (1 << operand)
			} else {
				return false
			}
			ip += 2
		case 1: // bxl: rxb XOR operand
			rxb ^= operand
			ip += 2
		case 2: // bst: operand % 8
			switch operand {
			case 4:
				operand = rxa
			case 5:
				operand = rxb
			case 6:
				operand = rxc
			}
			rxb = operand % 8
			ip += 2
		case 3: // jnz: jump if rxa != 0
			if rxa != 0 {
				ip = operand
			} else {
				ip += 2
			}
		case 4: // bxc: rxb XOR rxc
			rxb ^= rxc
			ip += 2
		case 5: // out: operand % 8
			switch operand {
			case 4:
				operand = rxa
			case 5:
				operand = rxb
			case 6:
				operand = rxc
			}
			output = append(output, operand%8)
			ip += 2
		default:
			return false
		}

		// Check if output mismatches the program so far
		if len(output) > len(prog) {
			return false
		}
		for i, val := range output {
			if val != prog[i] {
				return false
			}
		}
	}

	return len(output) == len(prog)
}

func workerDynamic(taskQueue <-chan int, prog []int, result chan int, done <-chan struct{}, wg *sync.WaitGroup) {
	defer wg.Done()
	for {
		select {
		case <-done:
			return
		case rxa, ok := <-taskQueue:
			if !ok {
				return
			}
			if QuickCheckRxa(rxa, prog) {
				select {
				case result <- rxa:
				default:
				}
				return
			}
		}
	}
}

func main() {
	prog := []int{2, 4, 1, 3, 7, 5, 4, 0, 1, 3, 0, 3, 5, 5, 3, 0}
	start := time.Now()
	taskQueue := make(chan int, 1000) // Buffered queue for tasks
	result := make(chan int, 1)
	done := make(chan struct{})
	var wg sync.WaitGroup
	numWorkers := runtime.NumCPU()

	// Status updater
	go func() {
		for {
			time.Sleep(1 * time.Second)
			fmt.Printf("Queue size: %d\n", len(taskQueue))
		}
	}()

	// Launch workers
	for i := 0; i < numWorkers; i++ {
		wg.Add(1)
		go workerDynamic(taskQueue, prog, result, done, &wg)
	}

	// Generate tasks dynamically
	go func() {
		for rxa := 0; ; rxa++ {
			select {
			case <-done:
				close(taskQueue)
				return
			default:
				taskQueue <- rxa
			}
		}
	}()

	// Wait for result
	foundRxa := <-result
	close(done) // Signal workers to stop
	wg.Wait()   // Ensure all workers exit cleanly
	duration := time.Since(start)

	fmt.Printf("Found rxa: %d\n", foundRxa)
	fmt.Printf("Execution time: %s\n", duration)
}