package modelserve

import (
	"context"

	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/trace"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
	"gitea.mynodes.xyz/ilker/llmgw/internal/obs"
)

// A job is a trace of its own, and each of its steps a span beneath it —
// registry, mirror, check, catalog, placement, smoke test — opened when the
// step begins and closed when it is done, failed or skipped. The steps are
// what a job already records for `ctl model jobs`; the trace adds where the
// time went, and from the placement step a link to the replica's own start,
// whose trace goes on down to the node.
//
// Begun by the first step rather than where a job is made, so every kind of
// job is traced by recording steps and nothing else.

const tracerName = "gitea.mynodes.xyz/ilker/llmgw/internal/control/modelserve"

// jobTrace is a job's spans. Read and written with Server.mu held.
type jobTrace struct {
	span     trace.Span
	step     trace.Span
	stepName string
}

// traceStep records a step on the job's trace. Called with s.mu held.
func (j *Job) traceStep(name, state, detail string) {
	if j.trace.span == nil {
		kind := j.Kind
		if kind == "" {
			kind = "serve"
		}
		_, j.trace.span = obs.Tracer(tracerName).Start(context.Background(), kind+" model",
			trace.WithAttributes(
				attribute.String("llmgw.job", j.ID),
				attribute.String("llmgw.repo", j.Request.Repo),
				attribute.String("llmgw.model", j.Model),
			))
	}
	if j.trace.step != nil && j.trace.stepName != name {
		// A step left running when the next began: over, whatever it said.
		j.trace.step.End()
		j.trace.step = nil
	}
	if j.trace.step == nil {
		ctx := trace.ContextWithSpan(context.Background(), j.trace.span)
		_, j.trace.step = obs.Tracer(tracerName).Start(ctx, name)
		j.trace.stepName = name
	}
	span := j.trace.step
	span.SetAttributes(attribute.String("llmgw.state", state), attribute.String("llmgw.detail", detail))
	switch state {
	case Running:
		return
	case Failed:
		span.SetStatus(codes.Error, detail)
	}
	span.End()
	j.trace.step = nil
}

// traceEnd closes the job's trace, failed when err is set. Called with s.mu
// held.
func (j *Job) traceEnd(err error) {
	if j.trace.step != nil {
		j.trace.step.End()
		j.trace.step = nil
	}
	if j.trace.span == nil {
		return
	}
	if j.Model != "" {
		j.trace.span.SetAttributes(attribute.String("llmgw.model", j.Model))
	}
	obs.End(j.trace.span, err)
	j.trace.span = nil
}

// linkReplica links the step in progress to the start of a replica it waited
// for, which is in the reconciler's trace rather than this one.
func (s *Server) linkReplica(j *Job, r reconciler.Replica) {
	sc := r.StartTrace()
	if !sc.IsValid() {
		return
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if j.trace.step != nil {
		j.trace.step.AddLink(trace.Link{SpanContext: sc,
			Attributes: []attribute.KeyValue{attribute.String("llmgw.replica", r.ID)}})
	}
}