package main

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func testFrontend(url string) frontend {
	return frontend{client: http.DefaultClient, gateway: url, prefix: "batch/", routePrefix: "native/", heartbeat: 10 * time.Millisecond}
}

func TestFrontendPayload(t *testing.T) {
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var body map[string]json.RawMessage
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			t.Error(err)
		}
		if string(body["model"]) != `"native/partner-model@001"` || string(body["tools"]) != `[{"opaque":true}]` || r.Header.Get("Authorization") != "Bearer synthetic" {
			t.Errorf("unexpected request: %s", body)
		}
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusBadGateway)
		_, _ = io.WriteString(w, `{"error":"synthetic failure"}`)
	}))
	defer upstream.Close()
	req := httptest.NewRequest(http.MethodPost, "/front", strings.NewReader(`{"model":"batch/partner-model@001","tools":[{"opaque":true}]}`))
	req.Header.Set("Authorization", "Bearer synthetic")
	res := httptest.NewRecorder()
	testFrontend(upstream.URL).ServeHTTP(res, req)
	if res.Code != http.StatusBadGateway || res.Body.String() != `{"error":"synthetic failure"}` {
		t.Fatalf("unexpected response: %d %s", res.Code, res.Body)
	}
}

func TestFrontendCancel(t *testing.T) {
	for _, headers := range []bool{false, true} {
		t.Run(map[bool]string{false: "pending_headers", true: "pending_body"}[headers], func(t *testing.T) {
			entered := make(chan struct{})
			cancelled := make(chan struct{})
			upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				_, _ = io.Copy(io.Discard, r.Body)
				if headers {
					w.Header().Set("Content-Type", "text/event-stream")
					_ = http.NewResponseController(w).Flush()
				}
				close(entered)
				<-r.Context().Done()
				close(cancelled)
			}))
			defer upstream.Close()
			defer upstream.CloseClientConnections()
			server := httptest.NewServer(testFrontend(upstream.URL))
			defer server.Close()
			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
			defer cancel()
			req, err := http.NewRequestWithContext(ctx, http.MethodPost, server.URL, strings.NewReader(`{"model":"batch/gemini-3.8-flash","stream":true}`))
			if err != nil {
				t.Fatal(err)
			}
			res, err := http.DefaultClient.Do(req)
			if err != nil {
				t.Fatal(err)
			}
			defer func() { _ = res.Body.Close() }()
			select {
			case <-entered:
			case <-ctx.Done():
				t.Fatal("upstream was not entered")
			}
			data := make([]byte, 1)
			if _, err := res.Body.Read(data); err != nil {
				t.Fatal(err)
			}
			if err := res.Body.Close(); err != nil {
				t.Fatal(err)
			}
			select {
			case <-cancelled:
			case <-ctx.Done():
				t.Fatal("disconnect did not cancel upstream")
			}
		})
	}
}

func TestFrontendStream(t *testing.T) {
	for _, complete := range []bool{false, true} {
		t.Run(map[bool]string{false: "incomplete", true: "complete"}[complete], func(t *testing.T) {
			upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", "text/event-stream")
				if complete {
					_, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"answer\"},\"finish_reason\":\"stop\"}]}\n\n")
				}
				_, _ = io.WriteString(w, "data: [DONE]\n\n")
			}))
			defer upstream.Close()
			req := httptest.NewRequest(http.MethodPost, "/front", strings.NewReader(`{"model":"batch/gemini-3.8-flash","stream":true}`))
			res := httptest.NewRecorder()
			testFrontend(upstream.URL).ServeHTTP(res, req)
			output := res.Body.String()
			if complete {
				if !strings.Contains(output, "answer") || !strings.Contains(output, "[DONE]") {
					t.Fatal(output)
				}
			} else if !strings.Contains(output, "batch_error") || strings.Contains(output, "[DONE]") {
				t.Fatal(output)
			}
		})
	}
}