// Package inprochttp adapts an in-process HTTP handler to net/http's
// RoundTripper contract without buffering response bodies.
package inprochttp

import (
	"context"
	"errors"
	"fmt"
	"github.com/tidbcloud/AgentTeam/internal/observability/logging"
	"io"
	"net/http"
	"net/url"
	"sync"
)

// RoundTripper dispatches requests directly to Handler. Response headers are
// returned as soon as the handler writes them, while the body streams through
// an io.Pipe with normal backpressure.
type RoundTripper struct {
	Handler http.Handler
	// Scheme optionally rewrites the synthetic request scheme before dispatch.
	// This keeps redirects emitted by ordinary HTTP routers on http(s).
	Scheme string
}

func (transport RoundTripper) RoundTrip(request *http.Request) (*http.Response, error) {
	if request == nil {
		return nil, fmt.Errorf("in-process HTTP request is nil")
	}
	requestBody := newOnceReadCloser(request.Body)
	if transport.Handler == nil {
		_ = requestBody.Close()
		return nil, fmt.Errorf("in-process HTTP handler is not configured")
	}

	ctx, cancel := context.WithCancel(request.Context())
	reader, writer := io.Pipe()
	responseWriter := &pipeResponseWriter{
		header:  make(http.Header),
		body:    writer,
		started: make(chan struct{}),
	}
	clone := request.Clone(ctx)
	clone.Body = requestBody
	if transport.Scheme != "" {
		clone.URL.Scheme = transport.Scheme
		clone.Host = "embedded.invalid"
		clone.URL.Host = "embedded.invalid"
	}
	logging.Go(nil, func() {
		defer func() {
			_ = requestBody.Close()
			if recovered := recover(); recovered != nil {
				responseWriter.fail(fmt.Errorf("in-process HTTP handler panicked: %v", recovered))
				return
			}
			responseWriter.finish()
		}()
		transport.Handler.ServeHTTP(responseWriter, clone)
	})

	select {
	case <-responseWriter.started:
		status, header := responseWriter.responseHead()
		if transport.Scheme != "" {
			if location, err := url.Parse(header.Get("Location")); err == nil && location.IsAbs() {
				location.Scheme = request.URL.Scheme
				header.Set("Location", location.String())
			}
		}
		return &http.Response{
			Status:        fmt.Sprintf("%d %s", status, http.StatusText(status)),
			StatusCode:    status,
			Proto:         "HTTP/1.1",
			ProtoMajor:    1,
			ProtoMinor:    1,
			Header:        header,
			Body:          &cancelReadCloser{ReadCloser: reader, cancel: cancel, requestBody: requestBody},
			ContentLength: -1,
			Request:       request,
		}, nil
	case <-ctx.Done():
		cancel()
		_ = requestBody.Close()
		_ = reader.CloseWithError(ctx.Err())
		return nil, ctx.Err()
	}
}

type pipeResponseWriter struct {
	mu          sync.Mutex
	header      http.Header
	body        *io.PipeWriter
	started     chan struct{}
	startedOnce sync.Once
	status      int
}

func (writer *pipeResponseWriter) Header() http.Header { return writer.header }

func (writer *pipeResponseWriter) WriteHeader(status int) {
	writer.mu.Lock()
	defer writer.mu.Unlock()
	if writer.status != 0 {
		return
	}
	writer.status = status
	writer.startedOnce.Do(func() { close(writer.started) })
}

func (writer *pipeResponseWriter) Write(payload []byte) (int, error) {
	writer.WriteHeader(http.StatusOK)
	return writer.body.Write(payload)
}

func (writer *pipeResponseWriter) Flush() { writer.WriteHeader(http.StatusOK) }

func (writer *pipeResponseWriter) finish() {
	writer.WriteHeader(http.StatusOK)
	_ = writer.body.Close()
}

func (writer *pipeResponseWriter) fail(err error) {
	writer.WriteHeader(http.StatusInternalServerError)
	_ = writer.body.CloseWithError(err)
}

func (writer *pipeResponseWriter) responseHead() (int, http.Header) {
	writer.mu.Lock()
	defer writer.mu.Unlock()
	return writer.status, writer.header.Clone()
}

type cancelReadCloser struct {
	io.ReadCloser
	cancel      context.CancelFunc
	requestBody io.Closer
}

func (body *cancelReadCloser) Close() error {
	body.cancel()
	requestErr := body.requestBody.Close()
	return errors.Join(body.ReadCloser.Close(), requestErr)
}

type onceReadCloser struct {
	io.Reader
	closer io.Closer
	once   sync.Once
	err    error
}

func newOnceReadCloser(body io.ReadCloser) *onceReadCloser {
	if body == nil {
		return &onceReadCloser{Reader: http.NoBody, closer: http.NoBody}
	}
	return &onceReadCloser{Reader: body, closer: body}
}

func (body *onceReadCloser) Close() error {
	body.once.Do(func() { body.err = body.closer.Close() })
	return body.err
}