package server

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync"
	"testing"
)

// fakeVertex is the image service: it records what it was asked and answers
// like the real one (a job to poll, a picture when done).
type fakeVertex struct {
	srv  *httptest.Server
	mu   sync.Mutex
	auth []string
	gen  map[string]any
	edit struct {
		prompt string
		image  []byte
		name   string
	}
	polls  int
	failAs string // an error_code the job ends with
}

func newFakeVertex(t *testing.T) *fakeVertex {
	f := &fakeVertex{}
	f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		f.mu.Lock()
		defer f.mu.Unlock()
		f.auth = append(f.auth, r.Header.Get("Authorization"))
		// The pictures are open by their address, as at the real service.
		if r.URL.Path == "/uploads/userphoto/ai/t.png" {
			w.Header().Set("Content-Type", "image/png")
			_, _ = w.Write([]byte("\x89PNG-original"))
			return
		}
		w.Header().Set("Content-Type", "application/json")
		if r.Header.Get("Authorization") != "Bearer vx_good" {
			w.WriteHeader(401)
			fmt.Fprint(w, `{"ok":false,"error_code":"unauthorized","error":"Неверный ключ"}`)
			return
		}
		switch {
		case r.URL.Path == "/api/vertex/v1/generate":
			_ = json.NewDecoder(r.Body).Decode(&f.gen)
			w.WriteHeader(202)
			fmt.Fprint(w, `{"ok":true,"job_id":"job_gen1","status":"queued","cost":5}`)
		case r.URL.Path == "/api/vertex/v1/edit":
			if err := r.ParseMultipartForm(32 << 20); err != nil {
				w.WriteHeader(400)
				fmt.Fprint(w, `{"ok":false,"error_code":"invalid_image","error":"нет фото"}`)
				return
			}
			f.edit.prompt = r.FormValue("prompt")
			file, hdr, err := r.FormFile("image")
			if err == nil {
				f.edit.image, _ = io.ReadAll(file)
				f.edit.name = hdr.Filename
			}
			w.WriteHeader(202)
			fmt.Fprint(w, `{"ok":true,"job_id":"job_edit1","status":"queued","cost":2}`)
		case strings.HasPrefix(r.URL.Path, "/api/vertex/v1/jobs/"):
			f.polls++
			if f.failAs != "" {
				fmt.Fprintf(w, `{"ok":true,"job_id":"x","status":"failed","error_code":%q,"error":"сбой"}`, f.failAs)
				return
			}
			if f.polls < 2 {
				fmt.Fprint(w, `{"ok":true,"job_id":"x","status":"running"}`)
				return
			}
			fmt.Fprintf(w, `{"ok":true,"job_id":"x","status":"done","photo_url":"%s/uploads/userphoto/ai/t.png","cost":5,"balance":42}`, f.srv.URL)
		default:
			http.NotFound(w, r)
		}
	}))
	t.Cleanup(f.srv.Close)
	return f
}

// setImage points the panel at the fake service with a key.
func setImage(t *testing.T, s *Server, f *fakeVertex, key string) {
	t.Helper()
	_ = s.db.SetConfig(cfgImageURL, f.srv.URL)
	_ = s.db.SetConfig(cfgImageKey, key)
}

func apiCall(t *testing.T, s *Server, c *http.Cookie, method, path, body string) string {
	t.Helper()
	req := httptest.NewRequest(method, path, strings.NewReader(body))
	req.AddCookie(c)
	w := httptest.NewRecorder()
	s.Handler().ServeHTTP(w, req)
	return w.Body.String()
}

// A picture from nothing: the prompt, model and ratio go to the service with
// the key; the job's progress and its picture come back, the key never.
func TestImageGenerate(t *testing.T) {
	s, c := newTestServer(t)
	f := newFakeVertex(t)
	setImage(t, s, f, "vx_good")
	got := apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"cat in space","model":"artist","ratio":"3:4"}`)
	if !strings.Contains(got, `"ok":true`) || !strings.Contains(got, `"job_id":"job_gen1"`) || strings.Contains(got, "vx_good") {
		t.Fatalf("create: %s", got)
	}
	if f.gen["prompt"] != "cat in space" || f.gen["model"] != "artist" || f.gen["ratio"] != "3:4" || f.gen["enhance"] != true {
		t.Fatalf("what the service got: %v", f.gen)
	}
	if got := apiCall(t, s, c, "GET", "/api/panel/ai/image/job_gen1?wait=1", ""); !strings.Contains(got, `"status":"running"`) {
		t.Fatalf("first poll: %s", got)
	}
	got = apiCall(t, s, c, "GET", "/api/panel/ai/image/job_gen1?wait=1", "")
	if !strings.Contains(got, `"status":"done"`) || !strings.Contains(got, "/uploads/userphoto/ai/t.png") || !strings.Contains(got, `"balance":42`) || strings.Contains(got, "vx_good") {
		t.Fatalf("done: %s", got)
	}
}

// Unknown models and ratios are not sent on: the defaults go instead.
func TestImageGenerateDefaults(t *testing.T) {
	s, c := newTestServer(t)
	f := newFakeVertex(t)
	setImage(t, s, f, "vx_good")
	apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x","model":"gpt","ratio":"7:1"}`)
	if f.gen["model"] != "fast" || f.gen["ratio"] != "1:1" {
		t.Fatalf("defaults: %v", f.gen)
	}
}

