package mockprovider

import (
	"bytes"
	"encoding/json"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"testing"
)

func TestChatSuccessAndErrorMatrix(t *testing.T) {
	handler := Handler{}
	for _, test := range []struct {
		prompt string
		status int
	}{{"hello", 200}, {"mock:429", 429}, {"mock:500", 500}, {"mock:400", 400}, {"mock:401", 401}, {"mock:policy", 400}} {
		t.Run(test.prompt, func(t *testing.T) {
			body, _ := json.Marshal(map[string]any{"messages": []any{map[string]string{"content": test.prompt}}})
			request := httptest.NewRequest(http.MethodPost, "http://mock/v1/chat/completions", bytes.NewReader(body))
			response := httptest.NewRecorder()
			handler.ServeHTTP(response, request)
			if response.Code != test.status {
				t.Fatalf("status=%d body=%s", response.Code, response.Body.String())
			}
		})
	}
}
func TestImageSuccess(t *testing.T) {
	var body bytes.Buffer
	writer := multipart.NewWriter(&body)
	_ = writer.WriteField("prompt", "edit")
	part, _ := writer.CreateFormFile("image", "input.png")
	part.Write(mockPNG())
	writer.Close()
	request := httptest.NewRequest(http.MethodPost, "http://mock/v1/images/edits", &body)
	request.Header.Set("Content-Type", writer.FormDataContentType())
	response := httptest.NewRecorder()
	Handler{}.ServeHTTP(response, request)
	if response.Code != 200 || !bytes.Contains(response.Body.Bytes(), []byte("b64_json")) {
		t.Fatalf("status=%d body=%s", response.Code, response.Body.String())
	}
}

func TestImagesGenerationAndGeminiEndpoints(t *testing.T) {
	imageRequest := httptest.NewRequest(http.MethodPost, "http://mock/v1/images/generations", bytes.NewBufferString(`{"model":"mock","prompt":"draw","n":1}`))
	imageRequest.Header.Set("Content-Type", "application/json")
	imageResponse := httptest.NewRecorder()
	Handler{}.ServeHTTP(imageResponse, imageRequest)
	if imageResponse.Code != http.StatusOK || !bytes.Contains(imageResponse.Body.Bytes(), []byte("b64_json")) {
		t.Fatalf("images status=%d body=%s", imageResponse.Code, imageResponse.Body.String())
	}

	geminiRequest := httptest.NewRequest(http.MethodPost, "http://mock/v1/models/mock:generateContent", bytes.NewBufferString(`{"contents":[{"parts":[{"text":"hello"}]}]}`))
	geminiResponse := httptest.NewRecorder()
	Handler{}.ServeHTTP(geminiResponse, geminiRequest)
	if geminiResponse.Code != http.StatusOK || !bytes.Contains(geminiResponse.Body.Bytes(), []byte("candidates")) {
		t.Fatalf("gemini status=%d body=%s", geminiResponse.Code, geminiResponse.Body.String())
	}
}