package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"

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

// modelTestKey is a fake API key.
const modelTestKey = "sk-test-not-a-real-key"

func TestClientSendsTheFixedRequestFields(t *testing.T) {
	server := newModelFake(t)
	client := newTestModelClient(t, server)

	_, err := client.Create(t.Context(), ModelRequest{
		Model:        "gpt-test",
		Effort:       "high",
		Instructions: "You are a test.",
		Tools: []ModelTool{{
			Type:       "function",
			Name:       "date",
			Parameters: map[string]any{"type": "object"},
		}},
		Input: []any{modelUserMessage("<chat>\n</chat>\n", "hello")},
	})
	require.NoError(t, err)

	call := server.Call(0)
	require.Equal(t, "Bearer "+modelTestKey, call.header.Get("Authorization"))
	require.Equal(t, "text/event-stream", call.header.Get("Accept"))
	require.Equal(t, "application/json", call.header.Get("Content-Type"))

	body := call.body
	require.Equal(t, "gpt-test", body["model"])
	require.Equal(t, "You are a test.", body["instructions"])
	require.Equal(t, false, body["store"])
	require.Equal(t, true, body["stream"])
	require.Equal(t, []any{"reasoning.encrypted_content"}, body["include"])
	require.Equal(t, "agent-key", body["prompt_cache_key"])
	require.Equal(t, map[string]any{"effort": "high"}, body["reasoning"])
	require.NotContains(t, body, "max_output_tokens")
	require.NotContains(t, body, "metadata")
	require.Equal(t, []any{map[string]any{
		"role": "user",
		"content": []any{
			map[string]any{"type": "input_text", "text": "<chat>\n</chat>\n"},
			map[string]any{"type": "input_text", "text": "hello"},
		},
	}}, body["input"])
	require.Len(t, body["tools"], 1)
}

func TestClientSendsTheSessionHeaderToChatGPT(t *testing.T) {
	login := newChatGPTFake(t)
	auth := newTestChatGPTAuth(t, login)
	chatgptWriteTestTokens(t, auth, login.token(time.Hour), "refresh-0")

	testCases := []struct {
		name string
		auth ModelAuthorizer
		want string
	}{
		{
			name: "ChatGPT backend",
			auth: auth,
			want: "agent-key",
		},
		{
			name: "OpenAI API",
			auth: ModelAPIKey(modelTestKey),
			want: "",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			server := newModelFake(t)
			client := newModelClient(server.URL, tc.auth, "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
			_, err := client.Create(t.Context(), modelTestRequest())
			require.NoError(t, err)

			call := server.Call(0)
			require.Equal(t, tc.want, call.header.Get("session_id"))
			require.Equal(t, "agent-key", call.body["prompt_cache_key"])
		})
	}
}

func TestClientReadsTheStream(t *testing.T) {
	server := newModelFake(t)
	server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
		modelWriteEvents(w,
			map[string]any{"type": "response.created", "response": map[string]any{}},
			map[string]any{"type": "response.output_item.done", "item": map[string]any{"type": "reasoning", "encrypted_content": "opaque"}},
			map[string]any{"type": "response.output_text.delta", "delta": "Hel"},
			map[string]any{"type": "response.output_item.done", "item": modelMessageItem("Hello")},
			map[string]any{"type": "response.output_item.done", "item": modelMessageItem(" world")},
			modelCompletedEvent(nil, 1200, 1024, 30),
		)
	}

	response, err := newTestModelClient(t, server).Create(t.Context(), modelTestRequest())
	require.NoError(t, err)
	require.Equal(t, "Hello world", response.Text)
	require.Len(t, response.Output, 3)
	require.JSONEq(t, `{"type": "reasoning", "encrypted_content": "opaque"}`, string(response.Output[0]))
	require.Equal(t, ModelUsage{
		Input:  1200,
		Cached: 1024,
		Output: 30,
	}, response.Usage)
}

func TestClientTakesTheOutputOfTheCompletedResponse(t *testing.T) {
	server := newModelFake(t)
	server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
		modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("only here")}, 10, 0, 2))
	}

	response, err := newTestModelClient(t, server).Create(t.Context(), modelTestRequest())
	require.NoError(t, err)
	require.Equal(t, "only here", response.Text)
}

