package query

import (
	"bufio"

	"code.justin.tv/release/trace/analysis/callgraph"
	"code.justin.tv/release/trace/analysis/tx"
	"code.justin.tv/release/trace/pbmsg"
	"github.com/golang/protobuf/proto"
)

const (
	decodeWorkerCount = 100
	decodeQueueSize   = 1000
	decodeBatchSize   = 100
)

// TraceStore provides access to transaction data
type TraceStore interface {
	// Add adds a new event to the TraceStore
	Add(*pbmsg.Event)
	// All returns a collection of all transactions
	All() tx.TransactionSet
	// GetTx returns a single transaction by ID. Should return nil if
	// the transaction is not found.
	GetTx(tx.TransactionID) *tx.Transaction
	// Len returns the total number of transactions in the TraceStore
	Len() int
}

type LoadStats struct {
	Messages, NumBytes, MissingRoot uint64
}

// Load iterates over a Scanner which provides pbmsg.EventSets and
// loads the events into a TraceStore. Stops if it has loaded limit
// Events. Set limit to -1 for unlimited loading.
func Load(s TraceStore, sc *bufio.Scanner, limit int64) (*LoadStats, error) {
	var (
		err   error
		stats *LoadStats
		i     int64
	)

	type decoded struct {
		buf []byte

		// a read from done indicates that proto decode worker is not writing
		// to es
		done chan struct{}
		es   *pbmsg.EventSet
		err  error
	}

	done := make(chan struct{})
	// send decode tasks to workers
	todo := make(chan decoded, decodeQueueSize)
	// send decode tasks (which will soon be complete) to the final consumer
	complete := make(chan decoded, decodeQueueSize)

	stats = &LoadStats{}
	go func() {
		defer close(todo)
		defer close(complete)

		for broken := false; !broken; {
			var buf []byte
			for i := 0; i < decodeBatchSize; i++ {
				if !sc.Scan() {
					broken = true
					break
				}
				// we don't own the output of sc.Bytes
				buf = append(buf, sc.Bytes()...)
				stats.Messages++
				stats.NumBytes += uint64(len(sc.Bytes()))
			}

			task := decoded{
				buf:  buf,
				done: make(chan struct{}),
				es:   new(pbmsg.EventSet),
			}

			select {
			case todo <- task:
			case <-done:
				return
			}

			select {
			case complete <- task:
			case <-done:
				return
			}
		}
	}()

	for i := 0; i < decodeWorkerCount; i++ {
		go func() {
			for task := range todo {
				task.err = proto.Unmarshal(task.buf, task.es)
				close(task.done)
			}
		}()
	}

	// a single goroutine reads complete decode tasks so the events remain
	// ordered

Scanning:
	for task := range complete {
		<-task.done

		es, err := task.es, task.err

		if err != nil {
			close(done)
			// make sure the scanner is no longer running
			for range complete {
			}
			return stats, err
		}
		for _, ev := range es.Event {
			s.Add(ev)
			i += 1
			if i == limit {
				break Scanning
			}
		}
	}

	close(done)
	// make sure the scanner is no longer running
	for range complete {
	}

	err = sc.Err()
	if err != nil {
		return stats, err
	}

	for _, t := range s.All().All() {
		if callgraph.Empty(t.Root) {
			stats.MissingRoot += 1
		}
	}

	return stats, nil
}

// TraceMemoryStore holds its transactions in memory in a single big map.
type TraceMemoryStore struct {
	txs map[tx.TransactionID]*tx.Transaction
}

func NewTraceMemoryStore() *TraceMemoryStore {
	return &TraceMemoryStore{
		txs: make(map[tx.TransactionID]*tx.Transaction),
	}
}

func (s *TraceMemoryStore) All() tx.TransactionSet {
	set := tx.NewInMemoryTransactionSet()
	for _, tx := range s.txs {
		set.Add(tx)
	}
	return set
}

func (s *TraceMemoryStore) GetTx(id tx.TransactionID) *tx.Transaction {
	return s.txs[id]
}

// Add appends an event to the MemoryStore, associating it with its
// transaction based on its tranaction ID
func (s *TraceMemoryStore) Add(event *pbmsg.Event) {
	var txid tx.TransactionID
	copy(txid[:], event.TransactionId)
	t := s.GetTx(txid)
	if t == nil {
		t = tx.NewTransaction(txid)
		s.txs[txid] = t
	}
	t.Add(event)
}

func (s *TraceMemoryStore) Len() int {
	return len(s.txs)
}
