package main

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

type runnerFunc func(context.Context, string, json.RawMessage) (json.RawMessage, error)

func (f runnerFunc) run(ctx context.Context, model string, body json.RawMessage) (json.RawMessage, error) {
	return f(ctx, model, body)
}

func TestHTTPStream(t *testing.T) {
	for _, failed := range []bool{false, true} {
		server := httptest.NewServer(handler{heartbeat: time.Millisecond, runner: runnerFunc(func(_ context.Context, model string, body json.RawMessage) (json.RawMessage, error) {
			if model != "partner-model@001" || string(body) != `{"contents":[]}` {
				t.Errorf("changed request %s %s", model, body)
			}
			time.Sleep(20 * time.Millisecond)
			if failed {
				return nil, errors.New("synthetic failure")
			}
			return json.RawMessage(`{"candidates":[{"finishReason":"STOP"}]}`), nil
		})})
		req, _ := http.NewRequest("POST", server.URL+"/v1/chat/completions", strings.NewReader(`{"contents":[]}`))
		req.Header.Set("x-batch-model", "partner-model@001")
		req.Header.Set("x-batch-stream", "true")
		response, err := server.Client().Do(req)
		if err != nil {
			t.Fatal(err)
		}
		data, err := io.ReadAll(response.Body)
		_ = response.Body.Close()
		server.Close()
		if err != nil {
			t.Fatal(err)
		}
		if !strings.Contains(string(data), ": batch pending") || strings.Contains(string(data), "data:") == failed {
			t.Fatalf("failed=%v body=%s", failed, data)
		}
	}
}

func TestCleanupWait(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	cancelled := make(chan struct{})
	cleanup := make(chan struct{})
	finished := make(chan struct{})
	h := handler{runner: runnerFunc(func(ctx context.Context, _ string, _ json.RawMessage) (json.RawMessage, error) {
		<-ctx.Done()
		close(cancelled)
		<-cleanup
		return nil, ctx.Err()
	})}
	req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"contents":[]}`)).WithContext(ctx)
	req.Header.Set("x-batch-model", "partner-model@001")
	req.Header.Set("x-batch-stream", "false")
	go func() {
		h.ServeHTTP(httptest.NewRecorder(), req)
		close(finished)
	}()
	cancel()
	<-cancelled
	select {
	case <-finished:
		t.Error("handler returned before cloud cleanup")
	case <-time.After(20 * time.Millisecond):
	}
	close(cleanup)
	select {
	case <-finished:
	case <-time.After(time.Second):
		t.Fatal("handler did not finish after cleanup")
	}
}