package main

import (
	"bytes"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"
)

type roundTripFunc func(*http.Request) (*http.Response, error)

func (f roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) {
	return f(request)
}

func TestStreamGeminiImageHashDoesNotNeedWholeResponse(t *testing.T) {
	payload := paddedPNG(4 << 20)
	response := fmt.Sprintf(`{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"%s"}}]}}]}`,
		base64.StdEncoding.EncodeToString(payload))
	size, digest, err := streamGeminiImageHash(bytes.NewBufferString(response))
	if err != nil {
		t.Fatalf("stream Gemini output: %v", err)
	}
	expected := sha256.Sum256(payload)
	if size != int64(len(payload)) || digest != hex.EncodeToString(expected[:]) {
		t.Fatalf("size=%d digest=%s", size, digest)
	}
}

func TestPaddedPNGVariantsAreExactSizeAndUnique(t *testing.T) {
	first := paddedPNGVariant(256<<10, 1)
	second := paddedPNGVariant(256<<10, 2)
	if len(first) != 256<<10 || len(second) != 256<<10 {
		t.Fatalf("variant sizes=%d/%d", len(first), len(second))
	}
	firstHash := sha256.Sum256(first)
	secondHash := sha256.Sum256(second)
	if firstHash == secondHash {
		t.Fatal("distinct variants have the same SHA-256")
	}
}

func TestStreamGeminiRequestBodyPreservesMultipleInputs(t *testing.T) {
	inputs := geminiInputs(3, 2<<20, "multi-image-test", 37)
	var body struct {
		Contents []struct {
			Parts []struct {
				InlineData *struct {
					MIMEType string `json:"mimeType"`
					Data     string `json:"data"`
				} `json:"inlineData"`
			} `json:"parts"`
		} `json:"contents"`
		GenerationConfig struct {
			ResponseModalities []string `json:"responseModalities"`
		} `json:"generationConfig"`
	}
	stream := streamGeminiRequestBody(inputs)
	defer stream.Close()
	if err := json.NewDecoder(stream).Decode(&body); err != nil {
		t.Fatalf("decode streamed Gemini request: %v", err)
	}
	if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 4 {
		t.Fatalf("unexpected Gemini body structure: %+v", body)
	}
	for index, input := range inputs {
		inlineData := body.Contents[0].Parts[index+1].InlineData
		if inlineData == nil || inlineData.MIMEType != "image/png" {
			t.Fatalf("image %d inline data=%+v", index, inlineData)
		}
		decoded, err := base64.StdEncoding.DecodeString(inlineData.Data)
		if err != nil {
			t.Fatalf("decode streamed input %d: %v", index, err)
		}
		if !bytes.Equal(decoded, input) {
			t.Fatalf("streamed input %d differs from source", index)
		}
	}
	if len(body.GenerationConfig.ResponseModalities) != 1 || body.GenerationConfig.ResponseModalities[0] != "IMAGE" {
		t.Fatalf("response modalities=%v", body.GenerationConfig.ResponseModalities)
	}
}

func TestGeminiRequestBodyCanBeReplayedAfterHTTP2Failure(t *testing.T) {
	inputs := geminiInputs(3, 768<<10, "replay-test", 91)
	request, err := newGeminiRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", inputs)
	if err != nil {
		t.Fatalf("new Gemini request: %v", err)
	}
	if request.GetBody == nil {
		t.Fatal("streaming Gemini request does not provide GetBody")
	}
	first, err := io.ReadAll(request.Body)
	if err != nil {
		t.Fatalf("read first request body: %v", err)
	}
	_ = request.Body.Close()
	replayedBody, err := request.GetBody()
	if err != nil {
		t.Fatalf("replay Gemini request body: %v", err)
	}
	replayed, err := io.ReadAll(replayedBody)
	if err != nil {
		t.Fatalf("read replayed request body: %v", err)
	}
	_ = replayedBody.Close()
	if !bytes.Equal(first, replayed) {
		t.Fatal("replayed Gemini request body differs from the original")
	}
	if request.ContentLength != int64(len(first)) {
		t.Fatalf("content length=%d, want %d", request.ContentLength, len(first))
	}
}

