package proxy

import (
	"io"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
)

type trackingBody struct {
	io.ReadCloser
	failed *atomic.Bool
}

func (b *trackingBody) Read(p []byte) (int, error) {
	n, err := b.ReadCloser.Read(p)
	if err != nil && err != io.EOF {
		b.failed.Store(true)
	}
	return n, err
}

func responseBodyTimeout(resp *http.Response, opts Options) time.Duration {
	if resp.StatusCode == http.StatusSwitchingProtocols {
		return 0
	}
	if isStreamResponse(resp.Header) {
		return opts.SSEIdleTimeout
	}
	return opts.ResponseTimeout
}

type timeoutBody struct {
	body      io.ReadCloser
	idle      bool
	timeout   time.Duration
	timer     *time.Timer
	mu        sync.Mutex
	done      bool
	sequence  uint64
	closeOnce sync.Once
	closeErr  error
}

func newTimeoutBody(body io.ReadCloser, timeout time.Duration, idle bool) *timeoutBody {
	t := &timeoutBody{body: body, idle: idle, timeout: timeout}
	t.resetLocked()
	return t
}

func (b *timeoutBody) Read(p []byte) (int, error) {
	n, err := b.body.Read(p)
	b.mu.Lock()
	if !b.done {
		if err != nil {
			b.done = true
			b.timer.Stop()
		} else if n > 0 && b.idle {
			b.resetLocked()
		}
	}
	b.mu.Unlock()
	return n, err
}

func (b *timeoutBody) Close() error {
	b.mu.Lock()
	if !b.done {
		b.done = true
		b.sequence++
		b.timer.Stop()
	}
	b.mu.Unlock()
	return b.closeUnderlying()
}

func (b *timeoutBody) resetLocked() {
	if b.timer != nil {
		b.timer.Stop()
	}
	b.sequence++
	sequence := b.sequence
	b.timer = time.AfterFunc(b.timeout, func() { b.expire(sequence) })
}

func (b *timeoutBody) expire(sequence uint64) {
	b.mu.Lock()
	if b.done || sequence != b.sequence {
		b.mu.Unlock()
		return
	}
	b.done = true
	b.mu.Unlock()
	_ = b.closeUnderlying()
}

func (b *timeoutBody) closeUnderlying() error {
	b.closeOnce.Do(func() { b.closeErr = b.body.Close() })
	return b.closeErr
}