package main

import (
	"encoding/json"
	"fmt"
	"math/rand/v2"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

// probeMessages is the length of the generated log for the live probe.
const probeMessages = 1000

// TestProbeLogHasTheSyntheticMix checks that the generated log has readable
// ASCII text with the mix of kinds and sizes of the synthetic log. When
// ONECHAT_PROBE_LOG names a file, the test also writes the log there, one
// message {kind, text} per line, for "onechat probe".
func TestProbeLogHasTheSyntheticMix(t *testing.T) {
	messages := probeLog(probeMessages)
	require.Len(t, messages, probeMessages)

	counts := make(map[string]int)
	long := 0
	for _, message := range messages {
		counts[message.Kind]++
		require.True(t, logKinds[message.Kind], message.Kind)
		require.NotEmpty(t, message.Text)
		for _, char := range message.Text {
			require.True(t, char == '\n' || (char >= ' ' && char <= '~'), "character %q", char)
		}

		if len(message.Text) >= 30000 {
			long++
		}

		if message.Kind == "work" {
			require.True(t, strings.HasPrefix(message.Text, "["), message.Text)
		}
	}

	require.Greater(t, counts["user"], 20)
	require.Greater(t, counts["tool"], 300)
	require.Greater(t, counts["echo"], 300)
	require.Greater(t, long, 2)
	require.Equal(t, messages, probeLog(probeMessages))

	path := os.Getenv("ONECHAT_PROBE_LOG")
	if path == "" {
		return
	}

	var data strings.Builder
	for _, message := range messages {
		line, err := json.Marshal(message)
		require.NoError(t, err)

		data.Write(line)
		data.WriteByte('\n')
	}

	require.NoError(t, os.WriteFile(path, []byte(data.String()), 0o600))
}

func TestProbeStopRules(t *testing.T) {
	quota := func(primary, secondary string) ModelUsageRecord {
		return ModelUsageRecord{
			Quota: map[string]string{
				"x-codex-primary-used-percent":   primary,
				"x-codex-secondary-used-percent": secondary,
				"x-codex-plan-type":              "plus",
			},
		}
	}

	testCases := []struct {
		name    string
		records []ModelUsageRecord
		want    string
	}{
		{
			name:    "no calls",
			records: nil,
			want:    "",
		},
		{
			name:    "growth within the limit",
			records: []ModelUsageRecord{quota("10", "40"), quota("25", "41")},
			want:    "",
		},
		{
			name:    "growth over the limit",
			records: []ModelUsageRecord{quota("10", "40"), quota("12", "55.5")},
			want:    "x-codex-secondary-used-percent grew from 40 to 55.5",
		},
		{
			name: "failed call",
			records: []ModelUsageRecord{quota("10", "40"), {
				Status: 400,
				Error:  "model call: HTTP 400",
			}},
			want: "call failed: model call: HTTP 400",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			require.Equal(t, tc.want, probeStop(tc.records))
		})
	}
}

func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) {
	server := newModelFake(t)
	server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
		w.Header().Set("X-Codex-Primary-Used-Percent", strconv.Itoa(10+call.number))
		modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem(strings.Repeat("s", 400))}, 1000, 800, 50))
	}

	run := newTestProbeRun(t, server)
	appended, stop := run.Run(t.Context(), probeLog(probeMessages))
	sizes, err := probeModelLines(run.memory)
	require.NoError(t, err)
	require.NoError(t, run.memory.Close())

	require.Regexp(t, `^x-codex-primary-used-percent grew from \d+ to \d+$`, stop)
	require.Less(t, appended, probeMessages)
	require.NotEmpty(t, sizes)

	records, err := probeReadUsage(run.usagePath)
	require.NoError(t, err)
	require.Greater(t, len(records), int(probeQuotaLimit))

	var report strings.Builder
	require.NoError(t, probeWriteReport(&report, appended, stop, records, sizes))
	require.Contains(t, report.String(), "stopped early: ")
	require.Contains(t, report.String(), "cached (80.0%)")
	require.Contains(t, report.String(), "line bytes: min 400, median 400, mean 400, max 400, 0 over 512")
	require.Contains(t, report.String(), "calls at once: at most ")
	require.Contains(t, report.String(), "quota at the first call: primary-used-percent=")
}

