// Package telemetry 记录 GenAI span。默认不采集提示词正文。
package telemetry

import (
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"git-ssh.u1680832.nyat.app/weichen/Logging-System.git/logsys"
)

// Span 是一次 invoke_agent、chat 或 execute_tool。
type Span struct {
	Name         string
	TraceID      string
	SpanID       string
	ParentSpanID string
	Start        time.Time
	End          time.Time
	Attrs        map[string]string
}

// 保留属性键：循环用它们把 run 级的 trace 关联传进工厂（不导出）。
const (
	AttrTraceID      = "_trace_id"
	AttrSpanID       = "_span_id"
	AttrParentSpanID = "_parent_span_id"
)

// NewTraceID / NewSpanID 生成一次 Run 的关联 ID（trace 16 字节、span 8 字节 hex）。
func NewTraceID() string { return newID(16) }
func NewSpanID() string  { return newID(8) }

// Option 是 StartFunc 的可选项。
type Option func(*factory)

type factory struct {
	endpoint    string
	serviceName string
	log         *logsys.Logger
}

// WithServiceName 设置 OTLP resource 的 service.name（默认 agent-native）。
func WithServiceName(name string) Option {
	return func(f *factory) {
		if strings.TrimSpace(name) != "" {
			f.serviceName = name
		}
	}
}

// WithEndpoint 设置 OTLP 端点（优先级高于环境变量；空 = 回退环境变量）。
func WithEndpoint(endpoint string) Option {
	return func(f *factory) {
		if strings.TrimSpace(endpoint) != "" {
			f.endpoint = strings.TrimRight(endpoint, "/")
		}
	}
}

// StartFunc 返回循环使用的追踪函数。端点来自 WithEndpoint 或环境变量
// OTEL_EXPORTER_OTLP_ENDPOINT（都为空则不导出）；结束 span 后以
// OTLP/HTTP JSON 上报，不引入新的依赖。默认不采集提示词正文。
func StartFunc(log *logsys.Logger, opts ...Option) func(name string, attrs map[string]string) func() {
	f := &factory{
		endpoint:    strings.TrimRight(os.Getenv("OTEL_EXPORTER_OTLP_ENDPOINT"), "/"),
		serviceName: "agent-native",
		log:         log,
	}
	for _, o := range opts {
		o(f)
	}
	return func(name string, attrs map[string]string) func() {
		sp := Span{
			Name: name, TraceID: newID(16), SpanID: newID(8),
			Start: time.Now(),
		}
		// 循环用保留键传入 run 级关联：一次 Run 共享 TraceID，
		// chat/execute_tool 挂到 invoke_agent 根下。
		if attrs != nil {
			if v := attrs[AttrTraceID]; v != "" {
				sp.TraceID = v
			}
			if v := attrs[AttrSpanID]; v != "" {
				sp.SpanID = v
			}
			if v := attrs[AttrParentSpanID]; v != "" {
				sp.ParentSpanID = v
			}
		}
		return func() {
			sp.End = time.Now()
			// 属性在结束时过滤：允许循环在 span 存续期间补充（如 usage）。
			sp.Attrs = publicAttrs(attrs)
			if f.log != nil {
				kvs := []any{"span", sp.Name, "trace_id", sp.TraceID, "span_id", sp.SpanID}
				if sp.ParentSpanID != "" {
					kvs = append(kvs, "parent_span_id", sp.ParentSpanID)
				}
				for k, v := range sp.Attrs {
					kvs = append(kvs, k, v)
				}
				f.log.Info("genai span", kvs...)
			}
			if f.endpoint != "" {
				if err := ExportOTLP(f.endpoint, sp, f.serviceName); err != nil && f.log != nil {
					f.log.Info("genai span export failed", "error", err)
				}
			}
		}
	}
}

// publicAttrs 丢掉提示词正文与保留键（_ 前缀的内部约定）。调用方即使
// 传入正文也不会进入 span。
func publicAttrs(in map[string]string) map[string]string {
	out := map[string]string{}
	for k, v := range in {
		if k == "prompt" || k == "messages" || k == "content" || strings.HasPrefix(k, "_") {
			continue
		}
		out[k] = v
	}
	return out
}

// ExportOTLP 把一个 span POST 到端点的 /v1/traces。
func ExportOTLP(endpoint string, sp Span, serviceName ...string) error {
	endpoint = strings.TrimRight(endpoint, "/")
	if !strings.HasSuffix(endpoint, "/v1/traces") {
		endpoint += "/v1/traces"
	}
	body, err := json.Marshal(otlpBody(sp, serviceName...))
	if err != nil {
		return err
	}
	req, err := http.NewRequest(http.MethodPost, endpoint, bytes.NewReader(body))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	return nil
}

func otlpBody(sp Span, serviceName ...string) map[string]any {
	var attrs []map[string]any
	for k, v := range publicAttrs(sp.Attrs) {
		// gen_ai.usage.* 是数值型语义（intValue），其余保持 stringValue。
		if strings.HasPrefix(k, "gen_ai.usage.") {
			if n, err := strconv.Atoi(v); err == nil {
				attrs = append(attrs, map[string]any{"key": k, "value": map[string]any{"intValue": n}})
				continue
			}
		}
		attrs = append(attrs, map[string]any{"key": k, "value": map[string]any{"stringValue": v}})
	}
	name := "agent-native"
	if len(serviceName) > 0 && strings.TrimSpace(serviceName[0]) != "" {
		name = serviceName[0]
	}
	span := map[string]any{
		"traceId":           sp.TraceID,
		"spanId":            sp.SpanID,
		"name":              sp.Name,
		"startTimeUnixNano": sp.Start.UnixNano(),
		"endTimeUnixNano":   sp.End.UnixNano(),
		"attributes":        attrs,
	}
	if sp.ParentSpanID != "" {
		span["parentSpanId"] = sp.ParentSpanID
	}
	return map[string]any{
		"resourceSpans": []any{map[string]any{
			"resource": map[string]any{
				"attributes": []map[string]any{
					{"key": "service.name", "value": map[string]any{"stringValue": name}},
				},
			},
			"scopeSpans": []any{map[string]any{
				"spans": []any{span},
			}},
		}},
	}
}

func newID(n int) string {
	b := make([]byte, n)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b)
}