// SPDX-License-Identifier: MIT

// Package execstream is the framed protobuf stream between the host process of
// ermite-runner and the executor running inside a sandbox.
//
// Wire format: each frame is a 4-byte big-endian length followed by that many
// bytes of a serialized execpb.Frame (see exec.proto). A frame is at most
// MaxFrame bytes; longer output is split by the sender. The stream is
// transport-agnostic: vsock in a micro-VM, a Unix socket in a container.
package execstream

//go:generate go tool buf generate ../.. --template ../../buf.gen.yaml

import (
	"bufio"
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"maps"
	"slices"
	"sync"
	"time"

	"google.golang.org/protobuf/proto"

	"bnema.dev/ermite-runner/internal/execstream/execpb"
)

// MaxFrame is the largest accepted serialized frame.
// It leaves room for the largest job result: outputs of 1 MiB in total (the
// limit of ermite) plus the overhead of the encoding.
const MaxFrame = 2 << 20

// ChunkSize is the largest output payload placed in one Output frame.
const ChunkSize = 32 * 1024

// ErrFrameTooLarge is returned for a frame above MaxFrame, on either side.
var ErrFrameTooLarge = errors.New("execstream: frame too large")

// Conn reads and writes frames on a byte stream. Send is safe for concurrent
// use; Recv must be called from one goroutine.
type Conn struct {
	rw  io.ReadWriter
	br  *bufio.Reader
	wmu sync.Mutex
}

// NewConn wraps a byte stream.
func NewConn(rw io.ReadWriter) *Conn {
	return &Conn{rw: rw, br: bufio.NewReaderSize(rw, 64*1024)}
}

// Send writes one frame.
func (c *Conn) Send(f *execpb.Frame) error {
	b, err := proto.Marshal(f)
	if err != nil {
		return fmt.Errorf("execstream: marshal: %w", err)
	}
	if len(b) > MaxFrame {
		return ErrFrameTooLarge
	}
	buf := make([]byte, 4+len(b))
	binary.BigEndian.PutUint32(buf, uint32(len(b)))
	copy(buf[4:], b)
	c.wmu.Lock()
	defer c.wmu.Unlock()
	_, err = c.rw.Write(buf)
	return err
}

// Recv reads one frame. It returns io.EOF only on a clean end of stream between
// frames; a stream cut inside a frame yields io.ErrUnexpectedEOF.
func (c *Conn) Recv() (*execpb.Frame, error) {
	var hdr [4]byte
	if _, err := io.ReadFull(c.br, hdr[:]); err != nil {
		return nil, err
	}
	n := binary.BigEndian.Uint32(hdr[:])
	if n > MaxFrame {
		return nil, ErrFrameTooLarge
	}
	b := make([]byte, n)
	if _, err := io.ReadFull(c.br, b); err != nil {
		if errors.Is(err, io.EOF) {
			err = io.ErrUnexpectedEOF
		}
		return nil, err
	}
	f := &execpb.Frame{}
	if err := proto.Unmarshal(b, f); err != nil {
		return nil, fmt.Errorf("execstream: unmarshal: %w", err)
	}
	return f, nil
}

// Drive is the host side of a session: it waits for the executor's Hello, sends
// run, passes every Output frame to onOutput, and returns the Exit frame.
//
// cancel, when non-nil, asks the executor to stop the script: the first value
// received is sent as one Cancel frame (the executor escalates SIGINT, SIGTERM,
// SIGKILL by itself); closing the channel without a value sends nothing. It is
// deliberately not a context: cancelling the session's context would also abort
// the reads that deliver the executor's Exit. Drive keeps waiting for Exit after
// a Cancel; the caller bounds that wait by stopping the sandbox, which closes
// the stream and makes Drive return an error.
func Drive(rw io.ReadWriter, run *execpb.Run, cancel <-chan *execpb.Cancel, onOutput func(*execpb.Output)) (*execpb.Hello, *execpb.Exit, error) {
	return DriveJob(rw, run, cancel, JobHandler{Output: onOutput})
}

// resolveJoinGrace is how long DriveJob waits for the resolution worker to
// finish once the session is over before it interrupts a write that is stuck
// (a guest that stopped reading): the write deadline is set to now, and the
// stream is closed if it cannot be interrupted otherwise. The session is over at
// that point, the caller closes the stream anyway.
var resolveJoinGrace = 2 * time.Second

// joinResolver waits for the worker. A worker in a Send that cannot finish
// because the peer does not read would block DriveJob for ever, so the write is
// interrupted after the grace period.
func joinResolver(rw io.ReadWriter, workers *sync.WaitGroup) {
	finished := make(chan struct{})
	go func() { workers.Wait(); close(finished) }()
	select {
	case <-finished:
		return
	case <-time.After(resolveJoinGrace):
	}
	if d, ok := rw.(interface{ SetWriteDeadline(time.Time) error }); ok {
		_ = d.SetWriteDeadline(time.Now())
	}
	select {
	case <-finished:
		return
	case <-time.After(resolveJoinGrace):
	}
	if cl, ok := rw.(io.Closer); ok {
		_ = cl.Close()
		<-finished
	}
}

// resolveQueue is how many ResolveRequest frames may wait for the worker.
const resolveQueue = 2

func resolveFrame(r *execpb.ResolveResponse) *execpb.Frame {
	return &execpb.Frame{Body: &execpb.Frame_ResolveResponse{ResolveResponse: r}}
}

