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

package ocr

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"image"
	"image/png"
	"io"
	"os/exec"
	"time"
)

var (
	ErrCapture       = errors.New("OCR capture unavailable")
	ErrSource        = errors.New("OCR source unavailable")
	ErrSourceChanged = errors.New("OCR source changed")
)

type SourceKind string

const (
	ObjectSource     SourceKind = "object"
	WindowSource     SourceKind = "window"
	DesktopSource    SourceKind = "desktop"
	ClipboardSource  SourceKind = "clipboard"
	maxCaptureHeader            = 4096
	captureTimeout              = 2 * time.Second
)

// CaptureRequest is an immutable job snapshot. Bounds are physical X11 root
// pixels, not GTK logical coordinates. The future target owner must convert
// provider coordinates before requesting capture and retain its object identity.
// Generation belongs to that owner, which must retire jobs on focus/lock/loss.
// Foreground is required even for clipboard capture. Window may name a child of
// Foreground (for example a Wine client); zero selects the foreground client.
type CaptureRequest struct {
	Kind       SourceKind
	Bounds     image.Rectangle
	Foreground uint64
	Window     uint64
	Generation uint64
}

func (request CaptureRequest) Validate() error {
	if request.Foreground == 0 || request.Foreground > 0xffffffff || request.Window > 0xffffffff || request.Generation == 0 {
		return ErrSource
	}
	switch request.Kind {
	case ObjectSource:
		if !valid_screen_rectangle(request.Bounds) {
			return ErrSource
		}
	case WindowSource, DesktopSource, ClipboardSource:
		if request.Bounds != (image.Rectangle{}) {
			return ErrSource
		}
	default:
		return ErrSource
	}
	if request.Kind != WindowSource && request.Window != 0 {
		return ErrSource
	}
	return nil
}

// CaptureHeader precedes a PNG on the private pipe. Rectangles describe physical
// root pixels. No monitor topology, object identity or activation authorization
// is inferred from them. Clipboard pixels have no screen mapping.
type CaptureHeader struct {
	Request      CaptureRequest
	SourceWindow uint64
	RootBounds   image.Rectangle
	SourceBounds image.Rectangle
	ScreenBounds image.Rectangle
	Status       string
}

type Capture struct {
	CaptureHeader
	Pixels    image.Image
	Transform Transform
}

// Capturer owns one helper at a time. It performs no startup I/O. Run Capture
// on a worker, never the input/event thread. The executable must be the reader
// binary supporting --ocr-capture; it is invoked directly without a shell.
type Capturer struct {
	gate    chan struct{}
	command func(context.Context, ...string) *exec.Cmd
}

func NewCapturer(executable string) *Capturer {
	return &Capturer{gate: make(chan struct{}, 1), command: func(ctx context.Context, args ...string) *exec.Cmd {
		return exec.CommandContext(ctx, executable, args...)
	}}
}

func (capturer *Capturer) Capture(ctx context.Context, request CaptureRequest) (Capture, error) {
	if err := request.Validate(); err != nil {
		return Capture{}, err
	}
	ctx, cancel := context.WithTimeout(ctx, captureTimeout)
	defer cancel()
	select {
	case capturer.gate <- struct{}{}:
		defer func() { <-capturer.gate }()
	case <-ctx.Done():
		return Capture{}, ctx.Err()
	}
	if err := ctx.Err(); err != nil {
		return Capture{}, err
	}
	input, err := json.Marshal(request)
	if err != nil {
		return Capture{}, ErrSource
	}
	output, _, err := run_process(ctx, capturer.command, input, maxImageBytes+maxCaptureHeader, "--ocr-capture")
	if err != nil {
		if errors.Is(err, ErrUnavailable) || errors.Is(err, ErrRecognition) {
			return Capture{}, ErrCapture
		}
		return Capture{}, err
	}
	capture, err := decode_capture(output, request)
	if ctx.Err() != nil {
		return Capture{}, ctx.Err()
	}
	return capture, err
}

// ReadCaptureRequest bounds and validates the helper's one request before any
// GTK/X11 initialization. Unknown fields and trailing messages are rejected.
func ReadCaptureRequest(input io.Reader) (CaptureRequest, error) {
	data, err := io.ReadAll(io.LimitReader(input, maxCaptureHeader+1))
	if err != nil || len(data) > maxCaptureHeader {
		return CaptureRequest{}, ErrLimit
	}
	var request CaptureRequest
	decoder := json.NewDecoder(bytes.NewReader(data))
	decoder.DisallowUnknownFields()
	if decoder.Decode(&request) != nil || decoder.Decode(new(any)) != io.EOF {
		return CaptureRequest{}, ErrSource
	}
	return request, request.Validate()
}

