package multimodal

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

func TestSpeechClientSendsMultipartAndReturnsTranscript(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/inference" || r.Method != http.MethodPost {
			t.Fatalf("unexpected speech request: %s %s", r.Method, r.URL.Path)
		}
		if err := r.ParseMultipartForm(1 << 20); err != nil {
			t.Fatal(err)
		}
		file, _, err := r.FormFile("file")
		if err != nil {
			t.Fatal(err)
		}
		defer file.Close()
		data, _ := io.ReadAll(file)
		if string(data) != "audio" {
			t.Fatalf("audio payload = %q", data)
		}
		_ = json.NewEncoder(w).Encode(map[string]string{"text": "распознанная фраза"})
	}))
	defer server.Close()
	text, err := (SpeechClient{BaseURL: server.URL, Client: server.Client()}).Transcribe(context.Background(), "voice.wav", "audio/wav", []byte("audio"))
	if err != nil || text != "распознанная фраза" {
		t.Fatalf("transcript=%q err=%v", text, err)
	}
}

func TestVisionClientSendsImageDataAndInstruction(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var payload struct {
			Messages []struct {
				Content []struct {
					Type     string `json:"type"`
					Text     string `json:"text"`
					ImageURL *struct {
						URL string `json:"url"`
					} `json:"image_url"`
				} `json:"content"`
			} `json:"messages"`
		}
		if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
			t.Fatal(err)
		}
		if len(payload.Messages) != 1 || len(payload.Messages[0].Content) != 2 || !strings.Contains(payload.Messages[0].Content[0].Text, "Опиши") || payload.Messages[0].Content[1].ImageURL == nil || !strings.HasPrefix(payload.Messages[0].Content[1].ImageURL.URL, "data:image/png;base64,") {
			t.Fatalf("unexpected vision payload: %+v", payload)
		}
		_ = json.NewEncoder(w).Encode(map[string]any{"choices": []any{map[string]any{"message": map[string]string{"content": "на изображении круг"}}}})
	}))
	defer server.Close()
	text, err := (VisionClient{BaseURL: server.URL + "/v1", Client: server.Client()}).Describe(context.Background(), "vision-model", "image/png", []byte{0, 1, 2}, "Опиши изображение")
	if err != nil || text != "на изображении круг" {
		t.Fatalf("description=%q err=%v", text, err)
	}
}

func TestVisionClientUsesReasoningContentWhenProviderOmitsContent(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewEncoder(w).Encode(map[string]any{"choices": []any{map[string]any{
			"message": map[string]string{"content": "", "reasoning_content": "чёрный квадрат"},
		}}})
	}))
	defer server.Close()
	text, err := (VisionClient{BaseURL: server.URL, Client: server.Client()}).Describe(context.Background(), "vision-model", "image/png", []byte("image"), "Опиши изображение")
	if err != nil || text != "чёрный квадрат" {
		t.Fatalf("description=%q err=%v", text, err)
	}
}

func TestImageGenerationClientSendsPromptAndReturnsBase64(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/images/generations" || r.Method != http.MethodPost {
			t.Fatalf("unexpected image request: %s %s", r.Method, r.URL.Path)
		}
		var payload struct {
			Model  string `json:"model"`
			Prompt string `json:"prompt"`
			Size   string `json:"size"`
		}
		if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
			t.Fatal(err)
		}
		if payload.Model != "image-model" || payload.Prompt != "green bird" || payload.Size != "1024x1024" {
			t.Fatalf("unexpected image payload: %+v", payload)
		}
		_ = json.NewEncoder(w).Encode(map[string]any{"data": []any{map[string]string{"b64_json": "aW1hZ2U="}}})
	}))
	defer server.Close()
	image, err := (ImageGenerationClient{BaseURL: server.URL, Client: server.Client()}).Generate(context.Background(), "image-model", "green bird", "1024x1024")
	if err != nil || image.DataBase64 != "aW1hZ2U=" || image.ContentType != "image/png" {
		t.Fatalf("image=%+v err=%v", image, err)
	}
}

