package clients

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"reflect"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

func TestGeminiVeoRunUsesOfficialLongRunningProtocol(t *testing.T) {
	const operationName = "models/veo-3.1-generate-preview/operations/test-operation"
	videoPayload := []byte("test mp4 payload")
	var server *httptest.Server
	var pollCount atomic.Int32
	server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("x-goog-api-key") != "test-gemini-key" {
			t.Errorf("missing Gemini API key header for %s", r.URL.Path)
		}
		switch {
		case r.Method == http.MethodPost && r.URL.Path == "/v1beta/models/veo-3.1-generate-preview:predictLongRunning":
			var body map[string]any
			if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
				t.Fatalf("decode Veo body: %v", err)
			}
			instances := body["instances"].([]any)
			if prompt := instances[0].(map[string]any)["prompt"]; prompt != "a paper boat on a lake" {
				t.Errorf("unexpected prompt: %v", prompt)
			}
			parameters := body["parameters"].(map[string]any)
			if parameters["durationSeconds"] != float64(4) || parameters["resolution"] != "720p" || parameters["aspectRatio"] != "16:9" {
				t.Errorf("unexpected parameters: %+v", parameters)
			}
			_ = json.NewEncoder(w).Encode(map[string]any{"name": operationName})
		case r.Method == http.MethodGet && r.URL.Path == "/v1beta/"+operationName:
			if pollCount.Add(1) == 1 {
				_ = json.NewEncoder(w).Encode(map[string]any{"name": operationName, "done": false})
				return
			}
			_ = json.NewEncoder(w).Encode(map[string]any{
				"name": operationName,
				"done": true,
				"response": map[string]any{"generateVideoResponse": map[string]any{"generatedSamples": []any{
					map[string]any{"video": map[string]any{"uri": server.URL + "/v1beta/files/generated-video:download", "mimeType": "video/mp4"}},
				}}},
			})
		case r.Method == http.MethodGet && r.URL.Path == "/v1beta/files/generated-video:download":
			w.Header().Set("Content-Type", "video/mp4")
			_, _ = w.Write(videoPayload)
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()

	var submitted string
	polled := 0
	response, err := (GeminiClient{}).Run(context.Background(), Request{
		Kind:      "videos.generations",
		ModelType: "video_generate",
		Model:     "veo-3.1",
		Body: map[string]any{
			"prompt":       "a paper boat on a lake",
			"duration":     4,
			"resolution":   "720p",
			"aspect_ratio": "16:9",
		},
		Candidate: testGeminiVeoCandidate(server.URL),
		OnRemoteTaskSubmitted: func(remoteTaskID string, _ map[string]any) error {
			submitted = remoteTaskID
			return nil
		},
		OnRemoteTaskPolled: func(remoteTaskID string, _ map[string]any) error {
			if remoteTaskID != operationName {
				t.Errorf("unexpected polled operation: %s", remoteTaskID)
			}
			polled++
			return nil
		},
	})
	if err != nil {
		t.Fatalf("run Gemini Veo: %v", err)
	}
	if submitted != operationName || polled != 2 || pollCount.Load() != 2 {
		t.Fatalf("unexpected async callbacks: submitted=%q polled=%d requests=%d", submitted, polled, pollCount.Load())
	}
	if response.UpstreamProtocol != ProtocolGeminiVeo || response.RequestID != operationName {
		t.Fatalf("unexpected response metadata: %+v", response)
	}
	data := response.Result["data"].([]any)
	item := data[0].(map[string]any)
	if got := item["video_bytes"].([]byte); !reflect.DeepEqual(got, videoPayload) {
		t.Fatalf("unexpected video payload: %q", got)
	}
	if _, exposed := item["uri"]; exposed {
		t.Fatalf("upstream authenticated URI must not be exposed: %+v", item)
	}
}

