package runs

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

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

const tid = "run_abcdefghijklmnopqrstuvwxyz"

func chunksOf(t *testing.T, s release.Store) []string {
	t.Helper()
	keys, err := ListChunks(context.Background(), s, "bab", "shop", tid, 0, 0)
	if err != nil {
		t.Fatal(err)
	}
	var out []string
	for _, k := range keys {
		raw, _, err := s.Get(context.Background(), k)
		if err != nil {
			t.Fatal(err)
		}
		_, stream, _ := ParseChunkKey(k, OutPrefix("bab", "shop", tid))
		out = append(out, stream+":"+string(raw))
	}
	return out
}

func TestChunkerKeepsReadOrderAcrossStreams(t *testing.T) {
	s := release.NewDirStore(t.TempDir())
	c := newChunker(context.Background(), s, "bab", "shop", tid, time.Hour, MaxOutputBytes)
	io.WriteString(c.Stream(StreamOut), "a")
	io.WriteString(c.Stream(StreamErr), "b")
	io.WriteString(c.Stream(StreamOut), "c")
	io.WriteString(c.Stream(StreamOut), "d")
	last, truncated, err := c.Close()
	if err != nil || truncated || last != 3 {
		t.Fatalf("Close = %d %v %v", last, truncated, err)
	}
	if got := strings.Join(chunksOf(t, s), " "); got != "o:a e:b o:cd" {
		t.Fatalf("chunks %q", got)
	}
}

func TestChunkerFlushesOnTheTickAndAtTheSizeLimit(t *testing.T) {
	s := release.NewDirStore(t.TempDir())
	c := newChunker(context.Background(), s, "bab", "shop", tid, 20*time.Millisecond, MaxOutputBytes)
	io.WriteString(c.Stream(StreamOut), "tick")
	deadline := time.Now().Add(2 * time.Second)
	for len(chunksOf(t, s)) == 0 && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
	if got := chunksOf(t, s); len(got) != 1 || got[0] != "o:tick" {
		t.Fatalf("after a tick: %q", got)
	}
	c2 := newChunker(context.Background(), s, "bab", "shop", "run_bbbbbbbbbbbbbbbbbbbbbbbbbb", time.Hour, MaxOutputBytes)
	io.WriteString(c2.Stream(StreamOut), strings.Repeat("x", ChunkBytes+1))
	keys, _ := ListChunks(context.Background(), s, "bab", "shop", "run_bbbbbbbbbbbbbbbbbbbbbbbbbb", 0, 0)
	if len(keys) != 1 {
		t.Fatalf("a full chunk must flush at once: %q", keys)
	}
	if last, _, _ := c2.Close(); last != 2 {
		t.Fatalf("last %d, want 2 (one full chunk, one byte)", last)
	}
	c.Close()
}

func TestChunkerStopsAtTheCapAndSaysSo(t *testing.T) {
	s := release.NewDirStore(t.TempDir())
	c := newChunker(context.Background(), s, "bab", "shop", tid, time.Hour, 10)
	n, err := io.WriteString(c.Stream(StreamOut), "0123456789abcdef")
	if n != 16 || err != nil {
		t.Fatalf("a write past the cap must not fail the stream: %d %v", n, err)
	}
	io.WriteString(c.Stream(StreamErr), "more")
	_, truncated, err := c.Close()
	if !truncated || err != nil {
		t.Fatalf("truncated %v err %v", truncated, err)
	}
	if got := strings.Join(chunksOf(t, s), ""); got != "o:0123456789" {
		t.Fatalf("chunks %q", got)
	}
}

// failingStore fails the first n Puts.
type failingStore struct {
	release.Store
	mu   sync.Mutex
	fail int
}

func (f *failingStore) Put(ctx context.Context, key string, body []byte, ifMatch string) (string, error) {
	f.mu.Lock()
	if f.fail > 0 {
		f.fail--
		f.mu.Unlock()
		return "", errors.New("slow down")
	}
	f.mu.Unlock()
	return f.Store.Put(ctx, key, body, ifMatch)
}

func TestChunkerRetriesAFailedPut(t *testing.T) {
	s := &failingStore{Store: release.NewDirStore(t.TempDir()), fail: 2}
	c := newChunker(context.Background(), s, "bab", "shop", tid, time.Hour, MaxOutputBytes)
	io.WriteString(c.Stream(StreamOut), "once")
	last, _, err := c.Close()
	if err != nil || last != 1 {
		t.Fatalf("Close = %d %v", last, err)
	}
	if got := chunksOf(t, s.Store); len(got) != 1 || got[0] != "o:once" {
		t.Fatalf("chunks %q", got)
	}
}

func TestChunkerGivesUpWhenItsContextEnds(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	s := &failingStore{Store: release.NewDirStore(t.TempDir()), fail: 1 << 30}
	c := newChunker(ctx, s, "bab", "shop", tid, time.Hour, MaxOutputBytes)
	go func() { time.Sleep(50 * time.Millisecond); cancel() }()
	io.WriteString(c.Stream(StreamOut), strings.Repeat("x", ChunkBytes))
	_, truncated, err := c.Close()
	if err == nil || !truncated {
		t.Fatalf("a chunk that never lands: truncated %v err %v", truncated, err)
	}
}

// dyingStore lets the first ok Puts through and fails every one after.
type dyingStore struct {
	release.Store
	mu sync.Mutex
	ok int
}

func (d *dyingStore) Put(ctx context.Context, key string, body []byte, ifMatch string) (string, error) {
	d.mu.Lock()
	allowed := d.ok > 0
	if allowed {
		d.ok--
	}
	d.mu.Unlock()
	if !allowed {
		return "", errors.New("gone")
	}
	return d.Store.Put(ctx, key, body, ifMatch)
}

