package main

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

func TestCloud(t *testing.T) {
	for _, test := range []struct {
		name, state, success string
		cancel               bool
		wantError            bool
	}{
		{"success", "JOB_STATE_SUCCEEDED", "1", false, false},
		{"failed row", "JOB_STATE_SUCCEEDED", "0", false, true},
		{"failed job", "JOB_STATE_FAILED", "0", false, true},
		{"cancel", "JOB_STATE_PENDING", "0", true, true},
	} {
		t.Run(test.name, func(t *testing.T) {
			var mu sync.Mutex
			objects := map[string]bool{}
			input, prefix := "", ""
			cancelled := false
			submitted := make(chan struct{})
			response := json.RawMessage(`{"candidates":[{"content":{"parts":[{"text":"answer"}]},"finishReason":"STOP"}],"usageMetadata":{"totalTokenCount":16}}`)
			native := json.RawMessage(`{"contents":[{"role":"user","parts":[{"text":"synthetic"},{"inlineData":{"mimeType":"image/png","data":"aGVsbG8="}}]}],"generationConfig":{"temperature":0.2}}`)
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				mu.Lock()
				defer mu.Unlock()
				w.Header().Set("Content-Type", "application/json")
				switch {
				case strings.HasPrefix(r.URL.Path, "/upload/"):
					input = r.URL.Query().Get("name")
					objects[input] = true
					var row struct {
						Request json.RawMessage `json:"request"`
					}
					if err := json.NewDecoder(r.Body).Decode(&row); err != nil {
						t.Error(err)
					}
					if string(row.Request) != string(native) {
						t.Errorf("native body changed: %s", row.Request)
					}
					_, _ = fmt.Fprint(w, `{}`)
				case strings.HasSuffix(r.URL.Path, "/batchPredictionJobs"):
					var body struct {
						Model  string `json:"model"`
						Output struct {
							GCS struct {
								Prefix string `json:"outputUriPrefix"`
							} `json:"gcsDestination"`
						} `json:"outputConfig"`
					}
					if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
						t.Error(err)
					}
					if body.Model != "publishers/google/models/gemini-test" {
						t.Error(body.Model)
					}
					prefix = strings.TrimPrefix(body.Output.GCS.Prefix, "gs://synthetic/")
					_, _ = fmt.Fprint(w, `{"name":"projects/synthetic/locations/global/batchPredictionJobs/1","state":"JOB_STATE_PENDING"}`)
					close(submitted)
				case strings.HasSuffix(r.URL.Path, ":cancel"):
					cancelled = true
					_, _ = fmt.Fprint(w, `{}`)
				case strings.Contains(r.URL.Path, "/batchPredictionJobs/"):
					objects[prefix+"predictions.jsonl"] = true
					_, _ = fmt.Fprintf(w, `{"name":"projects/synthetic/locations/global/batchPredictionJobs/1","state":%q,"completionStats":{"successfulCount":%q}}`, test.state, test.success)
				case r.Method == "DELETE":
					name := strings.SplitN(r.URL.Path, "/o/", 2)[1]
					delete(objects, name)
					w.WriteHeader(204)
				case r.URL.Query().Get("alt") == "media":
					_, _ = fmt.Fprintf(w, `{"status":"","response":%s}`, response)
				case strings.HasSuffix(r.URL.Path, "/o"):
					var items []map[string]string
					for name := range objects {
						if strings.HasPrefix(name, r.URL.Query().Get("prefix")) {
							items = append(items, map[string]string{"name": name})
						}
					}
					_ = json.NewEncoder(w).Encode(map[string]any{"items": items})
				default:
					t.Error(r.Method, r.URL.String())
					w.WriteHeader(404)
				}
			}))
			defer server.Close()
			client := &cloud{client: server.Client(), vertex: server.URL, storage: server.URL, project: "synthetic", bucket: "synthetic", poll: time.Millisecond}
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()
			if test.cancel {
				go func() { <-submitted; cancel() }()
			}
			result, err := client.run(ctx, "gemini-test", native)
			if (err != nil) != test.wantError {
				t.Fatalf("result %s, error %v", result, err)
			}
			if !test.wantError && string(result) != string(response) {
				t.Errorf("result changed: %s", result)
			}
			mu.Lock()
			defer mu.Unlock()
			if len(objects) != 0 || cancelled != test.cancel {
				t.Errorf("cleanup objects=%v cancelled=%v", objects, cancelled)
			}
		})
	}
}