func TestGeminiVeoRunResumesOperationWithoutSubmittingAgain(t *testing.T) {
	const operationName = "models/veo-3.1-generate-preview/operations/resumed"
	var postCount atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodPost {
			postCount.Add(1)
			http.Error(w, "unexpected submit", http.StatusInternalServerError)
			return
		}
		_ = json.NewEncoder(w).Encode(map[string]any{
			"name": operationName,
			"done": true,
			"response": map[string]any{"generateVideoResponse": map[string]any{"generatedSamples": []any{
				map[string]any{"video": map[string]any{"videoBytes": base64.StdEncoding.EncodeToString([]byte("resumed video")), "mimeType": "video/mp4"}},
			}}},
		})
	}))
	defer server.Close()

	response, err := (GeminiClient{}).Run(context.Background(), Request{
		Kind:         "videos.generations",
		ModelType:    "video_generate",
		Model:        "veo-3.1",
		Body:         map[string]any{"prompt": "ignored on resume"},
		Candidate:    testGeminiVeoCandidate(server.URL),
		RemoteTaskID: operationName,
	})
	if err != nil {
		t.Fatalf("resume Gemini Veo: %v", err)
	}
	if postCount.Load() != 0 || response.Result["id"] != operationName {
		t.Fatalf("resume submitted a new operation or returned wrong ID: posts=%d result=%+v", postCount.Load(), response.Result)
	}
}

func TestGeminiVeoBodyMapsImageAndParameters(t *testing.T) {
	first := "data:image/png;base64," + base64.StdEncoding.EncodeToString([]byte("first"))
	last := "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString([]byte("last"))
	reference := "data:image/webp;base64," + base64.StdEncoding.EncodeToString([]byte("reference"))
	body, err := geminiVeoBody(context.Background(), Request{Body: map[string]any{
		"prompt":            "animate",
		"first_frame":       first,
		"last_frame":        last,
		"reference_images":  []any{reference},
		"duration":          8,
		"resolution":        "2160p",
		"aspect_ratio":      "9:16",
		"negative_prompt":   "text",
		"person_generation": "allow_adult",
		"enhance_prompt":    false,
	}})
	if err != nil {
		t.Fatalf("build Gemini Veo body: %v", err)
	}
	instance := body["instances"].([]any)[0].(map[string]any)
	image := instance["image"].(map[string]any)
	lastFrame := instance["lastFrame"].(map[string]any)
	references := instance["referenceImages"].([]any)
	if image["bytesBase64Encoded"] != base64.StdEncoding.EncodeToString([]byte("first")) || image["mimeType"] != "image/png" {
		t.Fatalf("unexpected first frame: %+v", image)
	}
	if lastFrame["mimeType"] != "image/jpeg" || len(references) != 1 {
		t.Fatalf("unexpected last/reference images: last=%+v references=%+v", lastFrame, references)
	}
	parameters := body["parameters"].(map[string]any)
	if parameters["resolution"] != "4k" || parameters["durationSeconds"] != 8 || parameters["enhancePrompt"] != false {
		t.Fatalf("unexpected Gemini Veo parameters: %+v", parameters)
	}
}

func TestGeminiVeoDownloadsImageURLWhenBuildingOfficialPayload(t *testing.T) {
	payload := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 1, 2, 3, 4}
	fetchCount := 0
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fetchCount++
		w.Header().Set("Content-Type", "image/png")
		_, _ = w.Write(payload)
	}))
	defer server.Close()

	body, err := geminiVeoBody(context.Background(), Request{Body: map[string]any{
		"prompt":      "animate",
		"first_frame": server.URL + "/input.png",
		"duration":    8,
	}})
	if err != nil {
		t.Fatalf("build Gemini Veo body from URL: %v", err)
	}
	instance := body["instances"].([]any)[0].(map[string]any)
	image := instance["image"].(map[string]any)
	if fetchCount != 1 || image["bytesBase64Encoded"] != base64.StdEncoding.EncodeToString(payload) || image["mimeType"] != "image/png" {
		t.Fatalf("Veo URL should be fetched during official payload construction: fetches=%d image=%+v", fetchCount, image)
	}
}