// Editing sends the user's photo as a file with the instruction.
func TestImageEditFromAttachedPhoto(t *testing.T) {
	s, c := newTestServer(t)
	f := newFakeVertex(t)
	setImage(t, s, f, "vx_good")
	photo := []byte("\xff\xd8\xffJPEGDATA")
	b, _ := json.Marshal(map[string]any{"mode": "edit", "prompt": "сделай закат",
		"image": map[string]string{"mime": "image/jpeg", "data": base64.StdEncoding.EncodeToString(photo)}})
	got := apiCall(t, s, c, "POST", "/api/panel/ai/image", string(b))
	if !strings.Contains(got, `"job_id":"job_edit1"`) {
		t.Fatalf("edit: %s", got)
	}
	if f.edit.prompt != "сделай закат" || !bytes.Equal(f.edit.image, photo) || !strings.HasSuffix(f.edit.name, ".jpg") {
		t.Fatalf("what the service got: %q %q %q", f.edit.prompt, f.edit.image, f.edit.name)
	}
	// A format the service does not take is told so, not sent.
	b, _ = json.Marshal(map[string]any{"mode": "edit", "prompt": "x", "image": map[string]string{"mime": "image/gif", "data": "R0lG"}})
	if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", string(b)); !strings.Contains(got, `"ok":false`) || !strings.Contains(got, "png") {
		t.Fatalf("gif: %s", got)
	}
}

// A picture made before can be edited again: fetched from the service's own
// address only - never from another.
func TestImageEditFromEarlierPicture(t *testing.T) {
	s, c := newTestServer(t)
	f := newFakeVertex(t)
	setImage(t, s, f, "vx_good")
	b, _ := json.Marshal(map[string]any{"mode": "edit", "prompt": "темнее", "image_url": f.srv.URL + "/uploads/userphoto/ai/t.png"})
	if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", string(b)); !strings.Contains(got, `"ok":true`) {
		t.Fatalf("by url: %s", got)
	}
	if !bytes.Equal(f.edit.image, []byte("\x89PNG-original")) || !strings.HasSuffix(f.edit.name, ".png") {
		t.Fatalf("the picture sent: %q %q", f.edit.image, f.edit.name)
	}
	b, _ = json.Marshal(map[string]any{"mode": "edit", "prompt": "x", "image_url": "http://127.0.0.1:9/secret.png"})
	if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", string(b)); !strings.Contains(got, `"ok":false`) {
		t.Fatalf("a foreign address was fetched: %s", got)
	}
}

// What goes wrong is told in words: no key, a wrong key, the service's own
// codes, a job that failed.
func TestImageErrors(t *testing.T) {
	s, c := newTestServer(t)
	f := newFakeVertex(t)
	if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x"}`); !strings.Contains(got, "ключ") {
		t.Fatalf("no key: %s", got)
	}
	setImage(t, s, f, "vx_bad")
	if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x"}`); !strings.Contains(got, `"ok":false`) || !strings.Contains(got, "ключ") {
		t.Fatalf("wrong key: %s", got)
	}
	setImage(t, s, f, "vx_good")
	f.failAs = "insufficient_credits"
	if got := apiCall(t, s, c, "GET", "/api/panel/ai/image/job_gen1", ""); !strings.Contains(got, `"status":"failed"`) || !strings.Contains(got, "кредит") {
		t.Fatalf("failed job: %s", got)
	}
	if got := apiCall(t, s, c, "GET", "/api/panel/ai/image/..%2Fme", ""); !strings.Contains(got, `"ok":false`) {
		t.Fatalf("a job id out of shape: %s", got)
	}
	if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":""}`); !strings.Contains(got, `"ok":false`) {
		t.Fatalf("empty prompt: %s", got)
	}
}

