package whisperx

import (
	"context"
	"errors"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

func TestHTTPClientTransportReplaysAbandonedMultipartBody(t *testing.T) {
	var opens atomic.Int32
	var closes atomic.Int32
	transport := roundTripFunc(func(request *http.Request) (*http.Response, error) {
		partial := make([]byte, 17)
		if _, err := io.ReadFull(request.Body, partial); err != nil {
			return nil, err
		}
		if err := request.Body.Close(); err != nil {
			return nil, err
		}
		if request.GetBody == nil {
			return nil, errors.New("request GetBody is nil")
		}
		replay, err := request.GetBody()
		if err != nil {
			return nil, err
		}
		replayedBytes, err := io.ReadAll(replay)
		if err != nil {
			_ = replay.Close()
			return nil, err
		}
		if err := replay.Close(); err != nil {
			return nil, err
		}
		parts := parseMultipartParts(t, request.Header.Get("Content-Type"), replayedBytes)
		if len(parts) != 2 || parts[0].name != "file" || parts[0].filename != "audio.flac" || string(parts[0].data) != "audio-data" {
			t.Fatalf("replayed file part = %#v, want exact audio.flac payload", parts)
		}
		if parts[1].name != "language" || string(parts[1].data) != "en" {
			t.Fatalf("replayed language part = %#v, want language=en", parts[1])
		}
		return whisperXJSONResponse(request), nil
	})
	client := newReplayTestHTTPClient(t, transport)
	client.openAudio = func(string) (io.ReadCloser, error) {
		opens.Add(1)
		return &countedReadCloser{Reader: strings.NewReader("audio-data"), closes: &closes}, nil
	}
	outPath := filepath.Join(t.TempDir(), "raw.json")

	result, err := client.Transcribe(context.Background(), TranscribeRequest{
		AudioPath:               "audio.flac",
		OutputRawTranscriptPath: outPath,
	})
	if err != nil {
		t.Fatalf("Transcribe() error = %v", err)
	}
	if result.Attempts != 1 {
		t.Fatalf("application attempts = %d, want 1", result.Attempts)
	}
	verifyJSONFile(t, outPath)
	if got := opens.Load(); got != 2 {
		t.Fatalf("audio opens = %d, want original and replay", got)
	}
	if got := closes.Load(); got != 2 {
		t.Fatalf("audio closes = %d, want all body handles disposed", got)
	}
}

func TestHTTPClientTransportReturnsReplayOpenFailure(t *testing.T) {
	replayErr := errors.New("replay open failed")
	var opens atomic.Int32
	var closes atomic.Int32
	transport := roundTripFunc(func(request *http.Request) (*http.Response, error) {
		if _, err := io.ReadFull(request.Body, make([]byte, 1)); err != nil {
			return nil, err
		}
		_ = request.Body.Close()
		_, err := request.GetBody()
		return nil, err
	})
	client := newReplayTestHTTPClient(t, transport)
	client.openAudio = func(string) (io.ReadCloser, error) {
		if opens.Add(1) == 2 {
			return nil, replayErr
		}
		return &countedReadCloser{Reader: strings.NewReader("audio-data"), closes: &closes}, nil
	}
	outPath := filepath.Join(t.TempDir(), "raw.json")

	result, err := client.Transcribe(context.Background(), TranscribeRequest{
		AudioPath:               "audio.flac",
		OutputRawTranscriptPath: outPath,
	})
	if !errors.Is(err, replayErr) {
		t.Fatalf("Transcribe() error = %v, want replay open failure", err)
	}
	if result.Attempts != 1 {
		t.Fatalf("application attempts = %d, want 1", result.Attempts)
	}
	if got := closes.Load(); got != 1 {
		t.Fatalf("audio closes = %d, want original disposed", got)
	}
	if _, statErr := os.Stat(outPath); !errors.Is(statErr, os.ErrNotExist) {
		t.Fatalf("output should not exist, stat error = %v", statErr)
	}
}

func TestHTTPClientTransportCannotSwallowReplayReadFailure(t *testing.T) {
	replayErr := errors.New("replay read failed")
	var opens atomic.Int32
	transport := roundTripFunc(func(request *http.Request) (*http.Response, error) {
		if _, err := io.ReadFull(request.Body, make([]byte, 1)); err != nil {
			return nil, err
		}
		_ = request.Body.Close()
		replay, err := request.GetBody()
		if err != nil {
			return nil, err
		}
		_, _ = io.ReadAll(replay)
		_ = replay.Close()
		return whisperXJSONResponse(request), nil
	})
	client := newReplayTestHTTPClient(t, transport)
	client.openAudio = func(string) (io.ReadCloser, error) {
		if opens.Add(1) == 2 {
			return &failingReadCloser{first: []byte("partial"), err: replayErr}, nil
		}
		return io.NopCloser(strings.NewReader("audio-data")), nil
	}
	outPath := filepath.Join(t.TempDir(), "raw.json")

	_, err := client.Transcribe(context.Background(), TranscribeRequest{
		AudioPath:               "audio.flac",
		OutputRawTranscriptPath: outPath,
	})
	if !errors.Is(err, replayErr) {
		t.Fatalf("Transcribe() error = %v, want replay read failure", err)
	}
	if _, statErr := os.Stat(outPath); !errors.Is(statErr, os.ErrNotExist) {
		t.Fatalf("output should not exist, stat error = %v", statErr)
	}
}

func TestHTTPClientTransportEarlyResponseDisposesOriginalBody(t *testing.T) {
	var closes atomic.Int32
	client := newReplayTestHTTPClient(t, roundTripFunc(func(request *http.Request) (*http.Response, error) {
		return whisperXJSONResponse(request), nil
	}))
	client.openAudio = func(string) (io.ReadCloser, error) {
		return &countedReadCloser{Reader: strings.NewReader("audio-data"), closes: &closes}, nil
	}

	_, err := client.Transcribe(context.Background(), TranscribeRequest{
		AudioPath:               "audio.flac",
		OutputRawTranscriptPath: filepath.Join(t.TempDir(), "raw.json"),
	})
	if err != nil {
		t.Fatalf("Transcribe() error = %v", err)
	}
	if got := closes.Load(); got != 1 {
		t.Fatalf("audio closes = %d, want early-response body disposed", got)
	}
}

func TestHTTPClientParentCancellationDuringReplayClosesEveryBody(t *testing.T) {
	var opens atomic.Int32
	var originalCloses atomic.Int32
	replaySource := newSignalingBlockedReadCloser()
	transport := roundTripFunc(func(request *http.Request) (*http.Response, error) {
		if _, err := io.ReadFull(request.Body, make([]byte, 1)); err != nil {
			return nil, err
		}
		_ = request.Body.Close()
		replay, err := request.GetBody()
		if err != nil {
			return nil, err
		}
		_, _ = io.ReadAll(replay)
		_ = replay.Close()
		return nil, request.Context().Err()
	})
	client := newReplayTestHTTPClient(t, transport)
	client.openAudio = func(string) (io.ReadCloser, error) {
		if opens.Add(1) == 1 {
			return &countedReadCloser{Reader: strings.NewReader("audio-data"), closes: &originalCloses}, nil
		}
		return replaySource, nil
	}
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() {
		_, err := client.Transcribe(ctx, TranscribeRequest{
			AudioPath:               "audio.flac",
			OutputRawTranscriptPath: filepath.Join(t.TempDir(), "raw.json"),
		})
		done <- err
	}()

	select {
	case <-replaySource.blocked:
		cancel()
	case <-time.After(time.Second):
		cancel()
		t.Fatal("replay source did not block")
	}
	select {
	case err := <-done:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("Transcribe() error = %v, want context cancellation", err)
		}
	case <-time.After(time.Second):
		t.Fatal("Transcribe() did not return after replay cancellation")
	}
	if got := originalCloses.Load(); got != 1 {
		t.Fatalf("original audio closes = %d, want 1", got)
	}
	select {
	case <-replaySource.closed:
	case <-time.After(time.Second):
		t.Fatal("replay audio source was not closed")
	}
}

type roundTripFunc func(*http.Request) (*http.Response, error)

func (roundTrip roundTripFunc) RoundTrip(request *http.Request) (*http.Response, error) {
	return roundTrip(request)
}

func newReplayTestHTTPClient(t *testing.T, transport http.RoundTripper) *HTTPClient {
	t.Helper()
	client, err := NewHTTPClient(HTTPClientConfig{
		TranscribeURL: "http://whisperx.test/transcribe",
		Language:      "en",
		Timeout:       time.Second,
		HTTPClient:    &http.Client{Transport: transport},
	})
	if err != nil {
		t.Fatalf("NewHTTPClient() error = %v", err)
	}
	return client
}

func whisperXJSONResponse(request *http.Request) *http.Response {
	return &http.Response{
		StatusCode: http.StatusOK,
		Status:     "200 OK",
		Header:     make(http.Header),
		Body:       io.NopCloser(strings.NewReader(`{"source":"transport-replay","segments":[]}`)),
		Request:    request,
	}
}