func TestGeminiOfficialAPIClassification(t *testing.T) {
	tests := []struct {
		baseURL string
		want    bool
	}{
		{baseURL: "", want: true},
		{baseURL: "https://generativelanguage.googleapis.com", want: true},
		{baseURL: "https://generativelanguage.googleapis.com/v1beta", want: true},
		{baseURL: "https://generativelanguage.googleapis.com/v1beta/openai", want: false},
		{baseURL: "https://gemini-compatible.example.com/v1beta", want: false},
	}
	for _, test := range tests {
		if got := GeminiUsesOfficialAPI(store.RuntimeModelCandidate{BaseURL: test.baseURL}); got != test.want {
			t.Fatalf("GeminiUsesOfficialAPI(%q) = %t, want %t", test.baseURL, got, test.want)
		}
	}
}

func TestGeminiCompatibleConvertsFileDataURLToInlineBase64(t *testing.T) {
	payload := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 1, 2, 3, 4}
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "image/png")
		_, _ = w.Write(payload)
	}))
	defer server.Close()
	body := map[string]any{"contents": []any{map[string]any{"parts": []any{map[string]any{
		"fileData": map[string]any{"fileUri": server.URL + "/source.png", "mimeType": "image/png"},
	}}}}}

	prepared, err := prepareCompatibleGeminiInlineImages(context.Background(), body)
	if err != nil {
		t.Fatalf("prepare compatible Gemini inline image: %v", err)
	}
	contents := prepared["contents"].([]any)
	parts := contents[0].(map[string]any)["parts"].([]any)
	part := parts[0].(map[string]any)
	inline := part["inlineData"].(map[string]any)
	if _, exists := part["fileData"]; exists || inline["mimeType"] != "image/png" || inline["data"] != base64.StdEncoding.EncodeToString(payload) {
		t.Fatalf("compatible Gemini should use inlineData: %+v", part)
	}
}

func TestGeminiOfficialImageUploadUsesFilesAPI(t *testing.T) {
	imagePayload := []byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n', 1, 2, 3, 4}
	imageServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "image/png")
		_, _ = w.Write(imagePayload)
	}))
	defer imageServer.Close()

	var uploadServer *httptest.Server
	startCount := 0
	uploadCount := 0
	uploadServer = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/upload/v1beta/files":
			startCount++
			if r.Header.Get("x-goog-api-key") != "gemini-key" || r.Header.Get("X-Goog-Upload-Protocol") != "resumable" || r.Header.Get("X-Goog-Upload-Header-Content-Type") != "image/png" {
				t.Fatalf("unexpected Gemini Files start headers: %+v", r.Header)
			}
			w.Header().Set("X-Goog-Upload-URL", uploadServer.URL+"/upload/session")
			w.WriteHeader(http.StatusOK)
		case "/upload/session":
			uploadCount++
			raw, err := io.ReadAll(r.Body)
			if err != nil {
				t.Fatalf("read Gemini Files upload: %v", err)
			}
			if string(raw) != string(imagePayload) || r.Header.Get("X-Goog-Upload-Command") != "upload, finalize" {
				t.Fatalf("unexpected Gemini Files upload body=%v headers=%+v", raw, r.Header)
			}
			_ = json.NewEncoder(w).Encode(map[string]any{"file": map[string]any{
				"uri":      "https://generativelanguage.googleapis.com/v1beta/files/easyai-image",
				"mimeType": "image/png",
			}})
		default:
			http.NotFound(w, r)
		}
	}))
	defer uploadServer.Close()

	fileData, err := uploadOfficialGeminiImageFile(context.Background(), uploadServer.Client(), store.RuntimeModelCandidate{BaseURL: uploadServer.URL + "/v1beta"}, "gemini-key", imageServer.URL+"/source.png", "image/png")
	if err != nil {
		t.Fatalf("upload official Gemini image: %v", err)
	}
	if startCount != 1 || uploadCount != 1 || fileData["fileUri"] != "https://generativelanguage.googleapis.com/v1beta/files/easyai-image" || fileData["mimeType"] != "image/png" {
		t.Fatalf("unexpected Gemini Files result: starts=%d uploads=%d fileData=%+v", startCount, uploadCount, fileData)
	}
}

func TestGeminiVeoRejectsInvalidDurationBeforeSubmission(t *testing.T) {
	_, err := geminiVeoBody(context.Background(), Request{Body: map[string]any{"prompt": "test", "duration": 5}})
	if err == nil || ErrorCode(err) != "invalid_parameter" || ErrorParam(err) != "duration" || !strings.Contains(err.Error(), "4, 6, or 8") {
		t.Fatalf("unexpected invalid duration error: %v", err)
	}
}

