package main

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

type apiFixture struct {
	*studioFixture
	gw   http.Handler
	key  string // the API key (secret)
	key2 string
}

func newAPIFixture(t *testing.T) *apiFixture {
	f := newStudioFixture(t)
	f.s.keys = newKeyStore(filepath.Join(t.TempDir(), "keys.json"))
	k1, _ := f.s.keys.create("one")
	k2, _ := f.s.keys.create("two")
	return &apiFixture{studioFixture: f, gw: f.s.gatewayMux(), key: k1, key2: k2}
}

func (a *apiFixture) call(method, path, key, ctype string, body io.Reader) *httptest.ResponseRecorder {
	r := httptest.NewRequest(method, path, body)
	r.Host = "api.wxch.org"
	if key != "" {
		r.Header.Set("Authorization", "Bearer "+key)
	}
	if ctype != "" {
		r.Header.Set("Content-Type", ctype)
	}
	w := httptest.NewRecorder()
	a.gw.ServeHTTP(w, r)
	return w
}

func (a *apiFixture) postJSON(path, key, body string) *httptest.ResponseRecorder {
	return a.call("POST", path, key, "application/json", strings.NewReader(body))
}

func keyIDOf(a *apiFixture, secret string) string {
	k, _ := a.s.keys.authenticate(secret)
	return k.ID
}

