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

package i3ipc

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net"
	"time"
)

const queryTimeout = 2 * time.Second

type connection struct {
	net.Conn
	pending      []message
	pendingBytes int
}

// request preserves events interleaved with a reply on the subscribed socket.
// One worker owns both reads and comparisons; overflow forces a fresh baseline.
func (c *connection) request(kind uint32, body []byte, result any) error {
	if err := c.SetDeadline(time.Now().Add(queryTimeout)); err != nil {
		return err
	}
	defer c.SetDeadline(time.Time{})
	if err := writeMessage(c, kind, body); err != nil {
		return err
	}
	for {
		m, err := readMessage(c)
		if err != nil {
			return err
		}
		if m.kind == kind {
			return json.Unmarshal(m.body, result)
		}
		if m.kind&eventBit == 0 {
			return errors.New("unexpected i3 reply")
		}
		if len(c.pending) >= maxPending || c.pendingBytes+len(m.body) > maxPayload {
			return errors.New("i3 event backlog exceeded")
		}
		c.pending = append(c.pending, m)
		c.pendingBytes += len(m.body)
	}
}

func (c *connection) queryTree() (node, error) {
	var root node
	err := c.request(tree, nil, &root)
	if err == nil && (root.Type != "root" || root.ID == 0) {
		err = errors.New("invalid i3 tree")
	}
	return root, err
}

func session(ctx context.Context, path string, emit func(Event)) error {
	socket, err := (&net.Dialer{Timeout: queryTimeout}).DialContext(ctx, "unix", path)
	if err != nil {
		return err
	}
	c := &connection{Conn: socket}
	defer c.Close()
	// Closing the descriptor interrupts both idle reads and queries immediately.
	stop := context.AfterFunc(ctx, func() { c.Close() })
	defer stop()
	var subscribed struct {
		Success bool `json:"success"`
	}
	if err := c.request(subscribe, []byte(`["workspace","mode","window"]`), &subscribed); err != nil {
		return err
	}
	if !subscribed.Success {
		return errors.New("i3 subscription rejected")
	}
	var spaces []workspace
	if err := c.request(workspaces, nil, &spaces); err != nil {
		return err
	}
	root, err := c.queryTree()
	if err != nil {
		return err
	}
	var state state
	state.baseline(root, spaces)
	// Everything received through the baseline reply belongs to the snapshot.
	// This also suppresses reconnect snapshots and races during initialization.
	c.pending = nil
	c.pendingBytes = 0
	for ctx.Err() == nil {
		var m message
		if len(c.pending) != 0 {
			m = c.pending[0]
			c.pendingBytes -= len(m.body)
			c.pending[0] = message{}
			c.pending = c.pending[1:]
		} else {
			// Idle sessions have no timeout; a partial frame must finish promptly.
			// Read one byte first so an idle desktop doesn't reconnect periodically.
			var first [1]byte
			if _, err := c.Read(first[:]); err != nil {
				return err
			}
			if err := c.SetReadDeadline(time.Now().Add(queryTimeout)); err != nil {
				return err
			}
			m, err = readMessage(io.MultiReader(bytes.NewReader(first[:]), c))
			c.SetReadDeadline(time.Time{})
			if err != nil {
				return err
			}
		}
		events, err := state.apply(m, c.queryTree)
		if err != nil {
			return err
		}
		for _, event := range events {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			emit(event)
		}
	}
	return ctx.Err()
}

// Run rediscovers the active session after loss. All connection/cache state is
// local to one session and is discarded on errors. emit must be nonblocking.
// discover must verify the current window manager, not just an environment hint.
func Run(ctx context.Context, discover func() string, emit func(Event)) {
	backoff := 250 * time.Millisecond
	for ctx.Err() == nil {
		started := time.Now()
		if path := discover(); path != "" {
			_ = session(ctx, path, emit)
		}
		if time.Since(started) > 10*time.Second {
			backoff = 250 * time.Millisecond
		}
		timer := time.NewTimer(backoff)
		select {
		case <-ctx.Done():
			timer.Stop()
			return
		case <-timer.C:
		}
		backoff = min(5*time.Second, backoff*2)
	}
}