func TestGeminiVeoDownloadURLOnlyTrustsProviderAndGoogleAPIs(t *testing.T) {
	for _, test := range []struct {
		uri     string
		trusted bool
	}{
		{uri: "https://gemini-proxy.example.com/v1beta/files/video", trusted: true},
		{uri: "https://generativelanguage.googleapis.com/v1beta/files/video", trusted: true},
		{uri: "https://storage.googleapis.com/signed/video", trusted: false},
		{uri: "https://attacker.example.com/video", trusted: false},
	} {
		_, trusted, err := geminiVeoDownloadURL("https://gemini-proxy.example.com/v1beta", test.uri)
		if err != nil {
			t.Fatalf("parse %s: %v", test.uri, err)
		}
		if trusted != test.trusted {
			t.Fatalf("unexpected trust for %s: got %t want %t", test.uri, trusted, test.trusted)
		}
	}
}

func TestGeminiVeoLive(t *testing.T) {
	apiKey := strings.TrimSpace(os.Getenv("GEMINI_VEO_LIVE_API_KEY"))
	if apiKey == "" {
		t.Skip("GEMINI_VEO_LIVE_API_KEY is not configured")
	}
	baseURL := strings.TrimSpace(os.Getenv("GEMINI_VEO_LIVE_BASE_URL"))
	model := strings.TrimSpace(os.Getenv("GEMINI_VEO_LIVE_MODEL"))
	if model == "" {
		model = "veo-3.1-generate-preview"
	}
	ctx, cancel := context.WithTimeout(context.Background(), 12*time.Minute)
	defer cancel()
	operationName := strings.TrimSpace(os.Getenv("GEMINI_VEO_LIVE_OPERATION"))
	response, err := (GeminiClient{}).Run(ctx, Request{
		Kind:      "videos.generations",
		ModelType: "video_generate",
		Model:     model,
		Body: map[string]any{
			"prompt":       "A small red paper boat gently floating on calm water, static camera, natural daylight.",
			"duration":     4,
			"resolution":   "720p",
			"aspect_ratio": "16:9",
		},
		Candidate: store.RuntimeModelCandidate{
			Provider:          "gemini",
			BaseURL:           baseURL,
			ProviderModelName: model,
			Credentials:       map[string]any{"apiKey": apiKey},
			PlatformConfig:    map[string]any{"pollIntervalMs": 10000, "pollTimeoutMs": 660000},
		},
		RemoteTaskID: operationName,
		OnRemoteTaskSubmitted: func(remoteTaskID string, _ map[string]any) error {
			t.Logf("live Gemini Veo submitted operation %s", remoteTaskID)
			return nil
		},
	})
	if err != nil {
		t.Fatalf("live Gemini Veo request failed: %v", err)
	}
	data, _ := response.Result["data"].([]any)
	if len(data) != 1 {
		t.Fatalf("live Gemini Veo returned %d videos", len(data))
	}
	item, _ := data[0].(map[string]any)
	payload, _ := item["video_bytes"].([]byte)
	if len(payload) < 1024 {
		t.Fatalf("live Gemini Veo video is unexpectedly small: %d bytes", len(payload))
	}
	if mimeType := strings.TrimSpace(stringFromAny(item["mime_type"])); !strings.HasPrefix(mimeType, "video/") {
		t.Fatalf("live Gemini Veo returned unexpected MIME type: %q", mimeType)
	}
	t.Logf("live Gemini Veo accepted operation %s and downloaded %d video bytes", response.RequestID, len(payload))
}

func testGeminiVeoCandidate(baseURL string) store.RuntimeModelCandidate {
	return store.RuntimeModelCandidate{
		Provider:          "gemini",
		BaseURL:           baseURL,
		ProviderModelName: "veo-3.1-generate-preview",
		Credentials:       map[string]any{"apiKey": "test-gemini-key"},
		PlatformConfig:    map[string]any{"pollIntervalMs": 1, "pollTimeoutMs": 1000},
	}
}