package whisperx

import (
	"bytes"
	"context"
	"errors"
	"io"
	"mime"
	"mime/multipart"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func TestMultipartBodyReplayUsesIdenticalFraming(t *testing.T) {
	factory := newTestMultipartBodyFactory(t, context.Background(), "audio-data")
	defer factory.Close()

	original, err := factory.NewBody()
	if err != nil {
		t.Fatalf("NewBody() original error = %v", err)
	}
	replay, err := factory.NewBody()
	if err != nil {
		t.Fatalf("NewBody() replay error = %v", err)
	}
	originalBytes, err := io.ReadAll(original)
	if err != nil {
		t.Fatalf("ReadAll(original) error = %v", err)
	}
	replayBytes, err := io.ReadAll(replay)
	if err != nil {
		t.Fatalf("ReadAll(replay) error = %v", err)
	}
	if !bytes.Equal(originalBytes, replayBytes) {
		t.Fatal("original and replay multipart bytes differ")
	}

	parts := parseMultipartParts(t, factory.contentType, replayBytes)
	if len(parts) != 2 {
		t.Fatalf("part count = %d, want 2", len(parts))
	}
	if parts[0].name != "file" || parts[0].filename != "audio.flac" || string(parts[0].data) != "audio-data" {
		t.Fatalf("file part = %#v, want file/audio.flac with exact audio", parts[0])
	}
	if parts[1].name != "language" || string(parts[1].data) != "en" {
		t.Fatalf("language part = %#v, want language=en", parts[1])
	}
}

func TestMultipartBodyReplayRestartsAfterPartialOriginal(t *testing.T) {
	factory := newTestMultipartBodyFactory(t, context.Background(), "audio-data")
	defer factory.Close()

	original, err := factory.NewBody()
	if err != nil {
		t.Fatalf("NewBody() original error = %v", err)
	}
	partial := make([]byte, len(factory.prefix)+3)
	if _, err := io.ReadFull(original, partial); err != nil {
		t.Fatalf("ReadFull(original) error = %v", err)
	}
	if err := original.Close(); err != nil {
		t.Fatalf("Close(original) error = %v", err)
	}

	replay, err := factory.NewBody()
	if err != nil {
		t.Fatalf("NewBody() replay error = %v", err)
	}
	replayBytes, err := io.ReadAll(replay)
	if err != nil {
		t.Fatalf("ReadAll(replay) error = %v", err)
	}
	parts := parseMultipartParts(t, factory.contentType, replayBytes)
	if got := string(parts[0].data); got != "audio-data" {
		t.Fatalf("replayed audio = %q, want complete payload", got)
	}
	if err := factory.SourceError(false); err != nil {
		t.Fatalf("SourceError(false) = %v, want abandoned original ignored", err)
	}
}

func TestMultipartBodySimultaneousBodiesUseIndependentSources(t *testing.T) {
	var opens atomic.Int32
	var closes atomic.Int32
	factory, err := newMultipartBodyFactory(context.Background(), "audio.flac", "en", func(string) (io.ReadCloser, error) {
		opens.Add(1)
		return &countedReadCloser{Reader: strings.NewReader("audio-data"), closes: &closes}, nil
	})
	if err != nil {
		t.Fatalf("newMultipartBodyFactory() error = %v", err)
	}

	first, err := factory.NewBody()
	if err != nil {
		t.Fatalf("NewBody() first error = %v", err)
	}
	second, err := factory.NewBody()
	if err != nil {
		t.Fatalf("NewBody() second error = %v", err)
	}
	var firstBytes, secondBytes []byte
	var firstErr, secondErr error
	var group sync.WaitGroup
	group.Add(2)
	go func() {
		defer group.Done()
		firstBytes, firstErr = io.ReadAll(first)
	}()
	go func() {
		defer group.Done()
		secondBytes, secondErr = io.ReadAll(second)
	}()
	group.Wait()
	if firstErr != nil || secondErr != nil {
		t.Fatalf("concurrent reads returned (%v, %v)", firstErr, secondErr)
	}
	if !bytes.Equal(firstBytes, secondBytes) {
		t.Fatal("independent body bytes differ")
	}
	if got := opens.Load(); got != 2 {
		t.Fatalf("audio opens = %d, want 2", got)
	}
	if err := factory.Close(); err != nil {
		t.Fatalf("Close() error = %v", err)
	}
	if err := factory.Close(); err != nil {
		t.Fatalf("second Close() error = %v", err)
	}
	if got := closes.Load(); got != 2 {
		t.Fatalf("audio closes = %d, want 2", got)
	}
}

func TestMultipartBodyReportsSourceFailures(t *testing.T) {
	t.Run("read", func(t *testing.T) {
		sourceErr := errors.New("read source")
		factory, err := newMultipartBodyFactory(context.Background(), "audio.flac", "en", func(string) (io.ReadCloser, error) {
			return &failingReadCloser{first: []byte("partial"), err: sourceErr}, nil
		})
		if err != nil {
			t.Fatalf("newMultipartBodyFactory() error = %v", err)
		}
		body, err := factory.NewBody()
		if err != nil {
			t.Fatalf("NewBody() error = %v", err)
		}
		if _, err := io.ReadAll(body); !errors.Is(err, sourceErr) {
			t.Fatalf("ReadAll() error = %v, want source failure", err)
		}
		_ = body.Close()
		if err := factory.SourceError(false); !errors.Is(err, sourceErr) {
			t.Fatalf("SourceError(false) = %v, want source failure", err)
		}
	})

	t.Run("close after complete read", func(t *testing.T) {
		closeErr := errors.New("close source")
		factory, err := newMultipartBodyFactory(context.Background(), "audio.flac", "en", func(string) (io.ReadCloser, error) {
			return &errorCloseReadCloser{Reader: strings.NewReader("audio-data"), err: closeErr}, nil
		})
		if err != nil {
			t.Fatalf("newMultipartBodyFactory() error = %v", err)
		}
		body, err := factory.NewBody()
		if err != nil {
			t.Fatalf("NewBody() error = %v", err)
		}
		if _, err := io.ReadAll(body); err != nil {
			t.Fatalf("ReadAll() error = %v", err)
		}
		if err := body.Close(); !errors.Is(err, closeErr) {
			t.Fatalf("Close() error = %v, want source close failure", err)
		}
		if err := factory.SourceError(false); !errors.Is(err, closeErr) {
			t.Fatalf("SourceError(false) = %v, want source close failure", err)
		}
	})
}

func TestMultipartBodyCloseAndCancelReleaseBlockedSource(t *testing.T) {
	for _, test := range []struct {
		name    string
		release func(context.CancelFunc, io.ReadCloser)
	}{
		{name: "close", release: func(_ context.CancelFunc, body io.ReadCloser) { _ = body.Close() }},
		{name: "cancel", release: func(cancel context.CancelFunc, _ io.ReadCloser) { cancel() }},
	} {
		t.Run(test.name, func(t *testing.T) {
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()
			source := newSignalingBlockedReadCloser()
			factory, err := newMultipartBodyFactory(ctx, "audio.flac", "en", func(string) (io.ReadCloser, error) {
				return source, nil
			})
			if err != nil {
				t.Fatalf("newMultipartBodyFactory() error = %v", err)
			}
			defer factory.Close()
			body, err := factory.NewBody()
			if err != nil {
				t.Fatalf("NewBody() error = %v", err)
			}
			readDone := make(chan error, 1)
			go func() {
				_, err := io.ReadAll(body)
				readDone <- err
			}()
			select {
			case <-source.blocked:
			case <-time.After(time.Second):
				t.Fatal("audio source did not block")
			}
			test.release(cancel, body)
			select {
			case <-source.closed:
			case <-time.After(time.Second):
				t.Fatal("audio source was not closed")
			}
			select {
			case <-readDone:
			case <-time.After(time.Second):
				t.Fatal("blocked body read did not return")
			}
			if test.name == "cancel" {
				if _, err := factory.NewBody(); !errors.Is(err, context.Canceled) {
					t.Fatalf("NewBody() after cancellation error = %v, want context.Canceled", err)
				}
			}
		})
	}
}

func newTestMultipartBodyFactory(t *testing.T, ctx context.Context, audio string) *multipartBodyFactory {
	t.Helper()
	factory, err := newMultipartBodyFactory(ctx, "audio.flac", "en", func(string) (io.ReadCloser, error) {
		return io.NopCloser(strings.NewReader(audio)), nil
	})
	if err != nil {
		t.Fatalf("newMultipartBodyFactory() error = %v", err)
	}
	return factory
}

type multipartTestPart struct {
	name     string
	filename string
	data     []byte
}

func parseMultipartParts(t *testing.T, contentType string, body []byte) []multipartTestPart {
	t.Helper()
	mediaType, params, err := mime.ParseMediaType(contentType)
	if err != nil {
		t.Fatalf("ParseMediaType() error = %v", err)
	}
	if mediaType != "multipart/form-data" {
		t.Fatalf("media type = %q, want multipart/form-data", mediaType)
	}
	reader := multipart.NewReader(bytes.NewReader(body), params["boundary"])
	var parts []multipartTestPart
	for {
		part, err := reader.NextPart()
		if errors.Is(err, io.EOF) {
			return parts
		}
		if err != nil {
			t.Fatalf("NextPart() error = %v", err)
		}
		data, err := io.ReadAll(part)
		if err != nil {
			t.Fatalf("ReadAll(part %q) error = %v", part.FormName(), err)
		}
		parts = append(parts, multipartTestPart{name: part.FormName(), filename: part.FileName(), data: data})
	}
}

type countedReadCloser struct {
	io.Reader
	closes *atomic.Int32
	once   sync.Once
}

func (reader *countedReadCloser) Close() error {
	reader.once.Do(func() { reader.closes.Add(1) })
	return nil
}

type errorCloseReadCloser struct {
	io.Reader
	err error
}

func (reader *errorCloseReadCloser) Close() error { return reader.err }

type signalingBlockedReadCloser struct {
	blocked chan struct{}
	closed  chan struct{}
	once    sync.Once
	read    bool
}

func newSignalingBlockedReadCloser() *signalingBlockedReadCloser {
	return &signalingBlockedReadCloser{blocked: make(chan struct{}), closed: make(chan struct{})}
}

func (reader *signalingBlockedReadCloser) Read(buffer []byte) (int, error) {
	if !reader.read {
		reader.read = true
		return copy(buffer, "audio-data"), nil
	}
	reader.once.Do(func() { close(reader.blocked) })
	<-reader.closed
	return 0, errors.New("audio source closed")
}

func (reader *signalingBlockedReadCloser) Close() error {
	select {
	case <-reader.closed:
	default:
		close(reader.closed)
	}
	return nil
}