package api

import (
	"bufio"
	"bytes"
	"compress/gzip"
	"context"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/klauspost/compress/zstd"
	"github.com/router-for-me/CLIProxyAPI/v8/internal/config"
	"github.com/router-for-me/CLIProxyAPI/v8/internal/httpwire"
	log "github.com/sirupsen/logrus"
)

func compressionTestBody(t *testing.T, encoding string, body []byte) []byte {
	t.Helper()
	var reader io.Reader = bytes.NewReader(body)
	switch encoding {
	case "gzip":
		decoder, errNewReader := gzip.NewReader(reader)
		if errNewReader != nil {
			t.Fatal(errNewReader)
		}
		defer func() {
			if errClose := decoder.Close(); errClose != nil {
				log.Errorf("close test gzip decoder: %v", errClose)
			}
		}()
		reader = decoder
	case "zstd":
		decoder, errNewReader := zstd.NewReader(reader)
		if errNewReader != nil {
			t.Fatal(errNewReader)
		}
		defer decoder.Close()
		reader = decoder
	}
	decoded, errReadAll := io.ReadAll(reader)
	if errReadAll != nil {
		t.Fatal(errReadAll)
	}
	return decoded
}

func TestResponseCompression(t *testing.T) {
	body := "{\"message\":\"" + strings.Repeat("compressible text ", 128) + "\"}"
	for _, tc := range []struct {
		name             string
		accept           string
		encoding         string
		body             string
		contentType      string
		existingEncoding string
		status           int
		method           string
		contentLength    bool
		writeString      bool
		forceHeader      bool
	}{
		{name: "zstd", accept: "zstd", encoding: "zstd"},
		{name: "gzip", accept: "gzip", encoding: "gzip"},
		{name: "large_known_length", accept: "gzip", encoding: "gzip", contentLength: true},
		{name: "no_accept"},
		{name: "identity", accept: "identity"},
		{name: "excluded", accept: "gzip;q=0"},
		{name: "preference", accept: "gzip;q=1,zstd;q=0.5", encoding: "gzip"},
		{name: "preference_tie", accept: "zstd;q=0.5,gzip;q=0.5", encoding: "zstd"},
		{name: "small_identity_refused", accept: "gzip, identity;q=0", body: "{}", encoding: "gzip"},
		{name: "small_known_length_identity_refused", accept: "gzip, identity;q=0", body: "{}", encoding: "gzip", contentLength: true},
		{name: "small_wildcard_identity_refused", accept: "gzip, *;q=0", body: "{}", encoding: "gzip"},
		{name: "small_identity_overrides_wildcard", accept: "gzip, *;q=0, identity;q=1", body: "{}"},
		{name: "identity_refused_no_encoding", accept: "identity;q=0,*;q=0", body: "{}"},
		{name: "wildcard", accept: "*;q=0.5", encoding: "zstd"},
		{name: "wildcard_explicit_exclusion", accept: "*;q=1,zstd;q=0", encoding: "gzip"},
		{name: "all_excluded", accept: "*;q=1,zstd;q=0,gzip;q=0"},
		{name: "invalid_quality", accept: "zstd;q=invalid,gzip;q=2"},
		{name: "small", accept: "zstd", body: "{}"},
		{name: "small_known_length", accept: "gzip", body: "{}", contentLength: true},
		{name: "existing_encoding", accept: "gzip", existingEncoding: "br"},
		{name: "binary", accept: "zstd", contentType: "image/png"},
		{name: "json_suffix", accept: "gzip", contentType: "application/problem+json; charset=utf-8", encoding: "gzip"},
		{name: "partial", accept: "gzip", status: http.StatusPartialContent},
		{name: "no_content", accept: "gzip", status: http.StatusNoContent},
		{name: "not_modified", accept: "gzip", status: http.StatusNotModified},
		{name: "informational", accept: "gzip", status: http.StatusEarlyHints},
		{name: "head", accept: "gzip", method: http.MethodHead},
		{name: "string_writer", accept: "gzip", encoding: "gzip", writeString: true},
		{name: "explicit_header", accept: "gzip", body: "{}", forceHeader: true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			payload := body
			if tc.body != "" {
				payload = tc.body
			}
			engine := gin.New()
			engine.Use(responseCompressionMiddleware(&config.Config{}))
			engine.Any("/", func(c *gin.Context) {
				contentType := tc.contentType
				if contentType == "" {
					contentType = "application/json"
				}
				c.Header("Content-Type", contentType)
				c.Header("Vary", "Origin")
				if tc.existingEncoding != "" {
					c.Header("Content-Encoding", tc.existingEncoding)
				}
				if tc.contentLength {
					c.Header("Content-Length", strconv.Itoa(len(payload)))
				}
				if tc.status != 0 {
					c.Status(tc.status)
				}
				if tc.status != 0 && (tc.status < 200 || tc.status == http.StatusNoContent || tc.status == http.StatusNotModified) {
					c.Writer.WriteHeaderNow()
					return
				}
				if tc.forceHeader {
					c.Writer.WriteHeaderNow()
				}
				if tc.writeString {
					_, _ = io.WriteString(c.Writer, payload)
				} else {
					if _, errWrite := c.Writer.Write([]byte(payload[:len(payload)/2])); errWrite != nil {
						t.Errorf("first write: %v", errWrite)
					}
					if _, errWrite := c.Writer.Write([]byte(payload[len(payload)/2:])); errWrite != nil {
						t.Errorf("second write: %v", errWrite)
					}
				}
			})
			method := tc.method
			if method == "" {
				method = http.MethodGet
			}
			req := httptest.NewRequest(method, "/", nil)
			req.Header.Set("Accept-Encoding", tc.accept)
			response := httptest.NewRecorder()
			engine.ServeHTTP(response, req)
			wantEncoding := tc.encoding
			if tc.existingEncoding != "" {
				wantEncoding = tc.existingEncoding
			}
			if got := response.Header().Get("Content-Encoding"); got != wantEncoding {
				t.Fatalf("Content-Encoding = %q, want %q", got, wantEncoding)
			}
			if tc.encoding != "" {
				if got := response.Header().Get("Vary"); !strings.Contains(got, "Origin") || !strings.Contains(got, "Accept-Encoding") {
					t.Fatalf("Vary = %q", got)
				}
				if got := response.Header().Get("Content-Length"); got != "" {
					t.Fatalf("Content-Length = %q", got)
				}
			}
			if tc.contentLength && tc.encoding == "" && response.Header().Get("Content-Length") != strconv.Itoa(len(payload)) {
				t.Fatal("plaintext Content-Length was changed")
			}
			wantBody := payload
			if tc.status != 0 && (tc.status < 200 || tc.status == http.StatusNoContent || tc.status == http.StatusNotModified) {
				wantBody = ""
			}
			if got := string(compressionTestBody(t, tc.encoding, response.Body.Bytes())); got != wantBody {
				t.Fatalf("decoded body differs: got %d bytes, want %d", len(got), len(wantBody))
			}
		})
	}
}