func TestVideoUnderstandingClientSendsMultipartAndReturnsDescription(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/understand" || r.Method != http.MethodPost {
			t.Fatalf("unexpected video request: %s %s", r.Method, r.URL.Path)
		}
		if err := r.ParseMultipartForm(1 << 20); err != nil {
			t.Fatal(err)
		}
		if r.FormValue("model") != "video-model" || r.FormValue("instruction") != "describe movement" {
			t.Fatalf("unexpected video fields: model=%q instruction=%q", r.FormValue("model"), r.FormValue("instruction"))
		}
		file, _, err := r.FormFile("file")
		if err != nil {
			t.Fatal(err)
		}
		defer file.Close()
		data, _ := io.ReadAll(file)
		if string(data) != "video" {
			t.Fatalf("video payload=%q", data)
		}
		_ = json.NewEncoder(w).Encode(map[string]string{"text": "человек идёт"})
	}))
	defer server.Close()
	text, err := (VideoUnderstandingClient{BaseURL: server.URL, Client: server.Client()}).Describe(context.Background(), "video-model", "clip.mp4", "video/mp4", []byte("video"), "describe movement")
	if err != nil || text != "человек идёт" {
		t.Fatalf("description=%q err=%v", text, err)
	}
}

func TestComfyImageGenerationClientRunsWorkflowAndReturnsImage(t *testing.T) {
	polls := 0
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/prompt":
			var payload struct {
				Prompt map[string]any `json:"prompt"`
			}
			if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
				t.Fatal(err)
			}
			if payload.Prompt["6"].(map[string]any)["inputs"].(map[string]any)["text"] != "red bird" {
				t.Fatalf("prompt was not forwarded: %+v", payload.Prompt["6"])
			}
			_ = json.NewEncoder(w).Encode(map[string]string{"prompt_id": "test-prompt"})
		case "/history/test-prompt":
			polls++
			if polls == 1 {
				_, _ = w.Write([]byte(`{}`))
				return
			}
			_ = json.NewEncoder(w).Encode(map[string]any{"test-prompt": map[string]any{"outputs": map[string]any{"9": map[string]any{"images": []any{map[string]string{"filename": "out.png", "subfolder": "", "type": "output"}}}}}})
		case "/view":
			w.Header().Set("Content-Type", "image/png")
			_, _ = w.Write([]byte("png"))
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()
	image, err := (ComfyImageGenerationClient{BaseURL: server.URL, Client: server.Client(), PollInterval: time.Millisecond, PollTimeout: time.Second}).Generate(context.Background(), "", "red bird", "256x256")
	if err != nil || image.DataBase64 != "cG5n" || image.ContentType != "image/png" {
		t.Fatalf("image=%+v err=%v", image, err)
	}
}

func TestComfyImageGenerationClientDeletesOnlyCancelledPrompt(t *testing.T) {
	var deleted []string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/prompt":
			_ = json.NewEncoder(w).Encode(map[string]string{"prompt_id": "cancel-me"})
		case "/history/cancel-me":
			<-r.Context().Done()
		case "/queue":
			var payload struct {
				Delete []string `json:"delete"`
			}
			if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
				t.Fatal(err)
			}
			deleted = append(deleted, payload.Delete...)
			w.WriteHeader(http.StatusNoContent)
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()

	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() {
		_, err := (ComfyImageGenerationClient{BaseURL: server.URL, Client: server.Client(), PollInterval: time.Millisecond, PollTimeout: time.Minute}).Generate(ctx, "", "red bird", "256x256")
		done <- err
	}()
	time.Sleep(10 * time.Millisecond)
	cancel()
	if err := <-done; !errors.Is(err, context.Canceled) {
		t.Fatalf("err=%v", err)
	}
	if len(deleted) != 1 || deleted[0] != "cancel-me" {
		t.Fatalf("deleted=%v", deleted)
	}
}