func TestGeminiMultiImageRequestUsesThreeFileDataReferences(t *testing.T) {
	request, err := newGeminiFileDataRequest(t.Context(), "https://gateway.example/v1beta/models/test:generateContent", []string{
		"https://fixtures.example/one.png",
		"https://fixtures.example/two.png",
		"https://fixtures.example/three.png",
	})
	if err != nil {
		t.Fatalf("new multi-image request: %v", err)
	}
	var body struct {
		Contents []struct {
			Parts []struct {
				FileData *struct {
					FileURI string `json:"fileUri"`
				} `json:"fileData"`
			} `json:"parts"`
		} `json:"contents"`
	}
	if err := json.NewDecoder(request.Body).Decode(&body); err != nil {
		t.Fatalf("decode multi-image request: %v", err)
	}
	if len(body.Contents) != 1 || len(body.Contents[0].Parts) != 4 {
		t.Fatalf("unexpected multi-image body: %+v", body)
	}
	for index := 1; index < 4; index++ {
		if body.Contents[0].Parts[index].FileData == nil || body.Contents[0].Parts[index].FileData.FileURI == "" {
			t.Fatalf("missing fileData at part %d", index)
		}
	}
}

func TestVideoCombinationsProvide128UniqueInputs(t *testing.T) {
	images := make([]string, 16)
	for index := range images {
		images[index] = fmt.Sprintf("https://fixtures.example/image-%02d", index)
	}
	combinations := videoCombinations(images, 128)
	if len(combinations) != 128 {
		t.Fatalf("combinations=%d", len(combinations))
	}
	seen := map[string]struct{}{}
	for _, combination := range combinations {
		seen[fmt.Sprint(combination)] = struct{}{}
		if len(combination) != 3 && len(combination) != 6 && len(combination) != 9 {
			t.Fatalf("invalid combination size=%d", len(combination))
		}
	}
	if len(seen) != 128 {
		t.Fatalf("unique combinations=%d", len(seen))
	}
	if got := requestedVideoCombinationCount(combinations, "video-capacity", 144, 0, 1); got != 131 {
		t.Fatalf("requested combinations=%d, want 131", got)
	}
}

func TestVideoProviderQuotaCombinationsUseDistinctNormalSizedImages(t *testing.T) {
	images := []string{
		"https://fixtures.example/image-00.png",
		"https://fixtures.example/image-01.png",
		"https://fixtures.example/image-02.png",
		"https://fixtures.example/image-03.png",
	}
	combinations := videoProviderQuotaCombinations(images, 24)
	if got := requestedVideoCombinationCount(combinations, "video-provider-quota", 24, 0, 1); got != 24 {
		t.Fatalf("quota combinations=%d, want 24", got)
	}
	for _, combination := range combinations {
		if len(combination) != 3 {
			t.Fatalf("invalid quota combination: %v", combination)
		}
		for _, imageURL := range combination {
			if !strings.Contains(imageURL, "acceptance_quota_variant=") {
				t.Fatalf("quota image lacks unique variant: %s", imageURL)
			}
		}
	}
}

func TestGeminiLoadIsSplitAcrossTwoGatewayAPIs(t *testing.T) {
	output := paddedPNG(256 << 10)
	encoded := base64.StdEncoding.EncodeToString(output)
	var firstCalls atomic.Int64
	var secondCalls atomic.Int64
	var firstKeyCalls atomic.Int64
	var secondKeyCalls atomic.Int64
	newGateway := func(calls *atomic.Int64) *httptest.Server {
		return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			calls.Add(1)
			switch r.Header.Get("Authorization") {
			case "Bearer key-1":
				firstKeyCalls.Add(1)
			case "Bearer key-2":
				secondKeyCalls.Add(1)
			default:
				t.Errorf("unexpected authorization header")
			}
			if r.Header.Get(runHeader) != "run-1" || r.Header.Get(tokenHeader) != "token-1" {
				t.Errorf("missing acceptance headers")
			}
			var body map[string]any
			if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
				t.Errorf("decode Gemini body: %v", err)
			}
			_ = json.NewEncoder(w).Encode(map[string]any{
				"candidates": []any{map[string]any{"content": map[string]any{
					"parts": []any{map[string]any{"inlineData": map[string]any{
						"mimeType": "image/png", "data": encoded,
					}}},
				}}},
			})
		}))
	}
	first := newGateway(&firstCalls)
	defer first.Close()
	second := newGateway(&secondCalls)
	defer second.Close()

	opts := options{
		gateways: []string{first.URL, second.URL}, apiKeys: []string{"key-1", "key-2"}, runID: "run-1",
		runToken: "token-1", geminiModel: "gemini-image-test",
	}
	result := runGemini(t.Context(), http.DefaultClient, opts, "dual-api", 8, 1, 256<<10, 256<<10, false, 0, 1)
	if result.err != nil || result.report.Completed != 8 {
		t.Fatalf("result=%+v err=%v", result.report, result.err)
	}
	if firstCalls.Load() != 4 || secondCalls.Load() != 4 {
		t.Fatalf("gateway calls=%d/%d", firstCalls.Load(), secondCalls.Load())
	}
	if firstKeyCalls.Load() != 4 || secondKeyCalls.Load() != 4 {
		t.Fatalf("API key calls=%d/%d", firstKeyCalls.Load(), secondKeyCalls.Load())
	}
}