func TestResponseCompressionDisabled(t *testing.T) {
	disabled := false
	for _, cfg := range []*config.Config{nil, {ServerResponseCompression: config.ResponseCompressionConfig{Enabled: &disabled}}} {
		engine := gin.New()
		engine.Use(responseCompressionMiddleware(cfg))
		engine.GET("/", func(c *gin.Context) { c.String(http.StatusOK, strings.Repeat("text", 1024)) })
		req := httptest.NewRequest(http.MethodGet, "/", nil)
		req.Header.Set("Accept-Encoding", "zstd")
		response := httptest.NewRecorder()
		engine.ServeHTTP(response, req)
		if response.Header().Get("Content-Encoding") != "" {
			t.Fatal("disabled middleware compressed response")
		}
		if got := response.Result().Header.Values("Vary"); len(got) != 0 {
			t.Fatalf("disabled middleware added Vary=%q", got)
		}
	}
}

func TestResponseCompressionWebsocketUpgradeSkipped(t *testing.T) {
	engine := gin.New()
	var original gin.ResponseWriter
	engine.Use(func(c *gin.Context) { original = c.Writer; c.Next() })
	engine.Use(responseCompressionMiddleware(&config.Config{}))
	engine.GET("/", func(c *gin.Context) {
		if c.Writer != original {
			t.Error("WebSocket upgrade writer was wrapped")
		}
		c.String(http.StatusOK, strings.Repeat("text", 1024))
	})
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.Header.Set("Upgrade", "websocket")
	req.Header.Set("Accept-Encoding", "gzip")
	response := httptest.NewRecorder()
	engine.ServeHTTP(response, req)
	if response.Header().Get("Content-Encoding") != "" {
		t.Fatal("WebSocket upgrade response was compressed")
	}
	if got := response.Result().Header.Values("Vary"); len(got) != 0 {
		t.Fatalf("WebSocket upgrade middleware added Vary=%q", got)
	}
}