func TestProbeRunStopsAtTheCallLimit(t *testing.T) {
	server := newModelFake(t)
	run := newTestProbeRun(t, server)
	run.client.LimitCalls(5)

	entries := make([]ProbeEntry, 40)
	for n := range entries {
		entries[n] = ProbeEntry{
			Kind: "echo",
			Text: strings.Repeat("long output ", 100),
		}
	}

	_, stop := run.Run(t.Context(), entries)
	require.NoError(t, run.memory.Close())
	require.Equal(t, "call limit of 5 reached", stop)
	require.Equal(t, 5, server.Calls())
}

func TestProbeRunWaitsOnlyBeforeUserMessages(t *testing.T) {
	gate := make(chan struct{})
	var mutex sync.Mutex
	running, most := 0, 0
	server := newModelFake(t)
	server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
		mutex.Lock()
		running++
		most = max(most, running)
		mutex.Unlock()

		<-gate

		mutex.Lock()
		running--
		mutex.Unlock()

		modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("a line")}, 100, 0, 5))
	}

	run := newTestProbeRun(t, server)
	long := strings.Repeat("long output ", 100)
	entries := []ProbeEntry{
		{
			Kind: "echo",
			Text: long,
		},
		{
			Kind: "tool",
			Text: long,
		},
		{
			Kind: "user",
			Text: "next",
		},
		{
			Kind: "reply",
			Text: long,
		},
	}

	done := make(chan string)
	go func() {
		_, stop := run.Run(t.Context(), entries)
		done <- stop
	}()

	require.Eventually(t, func() bool {
		return server.Calls() == 2
	}, 10*time.Second, time.Millisecond)

	run.memory.mutex.Lock()
	appended := run.memory.log.Len()
	run.memory.mutex.Unlock()
	require.Equal(t, 2, appended)

	close(gate)
	require.Empty(t, <-done)
	require.NoError(t, run.memory.Close())

	mutex.Lock()
	defer mutex.Unlock()

	require.Equal(t, 2, most)
	require.Equal(t, 3, server.Calls())
}

func TestSlotSummarizerLimitsTheBuildsAtOnce(t *testing.T) {
	inner := newFakeSummarizer()
	inner.gate = make(chan struct{})
	summarizer := ProbeSlotSummarizer{
		inner: inner,
		slots: make(chan struct{}, 2),
	}

	var group sync.WaitGroup
	for n := range 6 {
		group.Go(func() {
			_, err := summarizer.Summarize(t.Context(), Task{
				Level: 0,
				Index: n,
				Input: []string{"user: " + strings.Repeat("x", 600)},
			})
			require.NoError(t, err)
		})
	}

	require.Eventually(t, func() bool {
		return inner.Active() == 2
	}, 10*time.Second, time.Millisecond)

	close(inner.gate)
	group.Wait()
	require.Equal(t, 2, inner.peak)
}

func TestReadUsageFileSortsByCallStart(t *testing.T) {
	path := filepath.Join(t.TempDir(), "usage.jsonl")
	usage := ModelUsageLog{
		path: path,
	}

	for _, date := range []string{"2026-01-01T12:00:02Z", "2026-01-01T12:00:00.5Z", "2026-01-01T12:00:01.25Z"} {
		require.NoError(t, usage.Write(ModelUsageRecord{
			Date: date,
		}))
	}

	records, err := probeReadUsage(path)
	require.NoError(t, err)
	require.Equal(t, []string{"2026-01-01T12:00:00.5Z", "2026-01-01T12:00:01.25Z", "2026-01-01T12:00:02Z"}, []string{records[0].Date, records[1].Date, records[2].Date})
}

