// Copyright (c) 2026 Storm Dragon
// Authors: Storm Dragon, Angus Kola
// SPDX-License-Identifier: GPL-3.0-or-later

package ocr

import (
	"bytes"
	"context"
	"encoding/binary"
	"encoding/json"
	"errors"
	"hash/crc32"
	"image"
	"image/color"
	"io"
	"math"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"syscall"
	"testing"
	"time"
)

func capture_request() CaptureRequest {
	return CaptureRequest{Kind: ObjectSource, Bounds: image.Rect(-20, 10, 30, 50), Foreground: 123, Generation: 7}
}

func capture_frame(t *testing.T, header CaptureHeader, pixels image.Image) []byte {
	t.Helper()
	var output bytes.Buffer
	if err := WriteCaptureResponse(&output, header, pixels); err != nil {
		t.Fatal(err)
	}
	return output.Bytes()
}

func TestCaptureRequestRejectsUnownedOrInvalidSources(t *testing.T) {
	base := capture_request()
	if err := base.Validate(); err != nil {
		t.Fatal(err)
	}
	mutations := []func(*CaptureRequest){
		func(r *CaptureRequest) { r.Foreground = 0 },
		func(r *CaptureRequest) { r.Foreground = 1 << 32 },
		func(r *CaptureRequest) { r.Generation = 0 },
		func(r *CaptureRequest) { r.Kind = "unknown" },
		func(r *CaptureRequest) { r.Bounds = image.Rectangle{} },
		func(r *CaptureRequest) { r.Bounds.Min.X = math.MinInt32 },
		func(r *CaptureRequest) { r.Bounds = image.Rect(0, 0, MaxPixels+1, 1) },
		func(r *CaptureRequest) { r.Window = 42 },
		func(r *CaptureRequest) { r.Kind = WindowSource },
	}
	for _, mutate := range mutations {
		request := base
		mutate(&request)
		if !errors.Is(request.Validate(), ErrSource) {
			t.Fatalf("invalid request accepted: %+v", request)
		}
	}
	for _, kind := range []SourceKind{WindowSource, DesktopSource, ClipboardSource} {
		request := base
		request.Kind, request.Bounds = kind, image.Rectangle{}
		if err := request.Validate(); err != nil {
			t.Fatal(err)
		}
	}
	data, _ := json.Marshal(base)
	request, err := ReadCaptureRequest(bytes.NewReader(data))
	if err != nil || request != base {
		t.Fatalf("request roundtrip: %+v, %v", request, err)
	}
	for _, invalid := range [][]byte{nil, []byte("{}"), append(append([]byte(nil), data...), data...), []byte(`{"Private":"generic"}`), bytes.Repeat([]byte{' '}, maxCaptureHeader+1)} {
		if _, err := ReadCaptureRequest(bytes.NewReader(invalid)); err == nil {
			t.Fatal("invalid private request accepted")
		}
	}
}

func TestCaptureFrameClippingAndGeneration(t *testing.T) {
	request := capture_request()
	header := CaptureHeader{Request: request, SourceWindow: request.Foreground,
		RootBounds: image.Rect(0, 0, 100, 100), SourceBounds: image.Rect(-10, 0, 80, 80), ScreenBounds: image.Rect(0, 10, 30, 50)}
	pixels := image.NewRGBA(image.Rect(0, 0, 30, 40))
	pixels.Set(1, 2, color.RGBA{12, 34, 56, 255})
	frame := capture_frame(t, header, pixels)
	result, err := decode_capture(frame, request)
	if err != nil {
		t.Fatal(err)
	}
	if color.RGBAModel.Convert(result.Pixels.At(1, 2)) != pixels.At(1, 2) {
		t.Fatal("pixels changed")
	}
	if _, err := decode_capture(append(append([]byte(nil), frame...), 'x'), request); !errors.Is(err, ErrCapture) {
		t.Fatalf("trailing payload accepted: %v", err)
	}
	mapped, ok := result.Transform.ScreenRectangle(image.Rect(1, 2, 3, 4))
	if !ok || mapped != image.Rect(1, 12, 3, 14) {
		t.Fatalf("mapping: %v, %v", mapped, ok)
	}
	request.Generation++
	if _, err := decode_capture(frame, request); !errors.Is(err, ErrCapture) {
		t.Fatalf("old generation accepted: %v", err)
	}
	request.Generation--
	for _, change := range []func(*CaptureHeader){
		func(h *CaptureHeader) { h.ScreenBounds.Min.X = -1 },
		func(h *CaptureHeader) { h.SourceBounds = image.Rect(50, 50, 80, 80) },
		func(h *CaptureHeader) { h.SourceWindow++ },
		func(h *CaptureHeader) { h.ScreenBounds.Max.X++ },
	} {
		bad := header
		change(&bad)
		if _, err := decode_capture(capture_frame(t, bad, pixels), request); err == nil {
			t.Fatalf("invalid mapping accepted: %+v", bad)
		}
	}
}