// nextJob waits until a job message reached the machine (scheduling as it goes) and returns its id.
func (a *apiFixture) nextJob(t *testing.T) string {
	for i := 0; i < 400; i++ {
		a.s.jobs.schedule(a.s.hub)
		a.mu.Lock()
		id := ""
		for k := range a.sent {
			if a.sent[k].T == "job" {
				id = a.sent[k].ID
			}
		}
		if id != "" {
			a.sent = nil
		}
		a.mu.Unlock()
		if id != "" {
			return id
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatal("no job reached the machine")
	return ""
}

// agentFinishes plays the machine: takes the next job, stores the files and reports the job done.
func (a *apiFixture) agentFinishes(t *testing.T, files map[string]string, text string) string {
	id := a.nextJob(t)
	var outs []JobOutput
	for name, data := range files {
		if w := a.agentReq("tok-pgx", "POST", "job="+id+"&name="+name+"&dir=out", strings.NewReader(data)); w.Code != 200 {
			t.Fatalf("upload %s: %d", name, w.Code)
		}
		outs = append(outs, JobOutput{Name: name, Mime: mimeFor(name)})
	}
	res, _ := json.Marshal(map[string]any{"outputs": outs, "text": text})
	a.s.jobs.onAgentMsg("pgx", Msg{T: "job_done", ID: id, OK: true, Status: 200, Data: res})
	return id
}

func TestJobAPINeedsAKey(t *testing.T) {
	a := newAPIFixture(t)
	for _, p := range []string{"/v1/capabilities", "/v1/jobs", "/v1/jobs/abc", "/v1/jobs/abc/files/x.png"} {
		if w := a.call("GET", p, "", "", nil); w.Code != 401 {
			t.Errorf("%s without a key: %d", p, w.Code)
		}
	}
	for _, p := range []string{"/v1/images/generations", "/v1/audio/transcriptions", "/v1/jobs"} {
		if w := a.postJSON(p, "", "{}"); w.Code != 401 {
			t.Errorf("POST %s without a key: %d", p, w.Code)
		}
	}
}

func TestJobAPIRefusesWhatNobodyOffers(t *testing.T) {
	a := newAPIFixture(t)
	w := a.postJSON("/v1/jobs", a.key, `{"capability":"image.generate","input":{"prompt":"x"}}`)
	if w.Code != 503 || !strings.Contains(w.Body.String(), "service_unavailable") {
		t.Fatalf("%d %s", w.Code, w.Body.String())
	}
	a.machine("pgx", "available", "image.generate")
	if w := a.postJSON("/v1/jobs", a.key, `{"capability":"nope"}`); w.Code != 400 {
		t.Errorf("unknown capability: %d", w.Code)
	}
	if w := a.postJSON("/v1/jobs", a.key, `{"capability":"image.generate","input":{}}`); w.Code != 400 {
		t.Errorf("a missing prompt: %d", w.Code)
	}
	if w := a.postJSON("/v1/jobs", a.key, `not json`); w.Code != 400 {
		t.Errorf("bad json: %d", w.Code)
	}
}

func TestJobAPIFullRoundTripAndPrivacy(t *testing.T) {
	a := newAPIFixture(t)
	a.machine("pgx", "available", "image.generate")
	w := a.postJSON("/v1/jobs", a.key, `{"capability":"image.generate","input":{"prompt":"a cat"}}`)
	if w.Code != 201 {
		t.Fatalf("%d %s", w.Code, w.Body.String())
	}
	var j apiJob
	json.Unmarshal(w.Body.Bytes(), &j)
	if j.Status != "queued" || j.Object != "job" {
		t.Fatalf("%+v", j)
	}
	a.agentFinishes(t, map[string]string{"out_1.png": "PNGDATA"}, "")
	w = a.call("GET", "/v1/jobs/"+j.ID, a.key, "", nil)
	json.Unmarshal(w.Body.Bytes(), &j)
	if j.Status != "done" || len(j.Outputs) != 1 || !strings.HasPrefix(j.Outputs[0].URL, "https://api.wxch.org/v1/jobs/"+j.ID+"/files/out_1.png?exp=") {
		t.Fatalf("%+v", j)
	}
	// the signed link needs no key
	link := strings.TrimPrefix(j.Outputs[0].URL, "https://api.wxch.org")
	if f := a.call("GET", link, "", "", nil); f.Code != 200 || f.Body.String() != "PNGDATA" || f.Header().Get("Content-Type") != "image/png" {
		t.Errorf("signed link: %d %q", f.Code, f.Body.String())
	}
	if f := a.call("GET", strings.Replace(link, "sig=", "sig=0", 1), "", "", nil); f.Code != 404 {
		t.Errorf("a changed signature: %d", f.Code)
	}
	if f := a.call("GET", strings.Replace(link, "out_1.png", "other.png", 1), "", "", nil); f.Code != 404 {
		t.Errorf("a signature for another file: %d", f.Code)
	}
	expired := "/v1/jobs/" + j.ID + "/files/out_1.png?exp=1&sig=" + a.s.fileSig(j.ID, "out_1.png", 1)
	if f := a.call("GET", expired, "", "", nil); f.Code != 404 {
		t.Errorf("an expired link: %d", f.Code)
	}
	// the owner can fetch with the key; another key and the studio cannot see the job at all
	if f := a.call("GET", "/v1/jobs/"+j.ID+"/files/out_1.png", a.key, "", nil); f.Code != 200 {
		t.Errorf("the owner with the key: %d", f.Code)
	}
	if f := a.call("GET", "/v1/jobs/"+j.ID+"/files/out_1.png", a.key2, "", nil); f.Code != 404 {
		t.Errorf("another key got the file: %d", f.Code)
	}
	if f := a.call("GET", "/v1/jobs/"+j.ID, a.key2, "", nil); f.Code != 404 {
		t.Errorf("another key saw the job: %d", f.Code)
	}
	var list struct{ Data []apiJob }
	json.Unmarshal(a.call("GET", "/v1/jobs", a.key2, "", nil).Body.Bytes(), &list)
	if len(list.Data) != 0 {
		t.Errorf("another key lists %d jobs", len(list.Data))
	}
	if st := a.do("GET", "/studio/api/jobs/"+j.ID, nil, "", true); st.Code != 404 {
		t.Errorf("the studio saw an API job: %d", st.Code)
	}
	var sl struct{ Jobs []Job }
	json.Unmarshal(a.do("GET", "/studio/api/jobs", nil, "", true).Body.Bytes(), &sl)
	if len(sl.Jobs) != 0 {
		t.Errorf("the studio lists API jobs: %d", len(sl.Jobs))
	}
	if d := a.call("DELETE", "/v1/jobs/"+j.ID, a.key, "", nil); d.Code != 200 {
		t.Errorf("delete: %d", d.Code)
	}
	if _, ok := a.s.jobs.get(j.ID); ok {
		t.Error("the job is still there")
	}
}

func TestJobAPIFilesInJSONAndTheLimitPerKey(t *testing.T) {
	a := newAPIFixture(t)
	a.machine("pgx", "owner", "image.edit") // nothing runs, jobs stay queued
	body := `{"capability":"image.edit","input":{"prompt":"blue"},"files":{"image":{"name":"../my pic.png","data":"` + base64.StdEncoding.EncodeToString([]byte("PICDATA")) + `"}}}`
	w := a.postJSON("/v1/jobs", a.key, body)
	if w.Code != 201 {
		t.Fatalf("%d %s", w.Code, w.Body.String())
	}
	var j apiJob
	json.Unmarshal(w.Body.Bytes(), &j)
	if _, has := j.Input["image"]; has {
		t.Error("the stored file name must not be shown")
	}
	stored, _ := a.s.jobs.get(j.ID)
	name, _ := stored.Input["image"].(string)
	if b, err := os.ReadFile(filepath.Join(a.s.jobs.jobDir(j.ID), "in", name)); err != nil || string(b) != "PICDATA" || strings.ContainsAny(name, "/\\") {
		t.Errorf("the upload: %q %v", name, err)
	}
	if w := a.postJSON("/v1/jobs", a.key, `{"capability":"image.edit","input":{"prompt":"x"},"files":{"image":{"name":"a.png","data":"!!!"}}}`); w.Code != 400 {
		t.Errorf("bad base64: %d", w.Code)
	}
	if w := a.postJSON("/v1/jobs", a.key, `{"capability":"image.generate","input":{"prompt":"x"},"files":{"image":{"name":"a.png","data":"QQ=="}}}`); w.Code != 400 {
		t.Errorf("a file for a capability that takes none: %d", w.Code)
	}
	for i := 1; i < maxActivePerKey; i++ {
		if w := a.postJSON("/v1/jobs", a.key, body); w.Code != 201 {
			t.Fatalf("job %d: %d %s", i+1, w.Code, w.Body.String())
		}
	}
	if w := a.postJSON("/v1/jobs", a.key, body); w.Code != 429 {
		t.Errorf("one key may not fill the queue: %d", w.Code)
	}
	if w := a.postJSON("/v1/jobs", a.key2, body); w.Code != 201 {
		t.Errorf("another key is not affected: %d", w.Code)
	}
}

func TestOpenAIImagesGenerationsWaitsForTheResult(t *testing.T) {
	a := newAPIFixture(t)
	a.machine("pgx", "available", "image.generate")
	res := make(chan *httptest.ResponseRecorder, 1)
	go func() {
		res <- a.postJSON("/v1/images/generations", a.key, `{"model":"x","prompt":"a cat","n":2,"size":"512x512","response_format":"b64_json"}`)
	}()
	id := a.agentFinishes(t, map[string]string{"a.png": "ONE", "b.png": "TWO"}, "")
	w := <-res
	if w.Code != 200 {
		t.Fatalf("%d %s", w.Code, w.Body.String())
	}
	var out struct {
		Created int64
		Data    []map[string]string
	}
	json.Unmarshal(w.Body.Bytes(), &out)
	if len(out.Data) != 2 {
		t.Fatalf("%s", w.Body.String())
	}
	got := map[string]bool{}
	for _, d := range out.Data {
		raw, _ := base64.StdEncoding.DecodeString(d["b64_json"])
		got[string(raw)] = true
	}
	if !got["ONE"] || !got["TWO"] {
		t.Errorf("%v", got)
	}
	j, _ := a.s.jobs.get(id)
	if j.Input["count"] != int64(2) || j.Input["size"] != "512x512" || j.Input["prompt"] != "a cat" {
		t.Errorf("the request was not translated: %v", j.Input)
	}
	// the default is a link
	res2 := make(chan *httptest.ResponseRecorder, 1)
	go func() { res2 <- a.postJSON("/v1/images/generations", a.key, `{"prompt":"dog"}`) }()
	a.agentFinishes(t, map[string]string{"c.png": "THREE"}, "")
	w2 := <-res2
	json.Unmarshal(w2.Body.Bytes(), &out)
	if len(out.Data) != 1 || !strings.Contains(out.Data[0]["url"], "/files/c.png?exp=") {
		t.Errorf("%s", w2.Body.String())
	}
	if w := a.postJSON("/v1/images/generations", a.key, `{"prompt":"x","size":"7x7"}`); w.Code != 400 {
		t.Errorf("a size that is not offered: %d", w.Code)
	}
}

func apiForm(fields map[string]string, file, fname, data string) (*bytes.Buffer, string) {
	var b bytes.Buffer
	mw := multipart.NewWriter(&b)
	for k, v := range fields {
		mw.WriteField(k, v)
	}
	fw, _ := mw.CreateFormFile(file, fname)
	fw.Write([]byte(data))
	mw.Close()
	return &b, mw.FormDataContentType()
}

func TestOpenAIImagesEditsAndAudio(t *testing.T) {
	a := newAPIFixture(t)
	a.machine("pgx", "available", "image.edit", "audio.transcribe")
	b, ct := apiForm(map[string]string{"prompt": "make it blue", "response_format": "b64_json"}, "image", "in.png", "PIC")
	res := make(chan *httptest.ResponseRecorder, 1)
	go func() { res <- a.call("POST", "/v1/images/edits", a.key, ct, b) }()
	id := a.agentFinishes(t, map[string]string{"e.png": "EDITED"}, "")
	w := <-res
	var out struct{ Data []map[string]string }
	json.Unmarshal(w.Body.Bytes(), &out)
	if w.Code != 200 || len(out.Data) != 1 {
		t.Fatalf("%d %s", w.Code, w.Body.String())
	}
	if raw, _ := base64.StdEncoding.DecodeString(out.Data[0]["b64_json"]); string(raw) != "EDITED" {
		t.Errorf("%q", raw)
	}
	if j, _ := a.s.jobs.get(id); j.Capability != "image.edit" {
		t.Errorf("%+v", j)
	}
	// speech, in every format
	files := map[string]string{"transcript.txt": "hello world\n", "transcript.srt": "1\n00:00:00,000 --> 00:00:01,000\nhello world\n\n", "transcript.vtt": "WEBVTT\n\n", "transcript.json": `{"text":"hello world","segments":[]}`}
	for format, want := range map[string]string{"": `{"text":"hello world"}`, "text": "hello world", "srt": "-->", "vtt": "WEBVTT", "verbose_json": "segments"} {
		fields := map[string]string{"model": "whisper-1", "language": "en"}
		if format != "" {
			fields["response_format"] = format
		}
		b, ct := apiForm(fields, "file", "talk.mp3", "AUDIO")
		res := make(chan *httptest.ResponseRecorder, 1)
		go func() { res <- a.call("POST", "/v1/audio/transcriptions", a.key, ct, b) }()
		id := a.agentFinishes(t, files, "hello world")
		w := <-res
		if w.Code != 200 || !strings.Contains(w.Body.String(), want) {
			t.Errorf("format %q: %d %q", format, w.Code, w.Body.String())
		}
		if j, _ := a.s.jobs.get(id); j.Capability != "audio.transcribe" || j.Input["task"] != "transcribe" || j.Input["language"] != "en" {
			t.Errorf("%+v", j.Input)
		}
	}
	// a failed job is reported as an error, not as an empty answer
	b2, ct2 := apiForm(map[string]string{}, "file", "talk.mp3", "AUDIO")
	res3 := make(chan *httptest.ResponseRecorder, 1)
	go func() { res3 <- a.call("POST", "/v1/audio/translations", a.key, ct2, b2) }()
	id3 := a.nextJob(t)
	a.s.jobs.onAgentMsg("pgx", Msg{T: "job_done", ID: id3, Status: 500, Err: "the engine failed: out of memory"})
	if w := <-res3; w.Code != 500 || !strings.Contains(w.Body.String(), "out of memory") {
		t.Errorf("%d %s", w.Code, w.Body.String())
	}
	if j, _ := a.s.jobs.get(id3); j.Input["task"] != "translate" {
		t.Errorf("translations must ask for the translate task: %v", j.Input)
	}
}

func TestOpenAIWrappersCancelTheJobWhenTheCallerGoesAway(t *testing.T) {
	a := newAPIFixture(t)
	a.machine("pgx", "owner", "image.generate") // it never starts
	r := httptest.NewRequest("POST", "/v1/images/generations", strings.NewReader(`{"prompt":"x"}`))
	r.Header.Set("Authorization", "Bearer "+a.key)
	ctx, cancel := context.WithCancel(r.Context())
	r = r.WithContext(ctx)
	done := make(chan struct{})
	go func() { a.gw.ServeHTTP(httptest.NewRecorder(), r); close(done) }()
	time.Sleep(300 * time.Millisecond)
	cancel()
	select {
	case <-done:
	case <-time.After(5 * time.Second):
		t.Fatal("the handler kept waiting after the caller left")
	}
	jobs := a.s.jobs.list("key:" + keyIDOf(a, a.key))
	if len(jobs) != 1 || jobs[0].Status != "cancelled" {
		t.Errorf("the abandoned job: %+v", jobs)
	}
}