package runs

import (
	"context"
	"errors"
	"io"
	"sync"
	"time"

	"github.com/carlosframework/platform/internal/release"
)

// errBadStream is what Close reports when a writer was made for a stream
// that is neither StreamOut nor StreamErr. ChunkKey trusts its stream, so the
// chunker refuses anything else before it reaches a key.
var errBadStream = errors.New("runs: chunker stream is not out or err")

// Chunker turns a run's two streams into chunks (run spec §5.5):
// out/<seq>-<o|e>, one seq for both streams in the order the bytes were
// read. A chunk holds one stream's bytes: a write on the other stream cuts
// it, so the follower interleaves exactly as the agent read. A chunk is
// written every second, or as soon as it holds 256 KB. Past 64 MB in total
// nothing more is written and Close says truncated; writes never fail, so
// the command is never failed by a refused pipe. Each chunk is a
// create-only Put retried until it lands or the context ends, and the
// retries run under the lock. A stalled bucket therefore does block the
// command's pipes, and the command with them once its pipe buffers fill,
// until the Put lands or the context ends. That is deliberate: it keeps
// seqs in order and stops output piling up in memory.
type Chunker struct {
	ctx           context.Context
	s             release.Store
	acct, app, id string
	every         time.Duration
	max           int64

	mu        sync.Mutex
	stream    string
	pending   []byte
	seq       int64
	total     int64
	truncated bool
	err       error
	closed    bool

	stop      chan struct{}
	done      chan struct{}
	closeOnce sync.Once
}

// NewChunker starts a chunker for run id with the spec's cadence and cap.
// An account, app or run id the store helpers would refuse (ErrBadApp,
// ErrBadID) gives a chunker that writes nothing and whose Close returns that
// error, so it can never write outside runs/<id>/out/.
func NewChunker(ctx context.Context, s release.Store, acct, app, id string) *Chunker {
	return newChunker(ctx, s, acct, app, id, ChunkEvery, MaxOutputBytes)
}

func newChunker(ctx context.Context, s release.Store, acct, app, id string, every time.Duration, max int64) *Chunker {
	c := &Chunker{ctx: ctx, s: s, acct: acct, app: app, id: id, every: every, max: max,
		stop: make(chan struct{}), done: make(chan struct{})}
	if err := checkRun(acct, app, id); err != nil {
		c.err, c.truncated = err, true
		close(c.done)
		return c
	}
	go c.tick()
	return c
}

func (c *Chunker) tick() {
	defer close(c.done)
	t := time.NewTicker(c.every)
	defer t.Stop()
	for {
		select {
		case <-c.stop:
			return
		case <-t.C:
			c.mu.Lock()
			c.flushLocked()
			c.mu.Unlock()
		}
	}
}

type chunkStream struct {
	c      *Chunker
	stream string
}

func (w chunkStream) Write(p []byte) (int, error) {
	w.c.write(w.stream, p)
	return len(p), nil
}

// Stream is the writer for one stream, StreamOut or StreamErr. Its Write
// never fails and never returns short.
func (c *Chunker) Stream(stream string) io.Writer { return chunkStream{c: c, stream: stream} }

func (c *Chunker) write(stream string, p []byte) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.closed {
		return
	}
	if stream != StreamOut && stream != StreamErr {
		if c.err == nil {
			c.err, c.truncated, c.pending = errBadStream, true, nil
		}
		return
	}
	if c.err != nil {
		c.truncated = c.truncated || len(p) > 0
		return
	}
	if room := c.max - c.total; int64(len(p)) > room {
		p = p[:room]
		c.truncated = true
	}
	for len(p) > 0 {
		if c.stream != stream && len(c.pending) > 0 {
			c.flushLocked()
			if c.err != nil {
				return
			}
		}
		c.stream = stream
		n := min(len(p), ChunkBytes-len(c.pending))
		c.pending = append(c.pending, p[:n]...)
		c.total += int64(n)
		p = p[n:]
		if len(c.pending) >= ChunkBytes {
			c.flushLocked()
			if c.err != nil {
				c.truncated = true
				return
			}
		}
	}
}

// flushLocked writes the pending chunk under the next seq. A Put that
// answers ErrCASConflict was written by an earlier attempt whose answer was
// lost; the seq is spent either way.
func (c *Chunker) flushLocked() {
	if len(c.pending) == 0 || c.err != nil {
		return
	}
	seq := c.seq + 1
	key := ChunkKey(c.acct, c.app, c.id, seq, c.stream)
	backoff := 100 * time.Millisecond
	for {
		_, err := c.s.Put(c.ctx, key, c.pending, "")
		if err == nil || errors.Is(err, release.ErrCASConflict) {
			break
		}
		select {
		case <-c.ctx.Done():
			c.err, c.truncated, c.pending = err, true, nil
			return
		case <-time.After(backoff):
		}
		backoff = min(backoff*2, 5*time.Second)
	}
	c.seq, c.pending = seq, nil
}

// Close stops the ticker and writes the last chunk. Call it after both
// streams reached their end (CmdHandle.Stop drains them); writes after Close
// are dropped. It does not return until every chunk is written or chunking
// gave up, so a caller may write result.json once it returns. It returns the
// last seq written (0: no output), whether output was cut off, and the
// failure that ended chunking, if one did. It is safe to call twice.
func (c *Chunker) Close() (int64, bool, error) {
	c.closeOnce.Do(func() {
		close(c.stop)
		<-c.done
	})
	c.mu.Lock()
	defer c.mu.Unlock()
	c.flushLocked()
	c.closed = true
	return c.seq, c.truncated, c.err
}