func TestDistributedShardIndexesCoverWorkloadWithoutOverlap(t *testing.T) {
	first := options{shardIndex: 0, shardCount: 2}
	second := options{shardIndex: 1, shardCount: 2}
	if first.shardRequestCount(1001) != 501 || second.shardRequestCount(1001) != 500 {
		t.Fatalf("shard counts=%d/%d", first.shardRequestCount(1001), second.shardRequestCount(1001))
	}
	seen := map[int]bool{}
	for local := 0; local < first.shardRequestCount(1001); local++ {
		seen[first.logicalRequestIndex(local)] = true
	}
	for local := 0; local < second.shardRequestCount(1001); local++ {
		index := second.logicalRequestIndex(local)
		if seen[index] {
			t.Fatalf("duplicate logical index %d", index)
		}
		seen[index] = true
	}
	if len(seen) != 1001 {
		t.Fatalf("covered indexes=%d", len(seen))
	}
}

func TestGeminiInputVariantIsRunScoped(t *testing.T) {
	if geminiInputVariant("run-a", 7) == geminiInputVariant("run-b", 7) {
		t.Fatal("different Run IDs produced the same input variant")
	}
	if geminiInputVariant("run-a", 7) == geminiInputVariant("run-a", 8) {
		t.Fatal("different logical indexes produced the same input variant")
	}
}

func TestAcceptanceIdempotencyKeyIncludesExecutionAndLogicalIndex(t *testing.T) {
	request := httptest.NewRequest(http.MethodPost, "https://gateway.example/v1", nil)
	opts := options{
		apiKeys: []string{"key-1"}, runID: "run-1", runToken: "token-1", executionID: "p24-2-gemini-baseline",
	}
	opts.setHeaders(request, 17, false)
	if got := request.Header.Get("Idempotency-Key"); got != "acceptance-run-1-p24-2-gemini-baseline-17" {
		t.Fatalf("idempotency key=%q", got)
	}
	poll := httptest.NewRequest(http.MethodGet, "https://gateway.example/result", nil)
	opts.setHeaders(poll, 17, false)
	if got := poll.Header.Get("Idempotency-Key"); got != "" {
		t.Fatalf("GET request unexpectedly has idempotency key %q", got)
	}
}

func TestValidateVideoAssetDownloadsFinalMedia(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		if request.Header.Get("Authorization") != "" {
			t.Fatal("acceptance credentials leaked to external media host")
		}
		w.Header().Set("Content-Type", "video/mp4")
		_, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'})
	}))
	defer server.Close()
	if err := validateVideoAsset(t.Context(), server.Client(), options{}, server.URL+"/result.mp4", 0); err != nil {
		t.Fatalf("validate video asset: %v", err)
	}
	if got := findMediaURL(map[string]any{"content": map[string]any{"video_url": server.URL}}); got != server.URL {
		t.Fatalf("media URL=%q", got)
	}
}

func TestValidateVideoAssetUsesIndependentTransportForExternalMedia(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "video/mp4")
		_, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'})
	}))
	defer server.Close()
	poisoned := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
		return nil, errors.New("gateway-only transport must not be used for external media")
	})}
	opts := options{gatewayTLSName: "gateway.easyai.local"}
	if err := validateVideoAsset(t.Context(), poisoned, opts, server.URL+"/result.mp4", 0); err != nil {
		t.Fatalf("validate external video with independent transport: %v", err)
	}
}

