package main

import (
	"container/heap"
	"flag"
	"fmt"
	"math"
	"math/rand"
	"os"
	"sort"
	"strconv"
	"time"
)

const usage = "smerge [count, count, ... count]"

var max = flag.Int("max", 1000, "maximum random number for each source")

func main() {
	flag.Parse()

	args := flag.Args()
	if len(args) < 1 {
		fmt.Fprintf(os.Stderr, "usage: %v\n", usage)
		os.Exit(1)
	}
	srcs := []<-chan int{}
	for _, arg := range args {
		i, err := strconv.Atoi(arg)
		if err != nil {
			fmt.Fprintf(os.Stderr, "%v\n", err)
			os.Exit(1)
		}

		srcs = append(srcs, source(i, *max))
	}
	f := fmt.Sprintf("%%%dd\n", int(math.Log10(float64(*max))))
	for i := range merge(srcs...) {
		fmt.Printf(f, i)
	}
}

func source(c, max int) <-chan int {
	out := make(chan int)
	go func() {
		vals := make([]int, c)
		src := rand.New(rand.NewSource(time.Now().UnixNano()))
		for i := 0; i < c; i++ {
			vals[i] = src.Intn(max)
		}

		sort.Ints(vals)

		for _, i := range vals {
			out <- i
		}
		close(out)
	}()
	return out
}

func merge(cs ...<-chan int) <-chan int {
	out := make(chan int)

	go func() {
		h := &items{}
		heap.Init(h)

		// prime the pumps
		for i, src := range cs {
			head, ok := <-src
			if !ok {
				continue
			}
			it := item{
				val: head,
				src: i,
			}
			heap.Push(h, it)
		}

		for h.Len() > 0 {
			top := heap.Pop(h).(item)
			out <- top.val
			if head, ok := <-cs[top.src]; ok {
				it := item{
					val: head,
					src: top.src,
				}
				heap.Push(h, it)
			}
		}

		close(out)
	}()
	return out
}

// item keeps track of an integer value and the index of its source.
//
// the index is used to chose the next source to use when a value has been
// pulled.
type item struct {
	val int
	src int
}

// items is a min-heap of item.
type items []item

func (h items) Len() int           { return len(h) }
func (h items) Less(i, j int) bool { return h[i].val < h[j].val }
func (h items) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *items) Push(x interface{}) {
	*h = append(*h, x.(item))
}

func (h *items) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}