func TestResponseCompressionSSEIncremental(t *testing.T) {
	for _, encoding := range []string{"zstd", "gzip"} {
		t.Run(encoding, func(t *testing.T) {
			advance := make(chan struct{})
			engine := gin.New()
			engine.Use(responseCompressionMiddleware(&config.Config{}))
			events := []string{"data: first\n\n", "data: second\n\n", "data: third\n\n"}
			engine.GET("/", func(c *gin.Context) {
				c.Header("Content-Type", "text/event-stream")
				for _, event := range events {
					if _, errWrite := io.WriteString(c.Writer, event); errWrite != nil {
						return
					}
					if errFlush := httpwire.FlushResponse(c.Writer); errFlush != nil {
						return
					}
					select {
					case <-advance:
					case <-c.Request.Context().Done():
						return
					}
				}
			})
			server := httptest.NewServer(engine)
			defer server.Close()
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()
			req, errNewRequest := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil)
			if errNewRequest != nil {
				t.Fatal(errNewRequest)
			}
			req.Header.Set("Accept-Encoding", encoding)
			response, errDo := server.Client().Do(req)
			if errDo != nil {
				t.Fatal(errDo)
			}
			defer func() {
				if errClose := response.Body.Close(); errClose != nil {
					log.Errorf("close SSE test response: %v", errClose)
				}
			}()
			if got := response.Header.Get("Content-Encoding"); got != encoding {
				t.Fatalf("Content-Encoding = %q, want %q", got, encoding)
			}
			var reader io.Reader = response.Body
			if encoding == "gzip" {
				decoder, errNewReader := gzip.NewReader(reader)
				if errNewReader != nil {
					t.Fatal(errNewReader)
				}
				defer func() {
					if errClose := decoder.Close(); errClose != nil {
						log.Errorf("close SSE gzip decoder: %v", errClose)
					}
				}()
				reader = decoder
			} else {
				decoder, errNewReader := zstd.NewReader(reader, zstd.WithDecoderConcurrency(1))
				if errNewReader != nil {
					t.Fatal(errNewReader)
				}
				defer decoder.Close()
				reader = decoder
			}
			buffered := bufio.NewReader(reader)
			var decoded strings.Builder
			for _, event := range events {
				line, errReadString := buffered.ReadString('\n')
				if errReadString != nil {
					t.Fatal(errReadString)
				}
				blank, errReadBlank := buffered.ReadString('\n')
				if errReadBlank != nil {
					t.Fatal(errReadBlank)
				}
				if line+blank != event {
					t.Fatalf("event = %q, want %q", line+blank, event)
				}
				decoded.WriteString(line + blank)
				advance <- struct{}{}
			}
			tail, errReadAll := io.ReadAll(buffered)
			if errReadAll != nil {
				t.Fatal(errReadAll)
			}
			decoded.Write(tail)
			if decoded.String() != strings.Join(events, "") {
				t.Fatal("decoded SSE differs")
			}
		})
	}
}

type deterministicResponseCompressorPool struct {
	newCompressor func() any
	item          any
}

func (p *deterministicResponseCompressorPool) Get() any {
	if p.item == nil {
		return p.newCompressor()
	}
	item := p.item
	p.item = nil
	return item
}

func (p *deterministicResponseCompressorPool) Put(item any) {
	p.item = item
}

