// Copyright 2023 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package trace

import (
	"bufio"
	"bytes"
	"cmp"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"slices"
	"strings"
	"time"

	"internal/trace/tracev2"
	"internal/trace/version"
)

// generation contains all the trace data for a single
// trace generation. It is purely data: it does not
// track any parse state nor does it contain a cursor
// into the generation.
type generation struct {
	gen        uint64
	batches    map[ThreadID][]batch
	batchMs    []ThreadID
	cpuSamples []cpuSample
	minTs      timestamp
	*evTable
}

// readGeneration buffers and decodes the structural elements of a trace generation
// out of r.
func readGeneration(r *bufio.Reader, ver version.Version) (*generation, error) {
	if ver < version.Go126 {
		return nil, errors.New("internal error: readGeneration called for <1.26 trace")
	}
	g := &generation{
		evTable: &evTable{
			pcs: make(map[uint64]frame),
		},
		batches: make(map[ThreadID][]batch),
	}

	// Read batches one at a time until we either hit the next generation.
	for {
		b, gen, err := readBatch(r)
		if err == io.EOF {
			if len(g.batches) != 0 {
				return nil, errors.New("incomplete generation found; trace likely truncated")
			}
			return nil, nil // All done.
		}
		if err != nil {
			return nil, err
		}
		if g.gen == 0 {
			// Initialize gen.
			g.gen = gen
		}
		if b.isEndOfGeneration() {
			break
		}
		if gen == 0 {
			// 0 is a sentinel used by the runtime, so we'll never see it.
			return nil, fmt.Errorf("invalid generation number %d", gen)
		}
		if gen != g.gen {
			return nil, fmt.Errorf("broken trace: missing end-of-generation event, or generations are interleaved")
		}
		if g.minTs == 0 || b.time < g.minTs {
			g.minTs = b.time
		}
		if err := processBatch(g, b, ver); err != nil {
			return nil, err
		}
	}

	// Check some invariants.
	if g.freq == 0 {
		return nil, fmt.Errorf("no frequency event found")
	}
	if !g.hasClockSnapshot {
		return nil, fmt.Errorf("no clock snapshot event found")
	}

	// N.B. Trust that the batch order is correct. We can't validate the batch order
	// by timestamp because the timestamps could just be plain wrong. The source of
	// truth is the order things appear in the trace and the partial order sequence
	// numbers on certain events. If it turns out the batch order is actually incorrect
	// we'll very likely fail to advance a partial order from the frontier.

	// Compactify stacks and strings for better lookup performance later.
	g.stacks.compactify()
	g.strings.compactify()

	// Validate stacks.
	if err := validateStackStrings(&g.stacks, &g.strings, g.pcs); err != nil {
		return nil, err
	}

	// Now that we have the frequency, fix up CPU samples.
	fixUpCPUSamples(g.cpuSamples, g.freq)
	return g, nil
}

// spilledBatch represents a batch that was read out for the next generation,
// while reading the previous one. It's passed on when parsing the next
// generation.
//
// Used only for trace versions < Go126.
type spilledBatch struct {
	gen uint64
	*batch
}

