package timing

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

// Transport records every upstream HTTP call against a stage.
//
// It is a RoundTripper rather than a helper each client calls because there are eight
// upstream clients here and the one thing they all genuinely share is that they make
// HTTP requests. Wrapping the transport means a new integration is instrumented by
// being constructed rather than by somebody remembering to measure it.
//
// The context it reads is the request's, so a call made outside a traced request —
// a scheduled sync, a health probe — records nothing and costs one nil check.
type Transport struct {
	Stage string
	Base  http.RoundTripper
}

func (t *Transport) RoundTrip(r *http.Request) (*http.Response, error) {
	base := t.Base
	if base == nil {
		base = http.DefaultTransport
	}
	trace := From(r.Context())
	if trace == nil {
		return base.RoundTrip(r)
	}
	stage := LabelFrom(r.Context(), t.Stage)
	began := time.Now()
	resp, err := base.RoundTrip(r)
	trace.Record(stage, time.Since(began))
	if err != nil || resp == nil || resp.Body == nil {
		return resp, err
	}
	// RoundTrip returns once the headers are in, and a home row is several hundred
	// kilobytes of JSON still on the wire at that point. Reading the body is part of
	// what the upstream cost, so the stage keeps accruing until the caller closes it —
	// otherwise the largest responses are the ones the breakdown under-reports.
	resp.Body = &timedBody{ReadCloser: resp.Body, trace: trace, stage: stage}
	return resp, nil
}

type timedBody struct {
	io.ReadCloser
	trace *Trace
	stage string
	spent time.Duration
	done  bool
}

func (b *timedBody) Read(p []byte) (int, error) {
	began := time.Now()
	n, err := b.ReadCloser.Read(p)
	b.spent += time.Since(began)
	return n, err
}

func (b *timedBody) Close() error {
	err := b.ReadCloser.Close()
	if !b.done {
		b.done = true
		b.trace.Record(b.stage, b.spent)
	}
	return err
}

// Instrument wraps a client's transport in place. Called during construction, before
// anything is serving, so it needs no synchronisation — the stance the Emby client's
// other setters take.
func Instrument(client *http.Client, stage string) *http.Client {
	if client == nil {
		return nil
	}
	client.Transport = &Transport{Stage: stage, Base: client.Transport}
	return client
}