func TestClientRecordsEachCall(t *testing.T) {
	server := newModelFake(t)
	server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
		w.Header().Set("X-Codex-Primary-Used-Percent", "12")
		w.Header().Set("X-Codex-Plan-Type", "plus")
		w.Header().Set("X-Other", "not quota")
		w.Header().Set("X-Codex-Turn-State", "opaque state")
		if call.number == 1 {
			w.WriteHeader(http.StatusTooManyRequests)
			_, _ = io.WriteString(w, `{"detail": "slow down"}`)

			return
		}

		modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("ok")}, 500, 400, 7))
	}

	client := newTestModelClient(t, server)
	request := modelTestRequest()
	request.Node = &nodeKey{2, 5}
	_, err := client.Create(t.Context(), request)
	require.NoError(t, err)
	require.NotContains(t, server.Call(0).body, "node")
	require.NotContains(t, string(server.Call(0).raw), "[2,5]")

	_, err = client.Create(t.Context(), modelTestRequest())
	require.ErrorContains(t, err, "HTTP 429")
	require.ErrorContains(t, err, "slow down")

	records := modelReadUsage(t, client.usage.path)
	require.Len(t, records, 2)

	quota := map[string]string{
		"x-codex-primary-used-percent": "12",
		"x-codex-plan-type":            "plus",
	}
	require.Equal(t, "gpt-test", records[0].Model)
	require.Equal(t, http.StatusOK, records[0].Status)
	require.Equal(t, 500, records[0].InputTokens)
	require.Equal(t, 400, records[0].CachedTokens)
	require.Equal(t, 7, records[0].OutputTokens)
	require.Equal(t, quota, records[0].Quota)
	require.Equal(t, &[2]int{2, 5}, records[0].Node)
	require.Equal(t, len("ok"), records[0].TextBytes)
	require.Greater(t, records[0].Seconds, 0.0)
	require.Empty(t, records[0].Error)
	require.NotEmpty(t, records[0].Date)

	require.Nil(t, records[1].Node)
	require.Zero(t, records[1].TextBytes)

	require.Equal(t, http.StatusTooManyRequests, records[1].Status)
	require.Equal(t, quota, records[1].Quota)
	require.Contains(t, records[1].Error, "HTTP 429")

	data, err := os.ReadFile(client.usage.path)
	require.NoError(t, err)
	require.NotContains(t, string(data), modelTestKey)

	info, err := os.Stat(client.usage.path)
	require.NoError(t, err)
	require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
}

func TestClientRefusesCallsOverTheLimit(t *testing.T) {
	server := newModelFake(t)
	client := newTestModelClient(t, server)
	client.LimitCalls(2)

	for range 2 {
		require.False(t, client.LimitReached())
		_, err := client.Create(t.Context(), modelTestRequest())
		require.NoError(t, err)
	}

	require.True(t, client.LimitReached())
	_, err := client.Create(t.Context(), modelTestRequest())
	require.ErrorIs(t, err, errModelCallLimit)
	require.Equal(t, 2, server.Calls())
	require.Len(t, modelReadUsage(t, client.usage.path), 2)
}

func TestClientFailsWhenTheResponseFails(t *testing.T) {
	testCases := []struct {
		name   string
		events []any
		want   string
	}{
		{
			name:   "failed response",
			events: []any{map[string]any{"type": "response.failed", "response": map[string]any{"error": map[string]any{"code": "server_error", "message": "busy"}}}},
			want:   "server_error: busy",
		},
		{
			name:   "incomplete response",
			events: []any{map[string]any{"type": "response.incomplete", "response": map[string]any{"incomplete_details": map[string]any{"reason": "content_filter"}}}},
			want:   "content_filter",
		},
		{
			name:   "error event",
			events: []any{map[string]any{"type": "error", "code": "rate_limit", "message": "too many"}},
			want:   "rate_limit: too many",
		},
		{
			name:   "stream ends early",
			events: []any{map[string]any{"type": "response.output_item.done", "item": modelMessageItem("cut")}},
			want:   "ended before the response completed",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			server := newModelFake(t)
			server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
				modelWriteEvents(w, tc.events...)
			}

			client := newTestModelClient(t, server)
			_, err := client.Create(t.Context(), modelTestRequest())
			require.ErrorContains(t, err, tc.want)

			records := modelReadUsage(t, client.usage.path)
			require.Len(t, records, 1)
			require.Contains(t, records[0].Error, tc.want)
		})
	}
}