// readGenerationWithSpill buffers and decodes the structural elements of a trace generation
// out of r. spill is the first batch of the new generation (already buffered and
// parsed from reading the last generation). Returns the generation and the first
// batch read of the next generation, if any.
//
// If gen is non-nil, it is valid and must be processed before handling the returned
// error.
func readGenerationWithSpill(r *bufio.Reader, spill *spilledBatch, ver version.Version) (*generation, *spilledBatch, error) {
	if ver >= version.Go126 {
		return nil, nil, errors.New("internal error: readGenerationWithSpill called for Go 1.26+ trace")
	}
	g := &generation{
		evTable: &evTable{
			pcs: make(map[uint64]frame),
		},
		batches: make(map[ThreadID][]batch),
	}
	// Process the spilled batch.
	if spill != nil {
		// Process the spilled batch, which contains real data.
		g.gen = spill.gen
		g.minTs = spill.batch.time
		if err := processBatch(g, *spill.batch, ver); err != nil {
			return nil, nil, err
		}
		spill = nil
	}
	// Read batches one at a time until we either hit the next generation.
	var spillErr error
	for {
		b, gen, err := readBatch(r)
		if err == io.EOF {
			break
		}
		if err != nil {
			if g.gen != 0 {
				// This may be an error reading the first batch of the next generation.
				// This is fine. Let's forge ahead assuming that what we've got so
				// far is fine.
				spillErr = err
				break
			}
			return nil, nil, err
		}
		if gen == 0 {
			// 0 is a sentinel used by the runtime, so we'll never see it.
			return nil, nil, fmt.Errorf("invalid generation number %d", gen)
		}
		if g.gen == 0 {
			// Initialize gen.
			g.gen = gen
		}
		if gen == g.gen+1 {
			// TODO: Increment the generation with wraparound the same way the runtime does.
			spill = &spilledBatch{gen: gen, batch: &b}
			break
		}
		if gen != g.gen {
			// N.B. Fail as fast as possible if we see this. At first it
			// may seem prudent to be fault-tolerant and assume we have a
			// complete generation, parsing and returning that first. However,
			// if the batches are mixed across generations then it's likely
			// we won't be able to parse this generation correctly at all.
			// Rather than return a cryptic error in that case, indicate the
			// problem as soon as we see it.
			return nil, nil, fmt.Errorf("generations out of order")
		}
		if g.minTs == 0 || b.time < g.minTs {
			g.minTs = b.time
		}
		if err := processBatch(g, b, ver); err != nil {
			return nil, nil, err
		}
	}

	// Check some invariants.
	if g.freq == 0 {
		return nil, nil, fmt.Errorf("no frequency event found")
	}
	if ver >= version.Go125 && !g.hasClockSnapshot {
		return nil, nil, fmt.Errorf("no clock snapshot event found")
	}

	// N.B. Trust that the batch order is correct. We can't validate the batch order
	// by timestamp because the timestamps could just be plain wrong. The source of
	// truth is the order things appear in the trace and the partial order sequence
	// numbers on certain events. If it turns out the batch order is actually incorrect
	// we'll very likely fail to advance a partial order from the frontier.

	// Compactify stacks and strings for better lookup performance later.
	g.stacks.compactify()
	g.strings.compactify()

	// Validate stacks.
	if err := validateStackStrings(&g.stacks, &g.strings, g.pcs); err != nil {
		return nil, nil, err
	}

	// Now that we have the frequency, fix up CPU samples.
	fixUpCPUSamples(g.cpuSamples, g.freq)
	return g, spill, spillErr
}

// processBatch adds the batch to the generation.
func processBatch(g *generation, b batch, ver version.Version) error {
	switch {
	case b.isStringsBatch():
		if err := addStrings(&g.strings, b); err != nil {
			return err
		}
	case b.isStacksBatch():
		if err := addStacks(&g.stacks, g.pcs, b); err != nil {
			return err
		}
	case b.isCPUSamplesBatch():
		samples, err := addCPUSamples(g.cpuSamples, b)
		if err != nil {
			return err
		}
		g.cpuSamples = samples
	case b.isSyncBatch(ver):
		if err := setSyncBatch(&g.sync, b, ver); err != nil {
			return err
		}
	case b.exp != tracev2.NoExperiment:
		if g.expBatches == nil {
			g.expBatches = make(map[tracev2.Experiment][]ExperimentalBatch)
		}
		if err := addExperimentalBatch(g.expBatches, b); err != nil {
			return err
		}
	case b.isEndOfGeneration():
		return errors.New("internal error: unexpectedly processing EndOfGeneration; broken trace?")
	default:
		if _, ok := g.batches[b.m]; !ok {
			g.batchMs = append(g.batchMs, b.m)
		}
		g.batches[b.m] = append(g.batches[b.m], b)
	}
	return nil
}

// validateStackStrings makes sure all the string references in
// the stack table are present in the string table.
func validateStackStrings(
	stacks *dataTable[stackID, stack],
	strings *dataTable[stringID, string],
	frames map[uint64]frame,
) error {
	var err error
	stacks.forEach(func(id stackID, stk stack) bool {
		for _, pc := range stk.pcs {
			frame, ok := frames[pc]
			if !ok {
				err = fmt.Errorf("found unknown pc %x for stack %d", pc, id)
				return false
			}
			_, ok = strings.get(frame.funcID)
			if !ok {
				err = fmt.Errorf("found invalid func string ID %d for stack %d", frame.funcID, id)
				return false
			}
			_, ok = strings.get(frame.fileID)
			if !ok {
				err = fmt.Errorf("found invalid file string ID %d for stack %d", frame.fileID, id)
				return false
			}
		}
		return true
	})
	return err
}