func TestTimeTextCountsTheCallsAtOnce(t *testing.T) {
	start := time.Date(2026, 1, 1, 12, 0, 0, 0, time.UTC)
	record := func(offset, seconds float64) ModelUsageRecord {
		return ModelUsageRecord{
			Date:    start.Add(time.Duration(offset * float64(time.Second))).Format(time.RFC3339Nano),
			Seconds: seconds,
		}
	}

	text := probeTimeText([]ModelUsageRecord{record(0, 10), record(2, 4), record(3, 2), record(10, 10)})
	require.Equal(t, "time: 20 s in all, 12.0 calls per minute, 6.5 s per call (median 10.0, max 10.0), calls at once: at most 3, 1.3 on average", text)
}

// ProbeWriter writes readable made-up messages.
type ProbeWriter struct {
	random *rand.Rand
}

func (g *ProbeWriter) message() (string, string) {
	roll := g.random.IntN(100)
	switch {
	case roll < 5:
		return "user", g.fill(g.userLine, 10+g.random.IntN(200))
	case roll < 10:
		return "reply", g.fill(g.replyLine, 100+g.random.IntN(1500))
	case roll < 12:
		return "note", g.fill(g.noteLine, 50+g.random.IntN(600))
	case roll < 13:
		return "work", "[" + g.pick(probePeople) + "] report: " + g.fill(g.replyLine, 300+g.random.IntN(3000))
	case roll < 56:
		return "tool", g.toolCall(60 + g.random.IntN(400))
	case roll < 57:
		return "echo", g.fill(g.outputLine, 30000+g.random.IntN(30000))
	default:
		return "echo", g.fill(g.outputLine, 20+g.random.IntN(4000))
	}
}

// fill joins lines of write until the text has about size bytes.
func (g *ProbeWriter) fill(write func() string, size int) string {
	var text strings.Builder
	for text.Len() < size {
		if text.Len() > 0 {
			text.WriteByte(' ')
		}

		text.WriteString(write())
	}

	return text.String()
}

func (g *ProbeWriter) userLine() string {
	return g.pick([]string{
		fmt.Sprintf("Please check why the %s %s.", g.pick(probeThings), g.pick(probeProblems)),
		fmt.Sprintf("Move the %s to %s before %s.", g.pick(probeThings), g.pick(probePlaces), g.pick(probeDays)),
		fmt.Sprintf("Thanks, that works. Next, %s.", g.pick(probeTasks)),
		fmt.Sprintf("No, keep the old %s. %s asked for it.", g.pick(probeThings), g.pick(probePeople)),
		fmt.Sprintf("How many %s orders came in on %s?", g.pick(probeBreads), g.pick(probeDays)),
	})
}

func (g *ProbeWriter) replyLine() string {
	return g.pick([]string{
		fmt.Sprintf("I looked at %s. The %s %s because %s.", g.pick(probeFiles), g.pick(probeThings), g.pick(probeProblems), g.pick(probeCauses)),
		fmt.Sprintf("I changed %s, and the tests pass now.", g.pick(probeChanges)),
		fmt.Sprintf("There were %d %s orders on %s, %d of them late.", 10+g.random.IntN(90), g.pick(probeBreads), g.pick(probeDays), g.random.IntN(10)),
		fmt.Sprintf("Open question: should %s get %s?", g.pick(probePeople), g.pick(probeTasks)),
	})
}

func (g *ProbeWriter) noteLine() string {
	return g.pick([]string{
		fmt.Sprintf("%s prefers %s.", g.pick(probePeople), g.pick(probePreferences)),
		fmt.Sprintf("The bakery opens at %d and closes at %d on %s.", 5+g.random.IntN(3), 17+g.random.IntN(4), g.pick(probeDays)),
	})
}

func (g *ProbeWriter) pick(words []string) string {
	return words[g.random.IntN(len(words))]
}

