package scanproto

import (
	"sync"
	"time"

	"code.justin.tv/release/trace/api"
	"code.justin.tv/release/trace/pbmsg"
)

type limitedTxSrc struct {
	base  TransactionSource
	limit time.Duration

	wg sync.WaitGroup

	timer *time.Timer
}

func (lts *limitedTxSrc) Run() {
	lts.wg.Add(1)
	lts.timer = time.AfterFunc(lts.limit, lts.Stop)
	go func() {
		defer lts.wg.Done()
		lts.base.Run()
	}()
	lts.wg.Wait()
}

func (lts *limitedTxSrc) Transactions() <-chan *api.Transaction {
	return lts.base.Transactions()
}

func (lts *limitedTxSrc) Errors() <-chan error {
	return lts.base.Errors()
}

func (lts *limitedTxSrc) Stop() {
	lts.base.Stop()
	lts.wg.Wait()
}

// Limit a TransactionSource to only run for d duration. It will call
// the base TransactionSource's Stop function after d time passes, which
// might take time, so it's not guaranteed to finish in less than d time.
func LimitedTransactionSource(s TransactionSource, d time.Duration) TransactionSource {
	return &limitedTxSrc{
		base:  s,
		limit: d,
	}
}

type legacySource struct {
	src TransactionSource
	ch  chan *pbmsg.Transaction
}

func LegacySource(src TransactionSource) LegacyTransactionSource {
	ls := &legacySource{
		src: src,
		ch:  make(chan *pbmsg.Transaction, cap(src.Transactions())),
	}
	return ls
}

func (ls *legacySource) Errors() <-chan error {
	return ls.src.Errors()
}

func (ls *legacySource) Run() {
	go func() {
		defer close(ls.ch)
		for tx := range ls.src.Transactions() {
			ls.ch <- pbmsg.DowngradeTransaction(tx)
		}
	}()
	ls.src.Run()
}

func (ls *legacySource) Stop() {
	ls.src.Stop()
}

func (ls *legacySource) Transactions() <-chan *pbmsg.Transaction {
	return ls.ch
}
