//go:build linux

package wayland

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strings"
	"time"
)

// Niri IPC screenshot stream implementation.
//
// Protocol reference (niri 26.04):
// - niri wiki IPC.md (Event Stream, Programmatic Access):
//   Connect to UNIX domain socket at $NIRI_SOCKET.
//   Send single-line JSON request "EventStream"\n.
//   Read single-line JSON reply {"Ok":"Handled"} (or {"Err":"..."}).
//   Read subsequent stream of newline-delimited JSON events.
// - niri-ipc/src/lib.rs:
//   Request::EventStream
//   Response::Handled
//   Event::ScreenshotCaptured { path: Option<String> }
// - Event ordering in niri (src/niri.rs save_screenshot()):
//   niri first dispatches the clipboard selection via calloop channel tx.send(buf),
//   writes the screenshot to disk if configured, and then dispatches the completion
//   event via event_tx.send(path_string) to fire state.ipc_screenshot_taken(path).
//   Because the Wayland connection and the Niri IPC socket are separate Unix streams,
//   the ScreenshotCaptured IPC event may arrive at chirami slightly before or after
//   the Wayland data-control selection offer arrives. When path is null or unreadable,
//   we briefly retry reading image/png from the clipboard.

func (b *backend) watchNiri(ctx context.Context, socketPath string) {
	backoff := 100 * time.Millisecond
	const maxBackoff = 5 * time.Second

	for {
		select {
		case <-ctx.Done():
			return
		default:
		}

		err := b.runNiriStream(ctx, socketPath)
		if ctx.Err() != nil {
			return
		}
		if debug && err != nil {
			debugf("wayland: niri ipc stream exited: %v (reconnecting in %v)", err, backoff)
		}

		select {
		case <-ctx.Done():
			return
		case <-time.After(backoff):
		}

		backoff *= 2
		if backoff > maxBackoff {
			backoff = maxBackoff
		}
	}
}

func (b *backend) runNiriStream(ctx context.Context, socketPath string) error {
	var dialer net.Dialer
	conn, err := dialer.DialContext(ctx, "unix", socketPath)
	if err != nil {
		return err
	}
	defer conn.Close()

	done := make(chan struct{})
	defer close(done)
	go func() {
		select {
		case <-ctx.Done():
			_ = conn.Close()
		case <-done:
		}
	}()

	// Write JSON request "EventStream"\n
	if _, err := conn.Write([]byte("\"EventStream\"\n")); err != nil {
		return fmt.Errorf("write EventStream request: %w", err)
	}

	r := bufio.NewReader(conn)
	replyLine, err := readNiriLine(r)
	if err != nil {
		return fmt.Errorf("read EventStream reply: %w", err)
	}

	var reply struct {
		Ok  string `json:"Ok"`
		Err string `json:"Err"`
	}
	if err := json.Unmarshal(replyLine, &reply); err != nil {
		return fmt.Errorf("unmarshal EventStream reply %q: %w", string(replyLine), err)
	}
	if reply.Err != "" {
		return fmt.Errorf("niri EventStream error: %s", reply.Err)
	}
	if reply.Ok != "Handled" {
		return fmt.Errorf("unexpected niri EventStream reply: %s", string(replyLine))
	}

	return b.readNiriEvents(ctx, r)
}

func (b *backend) readNiriEvents(ctx context.Context, r *bufio.Reader) error {
	for {
		line, err := readNiriLine(r)
		if err != nil {
			return err
		}

		// Fast filter to ignore other high-volume events (workspaces, window focus, etc.)
		if !bytes.Contains(line, []byte(`"ScreenshotCaptured"`)) {
			continue
		}

		var ev struct {
			ScreenshotCaptured *struct {
				Path *string `json:"path"`
			} `json:"ScreenshotCaptured"`
		}
		if err := json.Unmarshal(line, &ev); err != nil || ev.ScreenshotCaptured == nil {
			continue
		}

		b.handleScreenshotCaptured(ctx, ev.ScreenshotCaptured.Path)
	}
}

func (b *backend) handleScreenshotCaptured(ctx context.Context, path *string) {
	if path != nil && *path != "" {
		p := expandHome(*path)
		data, err := os.ReadFile(p)
		if err == nil && len(data) > 0 && isPNG(data) {
			b.deliverScreenshot(data)
			return
		}
	}

	// When path is not provided, null, or unreadable, read from clipboard.
	// We retry briefly because niri sends clipboard data and IPC events concurrently.
	img, err := b.ReadClipboardImageWithRetry(ctx, 1*time.Second)
	if err == nil && len(img) > 0 {
		b.deliverScreenshot(img)
	} else if debug {
		debugf("wayland: failed to read screenshot from clipboard: %v", err)
	}
}

// ReadClipboardImageWithRetry reads the clipboard's image/png once niri's
// new selection has arrived. niri sets the clipboard just before it sends
// ScreenshotCaptured, but the two travel on different sockets, so the
// selection may land a little before or after the event; an older one is an
// earlier copy and is never read, so a stale image is never OCR'd.
func (b *backend) ReadClipboardImageWithRetry(ctx context.Context, timeout time.Duration) ([]byte, error) {
	since := time.Now().Add(-500 * time.Millisecond)
	deadline := time.Now().Add(timeout)
	for {
		if time.Unix(0, b.lastClipOfferNanos.Load()).After(since) {
			img, err := b.ReadClipboardImage(ctx)
			if err == nil && isPNG(img) {
				return img, nil
			}
		}
		rem := time.Until(deadline)
		if rem <= 0 {
			return nil, errors.New("no new image/png on the clipboard")
		}
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-b.ctx.Done():
			return nil, b.ctx.Err()
		case <-b.clipboardNotify:
		case <-time.After(min(rem, 25*time.Millisecond)):
		}
	}
}

func (b *backend) deliverScreenshot(data []byte) {
	select {
	case b.shots <- data:
	default:
		if debug {
			debugf("wayland: dropped screenshot: shots channel buffer full")
		}
	}
}

var pngHeader = []byte{0x89, 'P', 'N', 'G', 0x0d, 0x0a, 0x1a, 0x0a}

func isPNG(data []byte) bool {
	return bytes.HasPrefix(data, pngHeader)
}

func expandHome(p string) string {
	if strings.HasPrefix(p, "~/") {
		if home := os.Getenv("HOME"); home != "" {
			return filepath.Join(home, p[2:])
		}
	}
	return p
}

func readNiriLine(r *bufio.Reader) ([]byte, error) {
	var buf bytes.Buffer
	for {
		chunk, isPrefix, err := r.ReadLine()
		if err != nil {
			return nil, err
		}
		buf.Write(chunk)
		if !isPrefix {
			break
		}
		if buf.Len() > 16*1024*1024 {
			return nil, errors.New("niri event line exceeds 16MB limit")
		}
	}
	return buf.Bytes(), nil
}