package common

import (
	"crypto/tls"
	"maps"
	"net/http/httptrace"
	"sync"
	"time"
)

const UpstreamTimingContextKey = "responses_upstream_timing"

// UpstreamTiming contains only transport metadata and request-relative times.
// HTTP callbacks can run concurrently with streaming and log checkpoints.
type UpstreamTiming struct {
	mu           sync.Mutex
	start        time.Time
	events       map[string]float64
	attempt      int
	channel      int
	bodyBytes    int64
	reused       bool
	gotConn      bool
	getConnCount int
	writeCount   int
	writeFailed  bool
	protocol     string
	status       int
}

func NewUpstreamTiming(start time.Time, attempt, channel int, bodyBytes int64) *UpstreamTiming {
	if start.IsZero() {
		start = time.Now()
	}
	t := &UpstreamTiming{start: start, events: make(map[string]float64), attempt: attempt, channel: channel, bodyBytes: bodyBytes}
	t.Mark("upstream_start")
	return t
}

func (t *UpstreamTiming) Mark(event string) {
	if t == nil {
		return
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	if _, exists := t.events[event]; !exists {
		t.events[event] = float64(time.Since(t.start).Microseconds()) / 1000
	}
}

func (t *UpstreamTiming) ClientTrace() *httptrace.ClientTrace {
	return &httptrace.ClientTrace{
		GetConn: func(string) {
			t.Mark("get_connection")
			t.mu.Lock()
			t.getConnCount++
			t.mu.Unlock()
		},
		DNSStart:          func(httptrace.DNSStartInfo) { t.Mark("dns_start") },
		DNSDone:           func(httptrace.DNSDoneInfo) { t.Mark("dns_done") },
		ConnectStart:      func(string, string) { t.Mark("tcp_start") },
		ConnectDone:       func(string, string, error) { t.Mark("tcp_done") },
		TLSHandshakeStart: func() { t.Mark("tls_start") },
		TLSHandshakeDone:  func(tls.ConnectionState, error) { t.Mark("tls_done") },
		GotConn: func(info httptrace.GotConnInfo) {
			t.Mark("connection_ready")
			t.mu.Lock()
			t.reused, t.gotConn = info.Reused, true
			t.mu.Unlock()
		},
		WroteRequest: func(info httptrace.WroteRequestInfo) {
			t.mu.Lock()
			defer t.mu.Unlock()
			t.writeCount++
			t.writeFailed = t.writeFailed || info.Err != nil
			if info.Err == nil {
				t.events["request_sent"] = float64(time.Since(t.start).Microseconds()) / 1000
			}
		},
		GotFirstResponseByte: func() { t.Mark("response_first_byte") },
	}
}

func (t *UpstreamTiming) Response(protocol string, status int) {
	if t == nil {
		return
	}
	t.Mark("response_headers")
	t.mu.Lock()
	defer t.mu.Unlock()
	t.protocol, t.status = protocol, status
}

func (t *UpstreamTiming) Snapshot() map[string]any {
	if t == nil {
		return nil
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	events := maps.Clone(t.events)
	result := map[string]any{
		"events_ms": events, "attempt": t.attempt, "channel_id": t.channel,
		"request_body_bytes": t.bodyBytes, "get_connection_count": t.getConnCount,
		"write_count": t.writeCount, "write_failed": t.writeFailed,
	}
	if t.gotConn {
		result["connection_reused"] = t.reused
	}
	if t.protocol != "" {
		result["protocol"], result["status"] = t.protocol, t.status
	}
	return result
}