// addStrings takes a batch whose first byte is an EvStrings event
// (indicating that the batch contains only strings) and adds each
// string contained therein to the provided strings map.
func addStrings(stringTable *dataTable[stringID, string], b batch) error {
	if !b.isStringsBatch() {
		return fmt.Errorf("internal error: addStrings called on non-string batch")
	}
	r := bytes.NewReader(b.data)
	hdr, err := r.ReadByte() // Consume the EvStrings byte.
	if err != nil || tracev2.EventType(hdr) != tracev2.EvStrings {
		return fmt.Errorf("missing strings batch header")
	}

	var sb strings.Builder
	for r.Len() != 0 {
		// Read the header.
		ev, err := r.ReadByte()
		if err != nil {
			return err
		}
		if tracev2.EventType(ev) != tracev2.EvString {
			return fmt.Errorf("expected string event, got %d", ev)
		}

		// Read the string's ID.
		id, err := binary.ReadUvarint(r)
		if err != nil {
			return err
		}

		// Read the string's length.
		len, err := binary.ReadUvarint(r)
		if err != nil {
			return err
		}
		if len > tracev2.MaxEventTrailerDataSize {
			return fmt.Errorf("invalid string size %d, maximum is %d", len, tracev2.MaxEventTrailerDataSize)
		}

		// Copy out the string.
		n, err := io.CopyN(&sb, r, int64(len))
		if n != int64(len) {
			return fmt.Errorf("failed to read full string: read %d but wanted %d", n, len)
		}
		if err != nil {
			return fmt.Errorf("copying string data: %w", err)
		}

		// Add the string to the map.
		s := sb.String()
		sb.Reset()
		if err := stringTable.insert(stringID(id), s); err != nil {
			return err
		}
	}
	return nil
}

// addStacks takes a batch whose first byte is an EvStacks event
// (indicating that the batch contains only stacks) and adds each
// string contained therein to the provided stacks map.
func addStacks(stackTable *dataTable[stackID, stack], pcs map[uint64]frame, b batch) error {
	if !b.isStacksBatch() {
		return fmt.Errorf("internal error: addStacks called on non-stacks batch")
	}
	r := bytes.NewReader(b.data)
	hdr, err := r.ReadByte() // Consume the EvStacks byte.
	if err != nil || tracev2.EventType(hdr) != tracev2.EvStacks {
		return fmt.Errorf("missing stacks batch header")
	}

	for r.Len() != 0 {
		// Read the header.
		ev, err := r.ReadByte()
		if err != nil {
			return err
		}
		if tracev2.EventType(ev) != tracev2.EvStack {
			return fmt.Errorf("expected stack event, got %d", ev)
		}

		// Read the stack's ID.
		id, err := binary.ReadUvarint(r)
		if err != nil {
			return err
		}

		// Read how many frames are in each stack.
		nFrames, err := binary.ReadUvarint(r)
		if err != nil {
			return err
		}
		if nFrames > tracev2.MaxFramesPerStack {
			return fmt.Errorf("invalid stack size %d, maximum is %d", nFrames, tracev2.MaxFramesPerStack)
		}

		// Each frame consists of 4 fields: pc, funcID (string), fileID (string), line.
		frames := make([]uint64, 0, nFrames)
		for i := uint64(0); i < nFrames; i++ {
			// Read the frame data.
			pc, err := binary.ReadUvarint(r)
			if err != nil {
				return fmt.Errorf("reading frame %d's PC for stack %d: %w", i+1, id, err)
			}
			funcID, err := binary.ReadUvarint(r)
			if err != nil {
				return fmt.Errorf("reading frame %d's funcID for stack %d: %w", i+1, id, err)
			}
			fileID, err := binary.ReadUvarint(r)
			if err != nil {
				return fmt.Errorf("reading frame %d's fileID for stack %d: %w", i+1, id, err)
			}
			line, err := binary.ReadUvarint(r)
			if err != nil {
				return fmt.Errorf("reading frame %d's line for stack %d: %w", i+1, id, err)
			}
			frames = append(frames, pc)

			if _, ok := pcs[pc]; !ok {
				pcs[pc] = frame{
					pc:     pc,
					funcID: stringID(funcID),
					fileID: stringID(fileID),
					line:   line,
				}
			}
		}

		// Add the stack to the map.
		if err := stackTable.insert(stackID(id), stack{pcs: frames}); err != nil {
			return err
		}
	}
	return nil
}