// ModelFakeCall is one call to a fake endpoint, numbered from 0.
type ModelFakeCall struct {
	number int
	header http.Header
	raw    []byte
	body   map[string]any
}

// ModelFake is an httptest Responses endpoint. It keeps each call, and reply
// writes the response. The default reply is one short message.
type ModelFake struct {
	*httptest.Server
	mutex sync.Mutex
	calls []ModelFakeCall
	reply func(http.ResponseWriter, ModelFakeCall)
}

// Call returns call number n.
func (f *ModelFake) Call(n int) ModelFakeCall {
	f.mutex.Lock()
	defer f.mutex.Unlock()

	return f.calls[n]
}

// Calls returns the number of calls.
func (f *ModelFake) Calls() int {
	f.mutex.Lock()
	defer f.mutex.Unlock()

	return len(f.calls)
}

func (f *ModelFake) serve(w http.ResponseWriter, r *http.Request) {
	raw, err := io.ReadAll(r.Body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	call := ModelFakeCall{
		header: r.Header.Clone(),
		raw:    raw,
	}

	err = json.Unmarshal(raw, &call.body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	f.mutex.Lock()
	call.number = len(f.calls)
	f.calls = append(f.calls, call)
	reply := f.reply
	f.mutex.Unlock()

	reply(w, call)
}

func newModelFake(t *testing.T) *ModelFake {
	t.Helper()

	f := &ModelFake{
		reply: func(w http.ResponseWriter, call ModelFakeCall) {
			modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("a line")}, 100, 0, 5))
		},
	}

	f.Server = httptest.NewServer(http.HandlerFunc(f.serve))
	t.Cleanup(f.Close)

	return f
}

func newTestModelClient(t *testing.T, server *ModelFake) *ModelClient {
	t.Helper()

	return newModelClient(server.URL, ModelAPIKey(modelTestKey), "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
}

func modelTestRequest() ModelRequest {
	return ModelRequest{
		Model:  "gpt-test",
		Effort: "low",
		Input:  []any{modelUserMessage("hello")},
	}
}

// modelWriteEvents writes events as a server-sent event stream.
func modelWriteEvents(w http.ResponseWriter, events ...any) {
	w.Header().Set("Content-Type", "text/event-stream")
	for _, event := range events {
		data, err := json.Marshal(event)
		if err != nil {
			panic(err)
		}

		_, err = fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event.(map[string]any)["type"], data)
		if err != nil {
			panic(err)
		}
	}
}

func modelMessageItem(text string) map[string]any {
	return map[string]any{
		"type": "message",
		"role": "assistant",
		"content": []any{
			map[string]any{"type": "output_text", "text": text},
		},
	}
}

func modelCompletedEvent(output []any, input, cached, out int) map[string]any {
	if output == nil {
		output = []any{}
	}

	return map[string]any{
		"type": "response.completed",
		"response": map[string]any{
			"output": output,
			"usage": map[string]any{
				"input_tokens":         input,
				"input_tokens_details": map[string]any{"cached_tokens": cached},
				"output_tokens":        out,
			},
		},
	}
}

func modelReadUsage(t *testing.T, path string) []ModelUsageRecord {
	t.Helper()

	file, err := os.Open(path)
	require.NoError(t, err)

	defer func() {
		require.NoError(t, file.Close())
	}()

	var records []ModelUsageRecord
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		var record ModelUsageRecord
		require.NoError(t, json.Unmarshal(scanner.Bytes(), &record))
		records = append(records, record)
	}

	require.NoError(t, scanner.Err())

	return records
}

func modelReasoningItem(id, encrypted string) map[string]any {
	return map[string]any{
		"type":              "reasoning",
		"id":                id,
		"encrypted_content": encrypted,
		"summary":           []any{},
	}
}

func modelFunctionCallItem(id, callID, name, arguments string) map[string]any {
	return map[string]any{
		"type":      "function_call",
		"id":        id,
		"call_id":   callID,
		"name":      name,
		"arguments": arguments,
	}
}