// JobHandler receives the frames of a session. Nil functions ignore their
// frames. The functions are called from the goroutine of DriveJob, in the
// order of the stream.
type JobHandler struct {
	Hello         func(*execpb.Hello)
	Output        func(*execpb.Output)
	StepStarted   func(*execpb.StepStarted)
	StepCompleted func(*execpb.StepCompleted)
	Log           func(*execpb.Log)
	JobResult     func(*execpb.JobResult)
	// Resolve answers the executor's ResolveRequest frames, one at a time and in
	// order, from a single worker (the call goes to the server); the answer is
	// sent as a ResolveResponse frame carrying the id of the request. ctx ends
	// with the session: a call still running is cancelled and DriveJob waits for
	// it before it returns. At most resolveQueue requests wait; the next ones are
	// answered with an error at once, without a call.
	Resolve func(ctx context.Context, r *execpb.ResolveRequest) *execpb.ResolveResponse
	// StepAnnotations and StepSummary come before the StepCompleted of their
	// step.
	StepAnnotations func(*execpb.StepAnnotations)
	StepSummary     func(*execpb.StepSummary)
	// Secrets and Vars (Job.secrets, Job.vars) are sent before the Run frame,
	// one NamedValue frame each, so that their total size is not bound by the
	// size of one frame (MaxFrame). A single value must fit in a frame.
	Secrets, Vars map[string]string
}

// DriveJob is Drive for a session that carries a job: besides Output frames it
// delivers the step, log and result frames to h.
func DriveJob(rw io.ReadWriter, run *execpb.Run, cancel <-chan *execpb.Cancel, h JobHandler) (*execpb.Hello, *execpb.Exit, error) {
	c := NewConn(rw)
	first, err := c.Recv()
	if err != nil {
		return nil, nil, fmt.Errorf("execstream: waiting for hello: %w", err)
	}
	hello := first.GetHello()
	if hello == nil {
		return nil, nil, fmt.Errorf("execstream: first frame is not hello")
	}
	if h.Hello != nil {
		h.Hello(hello)
	}
	if err := sendValues(c, h); err != nil {
		return hello, nil, err
	}
	if err := c.Send(&execpb.Frame{Body: &execpb.Frame_Run{Run: run}}); err != nil {
		return hello, nil, fmt.Errorf("execstream: sending run: %w", err)
	}

	done := make(chan struct{})
	defer close(done)

	// One worker answers the resolutions; it ends, and is waited for, with the
	// session.
	resolveCh := make(chan *execpb.ResolveRequest, resolveQueue)
	rctx, rcancel := context.WithCancel(context.Background())
	var workers sync.WaitGroup
	workers.Add(1)
	go func() {
		defer workers.Done()
		for {
			select {
			case req := <-resolveCh:
				resp := h.Resolve(rctx, req)
				resp.Id = req.Id
				if rctx.Err() != nil {
					return // the session is over: nobody reads the answer
				}
				_ = c.Send(resolveFrame(resp))
			case <-rctx.Done():
				return
			}
		}
	}()
	defer func() { rcancel(); joinResolver(rw, &workers) }()
	if cancel != nil {
		go func() {
			select {
			case cn, ok := <-cancel:
				if ok {
					if cn == nil {
						cn = &execpb.Cancel{}
					}
					_ = c.Send(&execpb.Frame{Body: &execpb.Frame_Cancel{Cancel: cn}})
				}
			case <-done:
			}
		}()
	}

	for {
		f, err := c.Recv()
		if err != nil {
			return hello, nil, fmt.Errorf("execstream: reading: %w", err)
		}
		switch b := f.Body.(type) {
		case *execpb.Frame_Output:
			if h.Output != nil {
				h.Output(b.Output)
			}
		case *execpb.Frame_StepStarted:
			if h.StepStarted != nil {
				h.StepStarted(b.StepStarted)
			}
		case *execpb.Frame_StepCompleted:
			if h.StepCompleted != nil {
				h.StepCompleted(b.StepCompleted)
			}
		case *execpb.Frame_Log:
			if h.Log != nil {
				h.Log(b.Log)
			}
		case *execpb.Frame_JobResult:
			if h.JobResult != nil {
				h.JobResult(b.JobResult)
			}
		case *execpb.Frame_ResolveRequest:
			req := b.ResolveRequest
			if h.Resolve == nil {
				_ = c.Send(resolveFrame(&execpb.ResolveResponse{Id: req.Id, Error: "the host cannot resolve actions"}))
				break
			}
			select {
			case resolveCh <- req:
			default:
				_ = c.Send(resolveFrame(&execpb.ResolveResponse{Id: req.Id, Error: "too many action resolutions in flight"}))
			}
		case *execpb.Frame_StepAnnotations:
			if h.StepAnnotations != nil {
				h.StepAnnotations(b.StepAnnotations)
			}
		case *execpb.Frame_StepSummary:
			if h.StepSummary != nil {
				h.StepSummary(b.StepSummary)
			}
		case *execpb.Frame_Exit:
			return hello, b.Exit, nil
		}
	}
}

// sendValues sends the secrets, then the variables of the job, in name order,
// one frame each. The error never holds a value.
func sendValues(c *Conn, h JobHandler) error {
	for _, k := range slices.Sorted(maps.Keys(h.Secrets)) {
		f := &execpb.Frame{Body: &execpb.Frame_Secret{Secret: &execpb.NamedValue{Name: k, Value: h.Secrets[k]}}}
		if err := c.Send(f); err != nil {
			return fmt.Errorf("execstream: sending secret %q: %w", k, err)
		}
	}
	for _, k := range slices.Sorted(maps.Keys(h.Vars)) {
		f := &execpb.Frame{Body: &execpb.Frame_Var{Var: &execpb.NamedValue{Name: k, Value: h.Vars[k]}}}
		if err := c.Send(f); err != nil {
			return fmt.Errorf("execstream: sending variable %q: %w", k, err)
		}
	}
	return nil
}