// The assistant off in the profile means no pictures either.
func TestImageNeedsAssistantOn(t *testing.T) {
	s, c := newTestServer(t)
	f := newFakeVertex(t)
	setImage(t, s, f, "vx_good")
	a, _ := s.db.AdminByUsername("adm")
	_ = s.db.SetConfig(aiUserKey(a.ID), "0")
	if got := apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x"}`); !strings.Contains(got, "выключен") {
		t.Fatalf("off: %s", got)
	}
}

// The key is saved from Настройки, shown back to the admin, and the chat
// learns only that pictures are on.
func TestImageKeySetting(t *testing.T) {
	s, c := newTestServer(t)
	postAIModels(t, s, c, [][7]string{{"", "M", "gradio", "", "", "", ""}})
	f := url.Values{"csrf": {"tok"}, "m_type": {"gradio"}, "ai_image_key": {"vx_secret"}}
	req := httptest.NewRequest(http.MethodPost, "/settings/ai", strings.NewReader(f.Encode()))
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	req.AddCookie(c)
	req.AddCookie(&http.Cookie{Name: csrfCookie, Value: "tok"})
	s.Handler().ServeHTTP(httptest.NewRecorder(), req)
	if _, k := s.imageConf(); k != "vx_secret" {
		t.Fatalf("key not saved: %q", k)
	}
	if body := get(t, s, c, "/settings?frag=1").Body.String(); !strings.Contains(body, `value="vx_secret"`) {
		t.Fatal("the settings page does not show the key")
	}
	state := get(t, s, c, "/api/panel/ai").Body.String()
	if strings.Contains(state, "vx_secret") || !strings.Contains(state, `"image":true`) {
		t.Fatalf("state: %s", state)
	}
}

// With a key the model is told it can draw; without, it is not.
func TestImagePromptOnlyWithKey(t *testing.T) {
	s, _ := newTestServer(t)
	if strings.Contains(s.imagePrompt(), "```image") {
		t.Fatal("told about pictures with no key")
	}
	_ = s.db.SetConfig(cfgImageKey, "vx_good")
	p := s.imagePrompt()
	for _, want := range []string{"```image\n", "```image edit", "fast", "artist", "0.2plus", "16:9"} {
		if !strings.Contains(p, want) {
			t.Fatalf("the prompt lacks %q:\n%s", want, p)
		}
	}
}

// The chat gives the model the image instructions when pictures are on.
func TestAIChatMentionsImageTool(t *testing.T) {
	s, c := newTestServer(t)
	var system string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var b map[string]any
		_ = json.NewDecoder(r.Body).Decode(&b)
		system, _ = b["messages"].([]any)[0].(map[string]any)["content"].(string)
		fmt.Fprint(w, openAIChunk("content", "ok")+"data: [DONE]\n\n")
	}))
	defer srv.Close()
	postAIModels(t, s, c, [][7]string{{"", "M", "openai", srv.URL, "m", "", ""}})
	ask := func() {
		apiCall(t, s, c, "POST", "/api/panel/ai/chat", `{"message":"нарисуй кота","context":"x"}`)
	}
	ask()
	if strings.Contains(system, "```image") {
		t.Fatal("image tool offered with no key")
	}
	_ = s.db.SetConfig(cfgImageKey, "vx_good")
	ask()
	if !strings.Contains(system, "```image") {
		t.Fatalf("image tool not offered with a key:\n%s", system)
	}
}

// The description is improved by the service unless the model asks for it
// as it is.
func TestImageEnhanceSwitch(t *testing.T) {
	s, c := newTestServer(t)
	f := newFakeVertex(t)
	setImage(t, s, f, "vx_good")
	apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x"}`)
	if f.gen["enhance"] != true {
		t.Fatalf("by default: %v", f.gen)
	}
	apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x","enhance":false}`)
	if f.gen["enhance"] != false {
		t.Fatalf("switched off: %v", f.gen)
	}
	apiCall(t, s, c, "POST", "/api/panel/ai/image", `{"mode":"generate","prompt":"x","enhance":true}`)
	if f.gen["enhance"] != true {
		t.Fatalf("switched on: %v", f.gen)
	}
	{
		p := s.imagePrompt()
		_ = s.db.SetConfig(cfgImageKey, "vx_good")
		if !strings.Contains(s.imagePrompt(), "noenhance") {
			t.Fatalf("the model is not told of the switch:\n%s", p)
		}
	}
}