package metrics

import (
	"path/filepath"
	"strings"
	"time"
)

const (
	tokenTrackingFilename = "tokentracking.tsdb"

	// TokenSourceKindTool identifies an MCP, function, or runtime-native tool.
	TokenSourceKindTool = "tool"
	// TokenSourceKindCommand identifies a normalized shell command.
	TokenSourceKindCommand = "command"
	// TokenSourceKindContext identifies ordinary model context and conservative
	// remainders that cannot be assigned to one tool.
	TokenSourceKindContext = "context"
	// TokenSourceKindDirect identifies a model call with no preceding tool.
	TokenSourceKindDirect = "direct"
	// TokenSourceKindMultiple identifies an aggregate call preceded by multiple tools.
	TokenSourceKindMultiple = "multiple"
	// TokenSourceKindUnknown identifies usage that could not be safely classified.
	TokenSourceKindUnknown = "unknown"

	// tokenSourceTrendTopN bounds the per-source trend series; lower-volume
	// sources roll up into one global "other" bucket.
	tokenSourceTrendTopN = 15
)

var tokenTrackingPath string

func init() {
	tokenTrackingPath = filepath.Join(statsDir, tokenTrackingFilename)
}

// TokenSourceObservation is one sanitized provider model-call observation. The
// runtime and model are validated (to align with the Overview recording gate)
// but are not persisted as labels — token sources are attributed only by agent,
// source kind, and source name.
type TokenSourceObservation struct {
	Timestamp        int64
	AgentID          string
	Runtime          string
	Model            string
	SourceKind       string
	SourceName       string
	InputToken       uint64
	CachedInputToken uint64
	OutputToken      uint64
	// CallCount records actual tool/command calls represented by this row.
	// Context remainders carry token volume with zero calls.
	CallCount uint64
}

// TokenSourceRecorder accepts sanitized token-source observations.
type TokenSourceRecorder interface {
	RecordTokenTracking(TokenSourceObservation)
}

// TokenSourceQuerier returns bounded Machine Dashboard token-source data.
type TokenSourceQuerier interface {
	QueryTokenTracking(TokenSourceQuery) TokenSourceResult
}

// TokenSourceQuery scopes one aggregate query to an explicit Machine Agent set
// over a time window.
type TokenSourceQuery struct {
	AgentIDs []string
	Window   TimeWindow
}

// TokenSourceTrendPoint is one time bucket of total token consumption keyed by
// source name (plus the roll-up "other" key when present).
type TokenSourceTrendPoint struct {
	Timestamp int64             `json:"timestamp"`
	Values    map[string]uint64 `json:"values"`
}

// TokenSourceTrendSeries is one stacked-chart series. Key is an opaque, stable
// identifier (never the raw source name, so a real source named "other" cannot
// collide with the roll-up, and identical names of different kinds stay
// distinct); Label is the display name. SourceKind describes named sources and
// is empty for the global remainder, which can contain multiple kinds.
type TokenSourceTrendSeries struct {
	Key        string `json:"key"`
	Label      string `json:"label"`
	SourceKind string `json:"sourceKind"`
}

// TokenSourceTrend is the bounded stacked-chart response. Series lists the
// Top N source identities plus one global remainder; Data carries the per-bucket
// totals keyed by the opaque series Key.
type TokenSourceTrend struct {
	Series []TokenSourceTrendSeries `json:"series"`
	Data   []TokenSourceTrendPoint  `json:"data"`
}

// TokenSourceResult is the Machine Dashboard token-source contract. The
// bounded stacked trend is the only per-source payload.
type TokenSourceResult struct {
	Available bool             `json:"available"`
	Trend     TokenSourceTrend `json:"trend"`
}

func emptyTokenSourceResult(window TimeWindow, available bool) TokenSourceResult {
	_ = window
	return TokenSourceResult{
		Available: available,
		Trend: TokenSourceTrend{
			Series: []TokenSourceTrendSeries{},
			Data:   []TokenSourceTrendPoint{},
		},
	}
}

// UnavailableTokenSourceResult returns the stable nil-store Dashboard shape.
func UnavailableTokenSourceResult(window TimeWindow) TokenSourceResult {
	return emptyTokenSourceResult(window, false)
}

func sanitizeTokenSourceObservation(item TokenSourceObservation) (TokenSourceObservation, bool) {
	item.AgentID = strings.TrimSpace(item.AgentID)
	item.Runtime = normalizedTokenTrackingLabel(item.Runtime, 48)
	item.Model = normalizedTokenTrackingLabel(item.Model, 120)
	item.SourceKind = strings.ToLower(strings.TrimSpace(item.SourceKind))
	item.SourceName = normalizedTokenSourceName(item.SourceName, 120)

	if item.AgentID == "" || len(item.AgentID) > 120 || item.Runtime == "" || item.Model == "" {
		return TokenSourceObservation{}, false
	}
	if !validTokenSourceKind(item.SourceKind) {
		item.SourceKind = TokenSourceKindUnknown
		item.SourceName = TokenSourceKindUnknown
	}
	if item.SourceName == "" {
		item.SourceKind = TokenSourceKindUnknown
		item.SourceName = TokenSourceKindUnknown
	}
	// Preserve the pre-allocation recorder contract for tool/direct/unknown
	// observations that predate CallCount. Context is the only deliberate
	// token-bearing, zero-call source.
	if item.SourceKind != TokenSourceKindContext && item.CallCount == 0 {
		item.CallCount = 1
	}
	if item.Timestamp <= 0 {
		item.Timestamp = time.Now().UnixMilli()
	}
	if item.InputToken == 0 && item.CachedInputToken == 0 && item.OutputToken == 0 && item.CallCount == 0 {
		return TokenSourceObservation{}, false
	}
	return item, true
}

func normalizedTokenTrackingLabel(value string, limit int) string {
	value = strings.ToLower(strings.TrimSpace(value))
	if value == "" || len(value) > limit {
		return ""
	}
	for _, r := range value {
		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') ||
			r == '.' || r == ':' || r == '_' || r == '-' || r == '/' {
			continue
		}
		return ""
	}
	return value
}

// normalizedTokenSourceName sanitizes persisted source names. It deliberately
// rejects '/' (unlike runtime/model labels) so a path-like source name can
// never be written to the TSDB or returned to the UI — the persistence layer
// fails closed even if an upstream producer handed it a path.
func normalizedTokenSourceName(value string, limit int) string {
	value = strings.ToLower(strings.TrimSpace(value))
	if value == "" || len(value) > limit {
		return ""
	}
	for _, r := range value {
		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') ||
			r == '.' || r == ':' || r == '_' || r == '-' {
			continue
		}
		return ""
	}
	return value
}

func validTokenSourceKind(value string) bool {
	switch value {
	case TokenSourceKindTool, TokenSourceKindCommand, TokenSourceKindContext, TokenSourceKindDirect,
		TokenSourceKindMultiple, TokenSourceKindUnknown:
		return true
	default:
		return false
	}
}