package forwarder

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/api"
)

// MultipartRequest describes a streaming internal multipart request that uses
// the same trusted Human identity projection as ForwardJSON.
type MultipartRequest struct {
	Method      string
	Path        string
	WorkspaceID string
	Identity    TrustedIdentity
	ContentType string
	Body        io.Reader
}

// ForwardMultipart streams a multipart request to the configured upstream.
// The body is never materialized by AgentTeam.
func (handler *Handler) ForwardMultipart(
	ctx context.Context,
	request MultipartRequest,
	response any,
) error {
	if handler == nil || handler.baseURL == nil {
		return fmt.Errorf("API forwarder is not configured")
	}
	path := strings.TrimSpace(request.Path)
	if path == "" || request.Body == nil || strings.TrimSpace(request.ContentType) == "" {
		return api.ErrBadRequest
	}
	rawQuery := ""
	if before, after, ok := strings.Cut(path, "?"); ok {
		path, rawQuery = before, after
	}
	if !strings.HasPrefix(path, "/") {
		path = "/" + path
	}
	method := strings.ToUpper(strings.TrimSpace(request.Method))
	if method == "" {
		method = http.MethodPost
	}
	endpoint := *handler.baseURL
	endpoint.Path = singleJoiningSlash(handler.baseURL.Path, path)
	endpoint.RawQuery = rawQuery
	upstreamRequest, err := http.NewRequestWithContext(
		ctx, method, endpoint.String(), request.Body,
	)
	if err != nil {
		return fmt.Errorf("build upstream multipart request: %w", err)
	}
	upstreamRequest.Header.Set("Content-Type", request.ContentType)
	upstreamRequest.Header.Set("Accept", "application/json")
	jwt, err := handler.issueTrustedJWTForJSONRequest(JSONRequest{
		WorkspaceID: request.WorkspaceID,
		Identity:    request.Identity,
	})
	if err != nil {
		return err
	}
	authHeader := "Bearer " + jwt
	upstreamRequest.Header.Set("Authorization", authHeader)
	upstreamRequest.Header.Set(trustedJWTHeader, authHeader)
	if workspaceID := strings.TrimSpace(request.WorkspaceID); workspaceID != "" {
		upstreamRequest.Header.Set(api.HeaderWorkspaceID, workspaceID)
	}

	client := handler.streamingHTTPClient
	if client == nil {
		client = newMultipartStreamingHTTPClient(handler.httpClient, defaultProxyTimeout)
	}
	result, err := client.Do(upstreamRequest)
	if err != nil {
		return fmt.Errorf("forward upstream multipart request: %w", err)
	}
	defer result.Body.Close()
	if result.StatusCode < http.StatusOK || result.StatusCode >= http.StatusMultipleChoices {
		return upstreamHTTPError(result)
	}
	target := response
	if target == nil {
		target = new(any)
	}
	decoder := json.NewDecoder(result.Body)
	if err := decoder.Decode(target); err != nil {
		return fmt.Errorf("decode upstream multipart response: %w", err)
	}
	var trailing any
	if err := decoder.Decode(&trailing); err != io.EOF {
		if err == nil {
			return fmt.Errorf("decode upstream multipart response: multiple JSON values")
		}
		return fmt.Errorf("decode upstream multipart response: %w", err)
	}
	return nil
}

func newMultipartStreamingHTTPClient(
	base *http.Client,
	responseHeaderTimeout time.Duration,
) *http.Client {
	if base == nil {
		base = &http.Client{Timeout: defaultProxyTimeout}
	}
	streaming := *base
	streaming.Timeout = 0
	streaming.Transport = newMultipartResponseHeaderTransport(
		base.Transport, responseHeaderTimeout,
	)
	return &streaming
}

func newMultipartResponseHeaderTransport(
	base http.RoundTripper,
	responseHeaderTimeout time.Duration,
) http.RoundTripper {
	if base == nil {
		base = http.DefaultTransport
	}
	if transport, ok := base.(*http.Transport); ok && transport != nil {
		cloned := transport.Clone()
		cloned.ResponseHeaderTimeout = responseHeaderTimeout
		return cloned
	}
	if responseHeaderTimeout <= 0 {
		return base
	}
	return &multipartBodyEOFHeaderTimeoutTransport{
		base: base, timeout: responseHeaderTimeout,
	}
}

// multipartBodyEOFHeaderTimeoutTransport is the custom-RoundTripper fallback.
// Its watchdog starts only after request-body EOF, matching the standard
// http.Transport.ResponseHeaderTimeout upload boundary.
type multipartBodyEOFHeaderTimeoutTransport struct {
	base    http.RoundTripper
	timeout time.Duration
}

func (transport *multipartBodyEOFHeaderTimeoutTransport) CloseIdleConnections() {
	if closer, ok := transport.base.(interface{ CloseIdleConnections() }); ok {
		closer.CloseIdleConnections()
	}
}

func (transport *multipartBodyEOFHeaderTimeoutTransport) RoundTrip(
	request *http.Request,
) (*http.Response, error) {
	ctx, cancel := context.WithCancel(request.Context())
	forwarded := request.Clone(ctx)
	body := &multipartHeaderTimeoutBody{
		ReadCloser: forwarded.Body,
		timeout:    transport.timeout,
		cancel:     cancel,
	}
	if body.ReadCloser == nil || body.ReadCloser == http.NoBody {
		body.ReadCloser = http.NoBody
		body.start()
	}
	forwarded.Body = body
	response, err := transport.base.RoundTrip(forwarded)
	body.finish()
	if err != nil {
		cancel()
		return nil, err
	}
	if response == nil || response.Body == nil {
		cancel()
		return response, nil
	}
	response.Body = &multipartCancelBody{ReadCloser: response.Body, cancel: cancel}
	return response, nil
}

type multipartHeaderTimeoutBody struct {
	io.ReadCloser
	timeout time.Duration
	cancel  context.CancelFunc

	mu      sync.Mutex
	started bool
	done    bool
	timer   *time.Timer
}

func (body *multipartHeaderTimeoutBody) Read(buffer []byte) (int, error) {
	count, err := body.ReadCloser.Read(buffer)
	if err == io.EOF {
		body.start()
	}
	return count, err
}

func (body *multipartHeaderTimeoutBody) start() {
	body.mu.Lock()
	defer body.mu.Unlock()
	if body.done || body.started {
		return
	}
	body.started = true
	body.timer = time.AfterFunc(body.timeout, body.cancel)
}

func (body *multipartHeaderTimeoutBody) finish() {
	body.mu.Lock()
	defer body.mu.Unlock()
	body.done = true
	if body.timer != nil {
		body.timer.Stop()
	}
}

type multipartCancelBody struct {
	io.ReadCloser
	cancel context.CancelFunc
}

func (body *multipartCancelBody) Read(buffer []byte) (int, error) {
	count, err := body.ReadCloser.Read(buffer)
	if err != nil {
		body.cancel()
	}
	return count, err
}

func (body *multipartCancelBody) Close() error {
	body.cancel()
	return body.ReadCloser.Close()
}