package main

import (
	"bufio"
	"cmp"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"maps"
	"os"
	"path/filepath"
	"slices"
	"strconv"
	"strings"
	"time"
)

const (
	// probeName is the agent name of a probe run. The cache key comes from
	// it.
	probeName = "Probe"

	// probeQuotaLimit is how many points the used percent of a quota window
	// can grow in a probe run before the run stops.
	probeQuotaLimit = 15.0

	// probeCallLimit is the most model calls of a probe run.
	probeCallLimit = 660

	// probeLineBytes is the most bytes of one line of a probe log.
	probeLineBytes = 1 << 20
)

// probe runs "onechat probe LOG COUNT". It appends the first COUNT messages of
// LOG, one JSON object {kind, text} per line, to a new chat in probe/ in the
// data directory, and builds its tree with the real compactor. Before each
// user message it waits for the builds, as a turn does. It stops at the first
// failed call, after probeCallLimit calls, or when the used percent of a quota
// window grows by more than probeQuotaLimit points. At the end it reports the
// calls, the time per call, the calls at once, the cached share, the quota,
// the retries, and the line sizes.
func probe(ctx context.Context, args []string, out io.Writer) error {
	if len(args) != 2 {
		return errors.New("usage: onechat probe LOG COUNT")
	}

	count, err := strconv.Atoi(args[1])
	if err != nil || count <= 0 {
		return fmt.Errorf("COUNT %q is not a positive number", args[1])
	}

	entries, err := readProbeLog(args[0], count)
	if err != nil {
		return err
	}

	dir, err := dataDir()
	if err != nil {
		return err
	}

	config, err := loadConfig(dir)
	if err != nil {
		return err
	}

	address, auth, err := modelEndpoint(config, dir)
	if err != nil {
		return err
	}

	root := filepath.Join(dir, "probe", time.Now().Format("20060102-150405"))
	err = os.MkdirAll(root, 0o700)
	if err != nil {
		return fmt.Errorf("create probe directory: %w", err)
	}

	usagePath := filepath.Join(root, "usage.jsonl")
	client := newClient(address, auth, "onechat-"+probeName, usagePath)
	client.LimitCalls(probeCallLimit)
	memory, err := openMemory(ctx, root, time.Now, newCompactor(client, probeName, config.Compactor))
	if err != nil {
		return err
	}

	_, err = fmt.Fprintf(out, "probe: %d messages into %s with %s at %s effort\n", len(entries), root, config.Compactor.Model, config.Compactor.Effort)
	if err != nil {
		return errors.Join(err, memory.Close())
	}

	run := probeRun{
		memory:    memory,
		client:    client,
		usagePath: usagePath,
	}

	appended, stop := run.Run(ctx, entries)
	sizes, err := modelLines(memory)
	err = errors.Join(err, memory.Close())
	if err != nil {
		return err
	}

	records, err := readUsageFile(usagePath)
	if err != nil {
		return err
	}

	return writeProbeReport(out, appended, stop, records, sizes)
}

// probeRun appends the messages of a probe run to its chat and stops the run
// by the rules of probeStop.
type probeRun struct {
	memory    *Memory
	client    *Client
	usagePath string
}

// Run appends the entries. It waits for the builds only before a user message,
// as a turn does, so up to maxBuilds builds run at once in between, and it
// waits for all builds at the end. It returns the number of messages appended
// and why the run stopped early, or "" when it did not. After a stop it waits
// for the calls in flight, so no sent call is cancelled.
func (r *probeRun) Run(ctx context.Context, entries []probeEntry) (int, string) {
	appended, reason := r.append(ctx, entries)
	if reason == "" {
		return appended, ""
	}

	err := r.memory.Wait(ctx)
	if err != nil {
		reason += "; " + err.Error()
	}

	return appended, reason
}

// append appends the entries and stops at the first stop rule.
func (r *probeRun) append(ctx context.Context, entries []probeEntry) (int, string) {
	for n, entry := range entries {
		if entry.Kind == "user" {
			reason := r.wait(ctx)
			if reason != "" {
				return n, reason
			}
		}

		_, err := r.memory.Append(entry.Kind, entry.Text)
		if err != nil {
			return n, err.Error()
		}

		reason := r.check()
		if reason != "" {
			return n + 1, reason
		}
	}

	return len(entries), r.wait(ctx)
}