func TestValidateVideoAssetUsesGatewayForMaterializedPath(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
		if request.Host != "gateway.easyai.local" || request.Header.Get("Authorization") != "Bearer key-1" {
			t.Fatalf("host=%q authorization=%q", request.Host, request.Header.Get("Authorization"))
		}
		if request.URL.Path != "/static/generated/result.mp4" {
			t.Fatalf("path=%q", request.URL.Path)
		}
		w.Header().Set("Content-Type", "video/mp4")
		_, _ = w.Write([]byte{0, 0, 0, 16, 'f', 't', 'y', 'p', 'i', 's', 'o', 'm'})
	}))
	defer server.Close()
	opts := options{
		gateways: []string{server.URL}, apiKeys: []string{"key-1"},
		runID: "run-1", runToken: "token-1", gatewayTLSName: "gateway.easyai.local",
	}
	if err := validateVideoAsset(t.Context(), server.Client(), opts, "/static/generated/result.mp4", 0); err != nil {
		t.Fatalf("validate materialized video: %v", err)
	}
	got := findMediaURL(map[string]any{
		"raw":  map[string]any{"video_url": "http://internal.invalid/video.mp4"},
		"data": []any{map[string]any{"video_url": "/static/generated/result.mp4"}},
	})
	if got != "/static/generated/result.mp4" {
		t.Fatalf("preferred media URL=%q", got)
	}
}

func TestAcceptanceReportErrorRedactsSecretsAndURLs(t *testing.T) {
	got := redactError(
		`token-1 failed at https://example.invalid/video.mp4?token=signed`,
		options{runToken: "token-1"},
	)
	if got != `[REDACTED] failed at [REDACTED_URL]` {
		t.Fatalf("redacted error=%q", got)
	}
}

func TestAcceptanceReportCannotOverwriteExistingRunArtifact(t *testing.T) {
	path := filepath.Join(t.TempDir(), "load-report.json")
	if err := writeReportExclusive(path, []byte(`{"runId":"first"}`)); err != nil {
		t.Fatalf("write first report: %v", err)
	}
	if err := writeReportExclusive(path, []byte(`{"runId":"second"}`)); err == nil {
		t.Fatal("second report unexpectedly overwrote the first report")
	}
	payload, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("read report: %v", err)
	}
	if string(payload) != "{\"runId\":\"first\"}\n" {
		t.Fatalf("report changed after rejected overwrite: %s", payload)
	}
}

func TestAcceptanceFailurePreservesOperationAndHTTPStatus(t *testing.T) {
	err := withOperation("video_poll", &httpStatusError{Status: http.StatusUnauthorized, Body: "unauthorized"})
	var operationErr *operationError
	if !errors.As(err, &operationErr) || operationErr.Operation != "video_poll" {
		t.Fatalf("operation error=%#v", operationErr)
	}
	var statusErr *httpStatusError
	if !errors.As(err, &statusErr) || statusErr.Status != http.StatusUnauthorized {
		t.Fatalf("HTTP status error=%#v", statusErr)
	}
}

func TestAcceptanceGatewayTLSNameSetsHostHeader(t *testing.T) {
	request := httptest.NewRequest(http.MethodGet, "https://127.0.0.1/api/v1/healthz", nil)
	opts := options{
		apiKeys: []string{"key-1"}, runID: "run-1", runToken: "token-1",
		gatewayTLSName: "ai.example.com",
	}
	opts.setHeaders(request, 0, false)
	if request.Host != "ai.example.com" {
		t.Fatalf("Host=%q", request.Host)
	}
}

func TestAcceptanceRootCAsRejectsSymlink(t *testing.T) {
	root := t.TempDir()
	target := filepath.Join(root, "ca.pem")
	if err := os.WriteFile(target, []byte("not a certificate"), 0o600); err != nil {
		t.Fatal(err)
	}
	link := filepath.Join(root, "ca-link.pem")
	if err := os.Symlink(target, link); err != nil {
		t.Fatal(err)
	}
	if _, err := acceptanceRootCAs(link); err == nil {
		t.Fatal("symlink CA file was accepted")
	}
}
