package main

import (
	"bytes"
	"compress/gzip"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"math"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
)

func testLocation(lat, lon int64, accuracy uint64) []byte {
	var out []byte
	out = append(out, writeTag(1, wireVarint)...)
	out = append(out, writeVarint(uint64(lat))...)
	out = append(out, writeTag(2, wireVarint)...)
	out = append(out, writeVarint(uint64(lon))...)
	out = append(out, writeTag(3, wireVarint)...)
	out = append(out, writeVarint(accuracy)...)
	return out
}

func testWifiDevice(loc []byte) []byte {
	mac := []byte("aa:bb:cc:dd:ee:ff")
	var out []byte
	out = append(out, writeLengthDelimited(1, mac)...)
	out = append(out, writeLengthDelimited(2, loc)...)
	return out
}

func testFrame(payload []byte) []byte {
	magic := []byte{0, 1, 0, 0, 0, 1, 0, 0}
	var lenBytes [2]byte
	binary.BigEndian.PutUint16(lenBytes[:], uint16(len(payload)))
	var out []byte
	out = append(out, magic...)
	out = append(out, lenBytes[:]...)
	out = append(out, payload...)
	return out
}

func TestPatchWifiLocation(t *testing.T) {
	payload := writeLengthDelimited(2, testWifiDevice(testLocation(100, 200, 25)))
	body := testFrame(payload)
	c := wlocCoords{Latitude: 31.230416, Longitude: 121.473701, Accuracy: 50}

	patched, stats, err := patchWlocBody(body, c)
	if err != nil {
		t.Fatal(err)
	}
	if stats.WiFi != 1 || stats.Locations != 1 {
		t.Fatalf("unexpected stats: %+v", stats)
	}
	if bytes.Equal(patched, body) {
		t.Fatal("body was not patched")
	}

	newLen := int(binary.BigEndian.Uint16(patched[8:10]))
	newPayload := patched[10 : 10+newLen]
	latBytes := append(writeTag(1, wireVarint), writeVarint(uint64(int64(math.Round(c.Latitude*1e8))))...)
	if !bytes.Contains(newPayload, latBytes) {
		t.Fatal("new latitude bytes not found")
	}
}

func TestPatchCellLocation(t *testing.T) {
	cell := writeLengthDelimited(5, testLocation(300, 400, 25))
	payload := writeLengthDelimited(22, cell)
	body := testFrame(payload)
	c := wlocCoords{Latitude: 22.544577, Longitude: 113.94114, Accuracy: 25}

	patched, stats, err := patchWlocBody(body, c)
	if err != nil {
		t.Fatal(err)
	}
	if stats.Cell != 1 || stats.Locations != 1 {
		t.Fatalf("unexpected stats: %+v", stats)
	}
	if bytes.Equal(patched, body) {
		t.Fatal("body was not patched")
	}
}

func TestPatchGzip(t *testing.T) {
	payload := writeLengthDelimited(2, testWifiDevice(testLocation(100, 200, 25)))
	var buf bytes.Buffer
	zw := gzip.NewWriter(&buf)
	if _, err := zw.Write(testFrame(payload)); err != nil {
		t.Fatal(err)
	}
	if err := zw.Close(); err != nil {
		t.Fatal(err)
	}
	c := wlocCoords{Latitude: 31.230416, Longitude: 121.473701, Accuracy: 50}

	patched, _, err := patchResponseBody(buf.Bytes(), c)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Equal(patched, testFrame(payload)) {
		t.Fatal("gzip body was not patched")
	}
}

func TestTransparentBodyUnchanged(t *testing.T) {
	body := []byte{1, 2, 3, 4}
	_, _, err := patchResponseBody(body, wlocCoords{Latitude: 31.230416, Longitude: 121.473701, Accuracy: 25})
	if err == nil {
		t.Fatal("expected non-patchable body to error")
	}
}

func TestPatchWlocResponsePassesThroughOversizedBody(t *testing.T) {
	payload := bytes.Repeat([]byte("x"), (1<<20)+1)
	req := httptest.NewRequest(http.MethodPost, "https://gs-loc.apple.com/clls/wloc", nil)
	resp := &http.Response{
		StatusCode:    http.StatusOK,
		Request:       req,
		Header:        make(http.Header),
		Body:          io.NopCloser(bytes.NewReader(payload)),
		ContentLength: int64(len(payload)),
	}

	patched := patchWlocResponse(resp, nil)
	got, err := io.ReadAll(patched.Body)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(got, payload) {
		t.Fatal("oversized WLOC response was changed or truncated")
	}
}

func TestServeLocalRequestsKeepsUnrelatedRequestBodyStreaming(t *testing.T) {
	const secret = "body-must-not-be-buffered-or-logged"
	req := httptest.NewRequest(http.MethodPost, "https://example.com/upload", bytes.NewBufferString(secret))
	returned, response := serveLocalRequests(req, nil)
	if response != nil {
		t.Fatalf("unexpected local response: %d", response.StatusCode)
	}
	got, err := io.ReadAll(returned.Body)
	if err != nil {
		t.Fatal(err)
	}
	if string(got) != secret {
		t.Fatalf("request body changed: got %q", got)
	}
	if logs := drainLogs(); bytes.Contains([]byte(logs), []byte(secret)) {
		t.Fatal("request body leaked into diagnostic logs")
	}
}

func TestCoordsEndpointReturnsAtomicSnapshot(t *testing.T) {
	stateMu.Lock()
	previousLat, previousLon := currentLat, currentLon
	previousEnabled, previousAccuracy := currentEnabled, currentAccuracy
	currentLat, currentLon, currentEnabled, currentAccuracy = 0, 0, false, 0
	stateMu.Unlock()
	t.Cleanup(func() {
		stateMu.Lock()
		currentLat, currentLon = previousLat, previousLon
		currentEnabled, currentAccuracy = previousEnabled, previousAccuracy
		stateMu.Unlock()
	})

	handler := newProxy(nil).NonproxyHandler
	const updates = 20_000
	const readers = 8
	const readsPerReader = 2_500

	var writers sync.WaitGroup
	writers.Add(1)
	go func() {
		defer writers.Done()
		for i := 1; i <= updates; i++ {
			stateMu.Lock()
			currentLat = float64(i)
			currentLon = -float64(i)
			currentEnabled = i%2 == 0
			currentAccuracy = i
			stateMu.Unlock()
		}
	}()

	errs := make(chan error, readers)
	var readersGroup sync.WaitGroup
	for range readers {
		readersGroup.Add(1)
		go func() {
			defer readersGroup.Done()
			for i := 0; i < readsPerReader; i++ {
				recorder := httptest.NewRecorder()
				handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "http://proxy.local/coords", nil))
				var snapshot struct {
					Enabled  bool    `json:"enabled"`
					Lat      float64 `json:"lat"`
					Lon      float64 `json:"lon"`
					Accuracy int     `json:"accuracy"`
				}
				if err := json.Unmarshal(recorder.Body.Bytes(), &snapshot); err != nil {
					errs <- err
					return
				}
				if snapshot.Lat != 0 && (snapshot.Lon != -snapshot.Lat || snapshot.Accuracy != int(snapshot.Lat)) {
					errs <- fmt.Errorf("torn coordinate snapshot: %+v", snapshot)
					return
				}
			}
		}()
	}
	readersGroup.Wait()
	writers.Wait()
	close(errs)
	for err := range errs {
		t.Error(err)
	}
}