// wait returns when no build runs, or when a stop rule ends the run. It checks
// the rules each second, so a run stops soon also during a long wait.
func (r *probeRun) wait(ctx context.Context) string {
	for {
		reason := r.check()
		if reason != "" {
			return reason
		}

		waitCtx, cancel := context.WithTimeout(ctx, time.Second)
		err := r.memory.Wait(waitCtx)
		cancel()
		if err == nil {
			return r.check()
		}

		if ctx.Err() != nil {
			return ctx.Err().Error()
		}
	}
}

// check returns why the run must stop, or "" when it can go on.
func (r *probeRun) check() string {
	if r.client.LimitReached() {
		return fmt.Sprintf("call limit of %d reached", r.client.limit)
	}

	records, err := readUsageFile(r.usagePath)
	if err != nil {
		return err.Error()
	}

	return probeStop(records)
}

// probeEntry is one message of a probe log.
type probeEntry struct {
	Kind string `json:"kind"`
	Text string `json:"text"`
}

// readProbeLog reads the first count messages of the probe log at path.
func readProbeLog(path string, count int) (entries []probeEntry, err error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, fmt.Errorf("open probe log: %w", err)
	}

	defer func() {
		err = errors.Join(err, file.Close())
	}()

	scanner := bufio.NewScanner(file)
	scanner.Buffer(make([]byte, 64*1024), probeLineBytes)
	for len(entries) < count && scanner.Scan() {
		var entry probeEntry
		err = json.Unmarshal(scanner.Bytes(), &entry)
		if err != nil {
			return nil, fmt.Errorf("probe log line %d: %w", len(entries)+1, err)
		}

		entries = append(entries, entry)
	}

	err = scanner.Err()
	if err != nil {
		return nil, fmt.Errorf("read probe log: %w", err)
	}

	if len(entries) < count {
		return nil, fmt.Errorf("probe log has %d messages, fewer than %d", len(entries), count)
	}

	return entries, nil
}

// probeStop returns why a probe run must stop: a failed call, or a quota window
// whose used percent grew by more than probeQuotaLimit points since the first
// call. It returns "" when the run can go on.
func probeStop(records []usageRecord) string {
	for _, record := range records {
		if record.Error != "" {
			return "call failed: " + record.Error
		}
	}

	if len(records) == 0 {
		return ""
	}

	first := usedPercents(records[0])
	for name, used := range usedPercents(records[len(records)-1]) {
		start, found := first[name]
		if found && used-start > probeQuotaLimit {
			return fmt.Sprintf("%s grew from %g to %g", name, start, used)
		}
	}

	return ""
}

// modelLines returns the sizes of the nodes that the model wrote, in tree
// order. A node that fits its source with no model call is not one of them.
func modelLines(memory *Memory) ([]int, error) {
	memory.mutex.Lock()
	defer memory.mutex.Unlock()

	keys, err := memory.tree.Keys()
	if err != nil {
		return nil, err
	}

	var sizes []int
	for _, key := range keys {
		node, _, err := memory.tree.Node(key.level, key.index)
		if err != nil {
			return nil, err
		}

		model, err := builtByModel(memory, key)
		if err != nil {
			return nil, err
		}

		if model {
			sizes = append(sizes, node.Size)
		}
	}

	return sizes, nil
}

// builtByModel tells if the node key needed a model call: a leaf whose source
// is over nodeBytes, or a merge whose children do not fit in one node.
func builtByModel(memory *Memory, key nodeKey) (bool, error) {
	if key.level == 0 {
		message, err := memory.log.Message(key.index)
		if err != nil {
			return false, err
		}

		_, free := freeLeaf(message)

		return !free, nil
	}

	left, right, err := memory.children(key)
	if err != nil {
		return false, err
	}

	_, free := freeMerge(left, right)

	return !free, nil
}

// readUsageFile reads the records of usage.jsonl at path.
func readUsageFile(path string) (records []usageRecord, err error) {
	file, err := os.Open(path)
	if errors.Is(err, os.ErrNotExist) {
		return nil, nil
	}

	if err != nil {
		return nil, fmt.Errorf("open usage file: %w", err)
	}

	defer func() {
		err = errors.Join(err, file.Close())
	}()

	scanner := bufio.NewScanner(file)
	scanner.Buffer(make([]byte, 64*1024), probeLineBytes)
	for scanner.Scan() {
		var record usageRecord
		err = json.Unmarshal(scanner.Bytes(), &record)
		if err != nil {
			return nil, fmt.Errorf("decode usage: %w", err)
		}

		records = append(records, record)
	}

	err = scanner.Err()
	if err != nil {
		return nil, fmt.Errorf("read usage file: %w", err)
	}

	return records, nil
}