func TestChunkerReportsTheLastSeqThatLanded(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	s := &dyingStore{Store: release.NewDirStore(t.TempDir()), ok: 1}
	c := newChunker(ctx, s, "bab", "shop", tid, time.Hour, MaxOutputBytes)
	go func() { time.Sleep(50 * time.Millisecond); cancel() }()
	io.WriteString(c.Stream(StreamOut), strings.Repeat("x", ChunkBytes)) // chunk 1 lands
	io.WriteString(c.Stream(StreamOut), strings.Repeat("y", ChunkBytes)) // chunk 2 gives up
	last, truncated, err := c.Close()
	if last != 1 || !truncated || err == nil {
		t.Fatalf("Close = %d %v %v, want 1 true and an error", last, truncated, err)
	}
	if got := chunksOf(t, s.Store); len(got) != 1 {
		t.Fatalf("chunks %d, want 1", len(got))
	}
}

// A give-up in the middle of one large write must end the write, not spin
// on it holding the lock.
func TestChunkerGivingUpMidWriteDoesNotHang(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	s := &failingStore{Store: release.NewDirStore(t.TempDir()), fail: 1 << 30}
	c := newChunker(ctx, s, "bab", "shop", tid, time.Hour, MaxOutputBytes)
	done := make(chan struct{})
	var last int64
	var truncated bool
	var err error
	go func() {
		defer close(done)
		n, werr := c.Stream(StreamOut).Write(make([]byte, 2*ChunkBytes+1))
		if n != 2*ChunkBytes+1 || werr != nil {
			t.Errorf("Write = %d %v", n, werr)
		}
		last, truncated, err = c.Close()
	}()
	select {
	case <-done:
	case <-time.After(10 * time.Second):
		t.Fatal("Write or Close hung after the chunker gave up")
	}
	if last != 0 || !truncated || err == nil {
		t.Fatalf("Close = %d %v %v", last, truncated, err)
	}
}

func TestChunkerRejoinsEachStreamWhenBothWriteAtOnce(t *testing.T) {
	s := release.NewDirStore(t.TempDir())
	c := newChunker(context.Background(), s, "bab", "shop", tid, 2*time.Millisecond, MaxOutputBytes)
	want := map[string]string{}
	var wg sync.WaitGroup
	for _, stream := range []string{StreamOut, StreamErr} {
		var b strings.Builder
		for i := 0; i < 2000; i++ {
			b.WriteString(stream)
			b.WriteString(strings.Repeat(string(rune('a'+i%26)), i%37))
			b.WriteString("\n")
		}
		data := b.String()
		want[stream] = data
		wg.Add(1)
		go func() {
			defer wg.Done()
			w := c.Stream(stream)
			for len(data) > 0 {
				n := min(len(data), 1+len(data)%50)
				w.Write([]byte(data[:n]))
				data = data[n:]
				time.Sleep(time.Microsecond)
			}
		}()
	}
	wg.Wait()
	if _, truncated, err := c.Close(); truncated || err != nil {
		t.Fatalf("Close: %v %v", truncated, err)
	}
	got := map[string]string{}
	for _, ch := range chunksOf(t, s) {
		got[ch[:1]] += ch[2:]
	}
	for _, stream := range []string{StreamOut, StreamErr} {
		if got[stream] != want[stream] {
			t.Fatalf("stream %s: got %d bytes, want %d, or the bytes differ", stream, len(got[stream]), len(want[stream]))
		}
	}
}

func TestChunkerRefusesAnInvalidRunAtConstruction(t *testing.T) {
	s := release.NewDirStore(t.TempDir())
	for _, tc := range []struct {
		acct, app, id string
		want          error
	}{
		{"bab", "shop", "run_../../x", ErrBadID},
		{"bab", "shop", "", ErrBadID},
		{"bab", "../other", tid, ErrBadApp},
		{"../bab", "shop", tid, ErrBadApp},
	} {
		c := NewChunker(context.Background(), s, tc.acct, tc.app, tc.id)
		n, err := io.WriteString(c.Stream(StreamOut), "x")
		if n != 1 || err != nil {
			t.Fatalf("%+v: a refused chunker still swallows writes: %d %v", tc, n, err)
		}
		if _, _, err := c.Close(); !errors.Is(err, tc.want) {
			t.Fatalf("%+v: Close err %v, want %v", tc, err, tc.want)
		}
	}
	if keys, _ := s.List(context.Background(), ""); len(keys) != 0 {
		t.Fatalf("a refused chunker wrote %q", keys)
	}
}

func TestChunkerRefusesAnUnknownStream(t *testing.T) {
	s := release.NewDirStore(t.TempDir())
	c := newChunker(context.Background(), s, "bab", "shop", tid, time.Hour, MaxOutputBytes)
	io.WriteString(c.Stream("../x"), "bad")
	_, _, err := c.Close()
	if err == nil {
		t.Fatal("an unknown stream must surface in Close")
	}
	if keys, _ := s.List(context.Background(), ""); len(keys) != 0 {
		t.Fatalf("wrote %q", keys)
	}
}

func TestChunkerDropsWritesAfterClose(t *testing.T) {
	s := release.NewDirStore(t.TempDir())
	c := newChunker(context.Background(), s, "bab", "shop", tid, time.Hour, MaxOutputBytes)
	io.WriteString(c.Stream(StreamOut), "a")
	last, _, _ := c.Close()
	io.WriteString(c.Stream(StreamOut), "late")
	last2, _, _ := c.Close()
	if last != 1 || last2 != 1 {
		t.Fatalf("last %d then %d", last, last2)
	}
	if got := chunksOf(t, s); len(got) != 1 || got[0] != "o:a" {
		t.Fatalf("chunks %q", got)
	}
}