func TestCaptureClipboardDoesNotMapToScreen(t *testing.T) {
	request := CaptureRequest{Kind: ClipboardSource, Foreground: 123, Generation: 1}
	header := CaptureHeader{Request: request}
	pixels := image.NewRGBA(image.Rect(0, 0, 10, 10))
	result, err := decode_capture(capture_frame(t, header, pixels), request)
	if err != nil || result.Transform.HasScreenMapping {
		t.Fatalf("clipboard: %+v, %v", result, err)
	}
	header.ScreenBounds = image.Rect(0, 0, 10, 10)
	if _, err := decode_capture(capture_frame(t, header, pixels), request); err == nil {
		t.Fatal("clipboard gained coordinates")
	}
}

func TestCaptureProtocolErrorsAndPixelLimit(t *testing.T) {
	request := capture_request()
	for status, want := range map[string]error{"unavailable": ErrCapture, "source": ErrSource, "changed": ErrSourceChanged, "limit": ErrLimit, "Private diagnostics": ErrCapture} {
		frame := capture_frame(t, CaptureHeader{Request: request, Status: status}, nil)
		_, err := decode_capture(frame, request)
		if !errors.Is(err, want) || strings.Contains(err.Error(), "Private") {
			t.Fatalf("status %q: %v", status, err)
		}
		if _, err := decode_capture(append(frame, 'x'), request); !errors.Is(err, ErrCapture) {
			t.Fatal("failure with payload accepted")
		}
	}
	clipboard := CaptureRequest{Kind: ClipboardSource, Foreground: 123, Generation: 1}
	frame := capture_frame(t, CaptureHeader{Request: clipboard}, image.NewRGBA(image.Rect(0, 0, 1, 1)))
	start := bytes.IndexByte(frame, '\n') + 1
	// A valid CRC and malicious IHDR are rejected before allocating pixel data.
	binary.BigEndian.PutUint32(frame[start+16:start+20], 10000)
	binary.BigEndian.PutUint32(frame[start+20:start+24], 10000)
	binary.BigEndian.PutUint32(frame[start+29:start+33], crc32.ChecksumIEEE(frame[start+12:start+29]))
	if _, err := decode_capture(frame, clipboard); !errors.Is(err, ErrLimit) {
		t.Fatalf("oversize PNG: %v", err)
	}
	for _, invalid := range [][]byte{nil, []byte("{}\nnot a png"), bytes.Repeat([]byte{'x'}, maxCaptureHeader+1)} {
		if _, err := decode_capture(invalid, request); !errors.Is(err, ErrCapture) {
			t.Fatalf("invalid response: %v", err)
		}
	}
}

func TestCapturerMissingHelperAndCancellation(t *testing.T) {
	request := capture_request()
	if _, err := NewCapturer("/nonexistent-ocr-capture-helper").Capture(context.Background(), request); !errors.Is(err, ErrCapture) {
		t.Fatalf("missing helper: %v", err)
	}
	executable, err := os.Executable()
	if err != nil {
		t.Fatal(err)
	}
	readyRead, readyWrite, err := os.Pipe()
	if err != nil {
		t.Fatal(err)
	}
	defer readyRead.Close()
	defer readyWrite.Close()
	capturer := NewCapturer(executable)
	capturer.command = func(ctx context.Context, args ...string) *exec.Cmd {
		command := exec.CommandContext(ctx, executable, "-test.run=^TestCaptureChildProcess$", "--", args[0])
		command.Env = append(os.Environ(), "CTHULHU_CAPTURE_TEST_CHILD=1")
		command.ExtraFiles = []*os.File{readyWrite}
		return command
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	completed := make(chan error, 1)
	go func() { _, err := capturer.Capture(ctx, request); completed <- err }()
	_ = readyRead.SetReadDeadline(time.Now().Add(5 * time.Second))
	data := make([]byte, 64)
	count, err := readyRead.Read(data)
	if err != nil {
		t.Fatal(err)
	}
	pid, err := strconv.Atoi(strings.TrimSpace(string(data[:count])))
	if err != nil {
		t.Fatal(err)
	}
	waitContext, waitCancel := context.WithCancel(context.Background())
	waitCancel()
	if _, err := capturer.Capture(waitContext, request); !errors.Is(err, context.Canceled) {
		t.Fatalf("waiting cancellation: %v", err)
	}
	cancel()
	select {
	case err := <-completed:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("running cancellation: %v", err)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("capture did not join cancelled helper")
	}
	if err := syscall.Kill(pid, 0); !errors.Is(err, syscall.ESRCH) {
		t.Fatalf("helper still alive or unreaped: %v", err)
	}
}

func TestCaptureChildProcess(t *testing.T) {
	if os.Getenv("CTHULHU_CAPTURE_TEST_CHILD") != "1" {
		return
	}
	if _, err := io.ReadAll(os.Stdin); err != nil {
		os.Exit(20)
	}
	ready := os.NewFile(3, "ready")
	_, _ = ready.WriteString(strconv.Itoa(os.Getpid()) + "\n")
	_ = ready.Close()
	time.Sleep(time.Hour)
	os.Exit(0)
}