func TestResponseCompressionReusesPooledEncoders(t *testing.T) {
	bodies := []string{
		"{\"message\":\"" + strings.Repeat("first response ", 200) + "\"}",
		"{\"message\":\"" + strings.Repeat("second, different response ", 150) + "\"}",
	}
	engine := gin.New()
	engine.Use(responseCompressionMiddleware(&config.Config{}))
	body := ""
	engine.GET("/", func(c *gin.Context) {
		c.Data(http.StatusOK, "application/json", []byte(body))
	})
	for _, encoding := range []string{"zstd", "gzip"} {
		original := responseCompressorPools[encoding]
		created := 0
		responseCompressorPools[encoding] = &deterministicResponseCompressorPool{
			newCompressor: func() any { created++; return original.Get() },
		}
		t.Cleanup(func() { responseCompressorPools[encoding] = original })
		const rounds = 16
		for round := 0; round < rounds; round++ {
			if round%2 == 1 {
				// A failed client write must not poison the encoder handed to the next response.
				context, _ := gin.CreateTestContext(httptest.NewRecorder())
				context.Header("Content-Type", "application/json")
				writer := &responseCompressionWriter{ResponseWriter: &compressionErrorWriter{ResponseWriter: context.Writer}, encoding: encoding, minBytes: 1}
				_, _ = writer.WriteString(bodies[0])
				writer.finish()
				continue
			}
			body = bodies[(round/2)%len(bodies)]
			req := httptest.NewRequest(http.MethodGet, "/", nil)
			req.Header.Set("Accept-Encoding", encoding)
			rec := httptest.NewRecorder()
			engine.ServeHTTP(rec, req)
			if got := rec.Header().Get("Content-Encoding"); got != encoding {
				t.Fatalf("%s round %d: Content-Encoding = %q", encoding, round, got)
			}
			if got := string(compressionTestBody(t, encoding, rec.Body.Bytes())); got != body {
				t.Fatalf("%s round %d: decoded body mismatch", encoding, round)
			}
		}
		if created != 1 {
			t.Fatalf("%s: created %d encoders for %d responses, want 1", encoding, created, rounds)
		}
	}
}

func TestResponseCompressionReusesEncoderAfterWriteFailure(t *testing.T) {
	for _, encoding := range []string{"zstd", "gzip"} {
		t.Run(encoding, func(t *testing.T) {
			original := responseCompressorPools[encoding]
			responseCompressorPools[encoding] = &deterministicResponseCompressorPool{newCompressor: original.Get}
			t.Cleanup(func() { responseCompressorPools[encoding] = original })

			failedContext, _ := gin.CreateTestContext(httptest.NewRecorder())
			failedContext.Header("Content-Type", "text/event-stream")
			failedWriter := &responseCompressionWriter{
				ResponseWriter: &compressionErrorWriter{ResponseWriter: failedContext.Writer},
				encoding:       encoding,
				minBytes:       1,
			}
			_, errWrite := failedWriter.WriteString("data: failed\n\n")
			if errWrite == nil {
				errWrite = failedWriter.FlushError()
			}
			if !errors.Is(errWrite, io.ErrClosedPipe) {
				t.Fatalf("failed response error = %v, want %v", errWrite, io.ErrClosedPipe)
			}
			failedCompressor := failedWriter.compressor
			if failedCompressor == nil {
				t.Fatal("failed response has no compressor")
			}
			failedWriter.finish()

			recorder := httptest.NewRecorder()
			context, _ := gin.CreateTestContext(recorder)
			context.Header("Content-Type", "application/json")
			writer := &responseCompressionWriter{ResponseWriter: context.Writer, encoding: encoding, minBytes: 1}
			body := `{"message":"response after failure"}`
			if _, errWrite = writer.WriteString(body); errWrite != nil {
				t.Fatal(errWrite)
			}
			if writer.compressor != failedCompressor {
				t.Fatal("the encoder from the failed response was not reused")
			}
			writer.finish()
			if got := string(compressionTestBody(t, encoding, recorder.Body.Bytes())); got != body {
				t.Fatalf("decoded body = %q, want %q", got, body)
			}
		})
	}
}

func TestResponseCompressionWriteAfterFinishFails(t *testing.T) {
	for _, encoding := range []string{"zstd", "gzip"} {
		context, _ := gin.CreateTestContext(httptest.NewRecorder())
		context.Header("Content-Type", "text/event-stream")
		writer := &responseCompressionWriter{ResponseWriter: context.Writer, encoding: encoding, minBytes: 1}
		if _, errWrite := writer.WriteString("data: one\n\n"); errWrite != nil {
			t.Fatal(errWrite)
		}
		writer.finish()
		if _, errWrite := writer.WriteString("data: late\n\n"); errWrite == nil {
			t.Fatalf("%s: a write after finish must fail, not go out uncompressed", encoding)
		}
	}
}