// toolCall returns a tool name and its JSON input of about size bytes.
func (g *ProbeWriter) toolCall(size int) string {
	name := g.pick([]string{"shell", "read", "write", "edit"})
	var input map[string]string
	switch name {
	case "shell":
		input = map[string]string{"command": "grep -n '" + g.pick(probeThings) + "' " + g.pick(probeFiles)}
	case "read":
		input = map[string]string{"path": g.pick(probeFiles)}
	default:
		input = map[string]string{
			"path": g.pick(probeFiles),
			"text": g.fill(g.replyLine, max(0, size-80)),
		}
	}

	data, err := json.Marshal(input)
	if err != nil {
		panic(err)
	}

	return name + " " + string(data)
}

func (g *ProbeWriter) outputLine() string {
	return g.pick([]string{
		fmt.Sprintf("%02d:%02d order %d for %d %s, packed by %s.", 6+g.random.IntN(12), g.random.IntN(60), 1000+g.random.IntN(9000), 1+g.random.IntN(24), g.pick(probeBreads), g.pick(probePeople)),
		fmt.Sprintf("%s:%d: %s %s.", g.pick(probeFiles), 1+g.random.IntN(400), g.pick(probeThings), g.pick(probeProblems)),
		fmt.Sprintf("ok   bakery/%s %d.%03ds", g.pick(probeModules), g.random.IntN(3), g.random.IntN(1000)),
	})
}

// probeLog returns count messages of made-up work on a small bakery order
// system, with the mix of kinds and sizes of replayMessage. The same count
// always gives the same log.
func probeLog(count int) []ProbeEntry {
	g := ProbeWriter{
		random: rand.New(rand.NewPCG(7, 11)),
	}

	messages := make([]ProbeEntry, count)
	for id := range messages {
		kind, text := g.message()
		messages[id] = ProbeEntry{
			Kind: kind,
			Text: text,
		}
	}

	return messages
}

func newTestProbeRun(t *testing.T, server *ModelFake) *ProbeRun {
	t.Helper()

	root := t.TempDir()
	usagePath := filepath.Join(root, "usage.jsonl")
	client := newModelClient(server.URL, ModelAPIKey(modelTestKey), "onechat-"+probeName, usagePath)
	memory, err := openMemory(t.Context(), root, newTestClock().Now, newCompactor(client, probePrompt, defaultConfig().Compactor), discardLog())
	require.NoError(t, err)

	return &ProbeRun{
		memory:    memory,
		client:    client,
		usagePath: usagePath,
	}
}

var (
	probePeople      = []string{"Mara", "Joel", "Ines", "Tomas", "Priya", "Lena"}
	probeBreads      = []string{"rye", "sourdough", "baguette", "brioche", "spelt", "seed loaf"}
	probeDays        = []string{"Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday"}
	probeThings      = []string{"order form", "price list", "delivery route", "stock count", "invoice export", "oven schedule"}
	probeProblems    = []string{"shows the wrong total", "skips the last row", "is slow at noon", "loses the customer note", "rounds the tax down"}
	probePlaces      = []string{"the front page", "the weekly report", "the shared folder", "the admin menu"}
	probeTasks       = []string{"add a field for allergies", "send the invoices on Friday", "split the route in two", "print labels for the shelf"}
	probeCauses      = []string{"the date has no time zone", "the cache keeps old prices", "the loop stops one row early", "the form trims the note"}
	probeChanges     = []string{"the total to use cents", "the route to sort by street", "the export to add the tax line", "the form to keep the note"}
	probePreferences = []string{"short answers", "a call before a change to prices", "invoices as PDF", "reports in the morning"}
	probeFiles       = []string{"orders/form.go", "prices/list.go", "routes/plan.go", "stock/count.go", "export/invoice.go", "ovens/schedule.go"}
	probeModules     = []string{"orders", "prices", "routes", "stock", "export", "ovens"}
)