package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"time"

	"github.com/gin-gonic/gin/binding"
	"github.com/pingcap/errors"
	"github.com/tidbcloud/dataflow-service/internal/config"
	"github.com/tidbcloud/pkgv2/infra/aws"
	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
)

type HttpRequester interface {
	RequestWithHeaders(ctx context.Context, method, url string, req any, extraHeaders map[string]string) (int, []byte, error)
	SetHttpTimeout(t time.Duration)
	Clone() HttpRequester
}

type AgentHttpRequester struct {
	Doer *http.Client
}

func NewAgentHttpRequester() *AgentHttpRequester {
	httpClient := &http.Client{}
	return &AgentHttpRequester{
		Doer: &http.Client{
			Transport: otelhttp.NewTransport(httpClient.Transport),
		},
	}
}

func NewAgentHttpRequesterForGardenerAndServerless() (*AgentHttpRequester, error) {
	tlsConfig, err := aws.ProxyTlsConfig(config.Global().AwsConfig.PCAArn)
	if err != nil {
		return nil, errors.Annotate(err, "failed to get proxy TLS config")
	}
	transport := &http.Transport{
		TLSClientConfig:   tlsConfig,
		DisableKeepAlives: true,
	}

	// Wrap the transport with OpenTelemetry
	otelTransport := otelhttp.NewTransport(transport)

	httpClient := &http.Client{
		Timeout:   defaultTimeout,
		Transport: otelTransport,
	}
	return &AgentHttpRequester{
		Doer: httpClient,
	}, nil
}

func (a *AgentHttpRequester) RequestWithHeaders(ctx context.Context, method, url string, req any, extraHeaders map[string]string) (int, []byte, error) {
	var (
		payload = []byte("")
		err     error
	)
	if req != nil {
		// Use protojson.Marshal for protobuf messages to correctly serialize oneof fields and enums
		if protoMsg, ok := req.(proto.Message); ok {
			payload, err = protojson.Marshal(protoMsg)
		} else {
			payload, err = json.Marshal(req)
		}
		if err != nil {
			return 0, nil, errors.Trace(err)
		}
	}
	httpReq, err := http.NewRequestWithContext(ctx, method, url, bytes.NewReader(payload))
	if err != nil {
		return 0, nil, errors.Trace(err)
	}
	httpReq.Header.Set("Content-Type", binding.MIMEJSON)
	for k, v := range extraHeaders {
		httpReq.Header.Set(k, v)
	}
	httpResp, err := a.Doer.Do(httpReq)
	if err != nil {
		return 0, nil, errors.Trace(err)
	}
	defer httpResp.Body.Close() // nolint: errcheck

	respBytes, err := io.ReadAll(httpResp.Body)
	if err != nil {
		return 0, nil, errors.Trace(err)
	}

	return httpResp.StatusCode, respBytes, nil
}

func (a *AgentHttpRequester) SetHttpTimeout(t time.Duration) {
	a.Doer.Timeout = t
}

func (a *AgentHttpRequester) Clone() HttpRequester {
	// Create a new client with the same configuration but ensure the OpenTelemetry transport is preserved
	clonedClient := &http.Client{
		Timeout:       a.Doer.Timeout,
		Transport:     a.Doer.Transport,
		CheckRedirect: a.Doer.CheckRedirect,
		Jar:           a.Doer.Jar,
	}
	return &AgentHttpRequester{
		Doer: clonedClient,
	}
}