// WriteCaptureResponse retains a complete bounded frame before publishing it.
// Failures contain a fixed status only; native diagnostics never cross the pipe.
func WriteCaptureResponse(output io.Writer, header CaptureHeader, pixels image.Image) error {
	data, err := json.Marshal(header)
	if err != nil || len(data)+1 > maxCaptureHeader {
		return ErrLimit
	}
	buffer := &bounded_buffer{ctx: context.Background(), limit: maxImageBytes + maxCaptureHeader}
	_, _ = buffer.Write(append(data, '\n'))
	if header.Status == "" {
		if pixels == nil || !valid_dimensions(pixels.Bounds().Dx(), pixels.Bounds().Dy()) {
			return ErrInvalidImage
		}
		if err := png.Encode(buffer, pixels); err != nil {
			return err
		}
	}
	_, err = output.Write(buffer.Bytes())
	return err
}

func decode_capture(data []byte, request CaptureRequest) (Capture, error) {
	separator := bytes.IndexByte(data, '\n')
	if separator < 0 || separator >= maxCaptureHeader || len(data) > maxImageBytes+maxCaptureHeader {
		return Capture{}, ErrCapture
	}
	var header CaptureHeader
	decoder := json.NewDecoder(bytes.NewReader(data[:separator]))
	decoder.DisallowUnknownFields()
	if decoder.Decode(&header) != nil || decoder.Decode(new(any)) != io.EOF || header.Request != request {
		return Capture{}, ErrCapture
	}
	if header.Status != "" {
		if len(data) != separator+1 {
			return Capture{}, ErrCapture
		}
		switch header.Status {
		case "unavailable":
			return Capture{}, ErrCapture
		case "source":
			return Capture{}, ErrSource
		case "changed":
			return Capture{}, ErrSourceChanged
		case "limit":
			return Capture{}, ErrLimit
		default:
			return Capture{}, ErrCapture
		}
	}
	payload := data[separator+1:]
	config, err := png.DecodeConfig(bytes.NewReader(payload))
	if err != nil {
		return Capture{}, ErrCapture
	}
	if !valid_dimensions(config.Width, config.Height) {
		return Capture{}, ErrLimit
	}
	transform := Transform{ImageContent: image.Rect(0, 0, config.Width, config.Height)}
	if request.Kind == ClipboardSource {
		if header.SourceWindow != 0 || header.RootBounds != (image.Rectangle{}) || header.SourceBounds != (image.Rectangle{}) || header.ScreenBounds != (image.Rectangle{}) {
			return Capture{}, ErrCapture
		}
	} else {
		if !valid_screen_rectangle(header.RootBounds) || !valid_screen_rectangle(header.SourceBounds) ||
			!valid_screen_rectangle(header.ScreenBounds) || !header.ScreenBounds.In(header.RootBounds) || !header.ScreenBounds.In(header.SourceBounds) ||
			header.ScreenBounds.Dx() != config.Width || header.ScreenBounds.Dy() != config.Height {
			return Capture{}, ErrCapture
		}
		sourceWindow := request.Foreground
		if request.Kind == WindowSource && request.Window != 0 {
			sourceWindow = request.Window
		}
		if request.Kind == DesktopSource {
			sourceWindow = 0
			if header.SourceBounds != header.RootBounds || header.ScreenBounds != header.RootBounds {
				return Capture{}, ErrCapture
			}
		} else if request.Kind == ObjectSource && header.ScreenBounds != request.Bounds.Intersect(header.SourceBounds).Intersect(header.RootBounds) {
			return Capture{}, ErrCapture
		}
		if request.Kind == WindowSource && header.ScreenBounds != header.SourceBounds.Intersect(header.RootBounds) {
			return Capture{}, ErrCapture
		}
		if header.SourceWindow != sourceWindow {
			return Capture{}, ErrCapture
		}
		transform.HasScreenMapping, transform.ScreenBounds = true, header.ScreenBounds
	}
	payloadReader := bytes.NewReader(payload)
	pixels, err := png.Decode(payloadReader)
	if err != nil || payloadReader.Len() != 0 {
		return Capture{}, ErrCapture
	}
	return Capture{CaptureHeader: header, Pixels: pixels, Transform: transform}, nil
}