// addCPUSamples takes a batch whose first byte is an EvCPUSamples event
// (indicating that the batch contains only CPU samples) and adds each
// sample contained therein to the provided samples list.
func addCPUSamples(samples []cpuSample, b batch) ([]cpuSample, error) {
	if !b.isCPUSamplesBatch() {
		return nil, fmt.Errorf("internal error: addCPUSamples called on non-CPU-sample batch")
	}
	r := bytes.NewReader(b.data)
	hdr, err := r.ReadByte() // Consume the EvCPUSamples byte.
	if err != nil || tracev2.EventType(hdr) != tracev2.EvCPUSamples {
		return nil, fmt.Errorf("missing CPU samples batch header")
	}

	for r.Len() != 0 {
		// Read the header.
		ev, err := r.ReadByte()
		if err != nil {
			return nil, err
		}
		if tracev2.EventType(ev) != tracev2.EvCPUSample {
			return nil, fmt.Errorf("expected CPU sample event, got %d", ev)
		}

		// Read the sample's timestamp.
		ts, err := binary.ReadUvarint(r)
		if err != nil {
			return nil, err
		}

		// Read the sample's M.
		m, err := binary.ReadUvarint(r)
		if err != nil {
			return nil, err
		}
		mid := ThreadID(m)

		// Read the sample's P.
		p, err := binary.ReadUvarint(r)
		if err != nil {
			return nil, err
		}
		pid := ProcID(p)

		// Read the sample's G.
		g, err := binary.ReadUvarint(r)
		if err != nil {
			return nil, err
		}
		goid := GoID(g)
		if g == 0 {
			goid = NoGoroutine
		}

		// Read the sample's stack.
		s, err := binary.ReadUvarint(r)
		if err != nil {
			return nil, err
		}

		// Add the sample to the slice.
		samples = append(samples, cpuSample{
			schedCtx: schedCtx{
				M: mid,
				P: pid,
				G: goid,
			},
			time:  Time(ts), // N.B. this is really a "timestamp," not a Time.
			stack: stackID(s),
		})
	}
	return samples, nil
}

// sync holds the per-generation sync data.
type sync struct {
	freq             frequency
	hasClockSnapshot bool
	snapTime         timestamp
	snapMono         uint64
	snapWall         time.Time
}

func setSyncBatch(s *sync, b batch, ver version.Version) error {
	if !b.isSyncBatch(ver) {
		return fmt.Errorf("internal error: setSyncBatch called on non-sync batch")
	}
	r := bytes.NewReader(b.data)
	if ver >= version.Go125 {
		hdr, err := r.ReadByte() // Consume the EvSync byte.
		if err != nil || tracev2.EventType(hdr) != tracev2.EvSync {
			return fmt.Errorf("missing sync batch header")
		}
	}

	lastTs := b.time
	for r.Len() != 0 {
		// Read the header
		ev, err := r.ReadByte()
		if err != nil {
			return err
		}
		et := tracev2.EventType(ev)
		switch {
		case et == tracev2.EvFrequency:
			if s.freq != 0 {
				return fmt.Errorf("found multiple frequency events")
			}
			// Read the frequency. It'll come out as timestamp units per second.
			f, err := binary.ReadUvarint(r)
			if err != nil {
				return err
			}
			// Convert to nanoseconds per timestamp unit.
			s.freq = frequency(1.0 / (float64(f) / 1e9))
		case et == tracev2.EvClockSnapshot && ver >= version.Go125:
			if s.hasClockSnapshot {
				return fmt.Errorf("found multiple clock snapshot events")
			}
			s.hasClockSnapshot = true
			// Read the EvClockSnapshot arguments.
			tdiff, err := binary.ReadUvarint(r)
			if err != nil {
				return err
			}
			lastTs += timestamp(tdiff)
			s.snapTime = lastTs
			mono, err := binary.ReadUvarint(r)
			if err != nil {
				return err
			}
			s.snapMono = mono
			sec, err := binary.ReadUvarint(r)
			if err != nil {
				return err
			}
			nsec, err := binary.ReadUvarint(r)
			if err != nil {
				return err
			}
			// TODO(felixge): In theory we could inject s.snapMono into the time
			// value below to make it comparable. But there is no API for this
			// in the time package right now.
			s.snapWall = time.Unix(int64(sec), int64(nsec))
		default:
			return fmt.Errorf("expected frequency or clock snapshot event, got %d", ev)
		}
	}
	return nil
}

// addExperimentalBatch takes an experimental batch and adds it to the list of experimental
// batches for the experiment its a part of.
func addExperimentalBatch(expBatches map[tracev2.Experiment][]ExperimentalBatch, b batch) error {
	if b.exp == tracev2.NoExperiment {
		return fmt.Errorf("internal error: addExperimentalBatch called on non-experimental batch")
	}
	expBatches[b.exp] = append(expBatches[b.exp], ExperimentalBatch{
		Thread: b.m,
		Data:   b.data,
	})
	return nil
}

func fixUpCPUSamples(samples []cpuSample, freq frequency) {
	// Fix up the CPU sample timestamps.
	for i := range samples {
		s := &samples[i]
		s.time = freq.mul(timestamp(s.time))
	}
	// Sort the CPU samples.
	slices.SortFunc(samples, func(a, b cpuSample) int {
		return cmp.Compare(a.time, b.time)
	})
}