// usedPercents returns the used percent of each quota window of record, by
// header name.
func usedPercents(record usageRecord) map[string]float64 {
	used := make(map[string]float64)
	for name, value := range record.Quota {
		if !strings.HasSuffix(name, "-used-percent") {
			continue
		}

		percent, err := strconv.ParseFloat(value, 64)
		if err == nil {
			used[name] = percent
		}
	}

	return used
}

// writeProbeReport writes the result of a probe run.
func writeProbeReport(out io.Writer, appended int, stop string, records []usageRecord, sizes []int) error {
	var input, cached, output, failed int
	for _, record := range records {
		input += record.InputTokens
		cached += record.CachedTokens
		output += record.OutputTokens
		if record.Error != "" {
			failed++
		}
	}

	var lines []string
	if stop != "" {
		lines = append(lines, "stopped early: "+stop)
	}

	lines = append(lines,
		fmt.Sprintf("messages: %d", appended),
		fmt.Sprintf("calls: %d (%d failed), %.2f per message", len(records), failed, ratio(len(records), appended)),
		fmt.Sprintf("model lines: %d, retries: %d, %.2f calls per model line", len(sizes), len(records)-failed-len(sizes), ratio(len(records), len(sizes))),
		fmt.Sprintf("tokens: %d input, %d cached (%.1f%%), %d output", input, cached, 100*ratio(cached, input), output),
	)

	if len(sizes) > 0 {
		lines = append(lines, sizeText(sizes))
	}

	if len(records) > 0 {
		lines = append(lines, timeText(records))
	}

	if len(records) > 0 {
		lines = append(lines,
			"quota at the first call: "+quotaText(records[0]),
			"quota at the last call: "+quotaText(records[len(records)-1]),
		)
	}

	_, err := io.WriteString(out, strings.Join(lines, "\n")+"\n")

	return err
}

// sizeText describes the sizes of the model lines.
func sizeText(sizes []int) string {
	sorted := slices.Sorted(slices.Values(sizes))
	total, over := 0, 0
	for _, size := range sorted {
		total += size
		if size > nodeBytes {
			over++
		}
	}

	return fmt.Sprintf("line bytes: min %d, median %d, mean %.0f, max %d, %d over %d", sorted[0], sorted[len(sorted)/2], ratio(total, len(sorted)), sorted[len(sorted)-1], over, nodeBytes)
}

// timeText describes the time per call and how many calls ran at once. A call
// runs from its date for its seconds.
func timeText(records []usageRecord) string {
	type event struct {
		at    time.Time
		delta int
	}

	var events []event
	seconds := make([]float64, 0, len(records))
	total := 0.0
	for _, record := range records {
		start, err := time.Parse(time.RFC3339Nano, record.Date)
		if err != nil {
			continue
		}

		length := time.Duration(record.Seconds * float64(time.Second))
		events = append(events, event{
			at:    start,
			delta: 1,
		}, event{
			at:    start.Add(length),
			delta: -1,
		})
		seconds = append(seconds, record.Seconds)
		total += record.Seconds
	}

	if len(events) == 0 {
		return "time: no dated calls"
	}

	slices.SortFunc(events, func(a, b event) int {
		return cmp.Or(a.at.Compare(b.at), cmp.Compare(a.delta, b.delta))
	})

	running, most := 0, 0
	for _, e := range events {
		running += e.delta
		most = max(most, running)
	}

	wall := events[len(events)-1].at.Sub(events[0].at).Seconds()
	slices.Sort(seconds)

	return fmt.Sprintf("time: %.0f s in all, %.1f s per call (median %.1f, max %.1f), calls at once: at most %d, %.1f on average", wall, total/float64(len(seconds)), seconds[len(seconds)/2], seconds[len(seconds)-1], most, total/max(wall, 1e-9))
}

// quotaText returns the quota headers of record, sorted by name.
func quotaText(record usageRecord) string {
	names := slices.Sorted(maps.Keys(record.Quota))
	parts := make([]string, len(names))
	for n, name := range names {
		parts[n] = strings.TrimPrefix(name, "x-codex-") + "=" + record.Quota[name]
	}

	return strings.Join(parts, " ")
}

// ratio returns a / b, or 0 when b is 0.
func ratio(a, b int) float64 {
	if b == 0 {
		return 0
	}

	return float64(a) / float64(b)
}