package agent

import (
	"encoding/json"
	"path/filepath"
	"strings"
	"time"
	"unicode"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/drivercontract"
)

const (
	tokenSourceKindTool     = "tool"
	tokenSourceKindCommand  = "command"
	tokenSourceKindContext  = "context"
	tokenSourceKindDirect   = "direct"
	tokenSourceKindMultiple = "multiple"
	tokenSourceKindUnknown  = "unknown"

	maxTokenSourceNameBytes = 120

	maxPendingTokenSourcesPerSession = 64
	maxTokenSourceSessions           = 16
)

// TokenSource is the bounded, persistence-safe identity of one tool source.
// It intentionally cannot carry tool arguments, paths, or output.
type TokenSource struct {
	Kind string
	Name string
}

// TokenSourceObservation is the agent-owned handoff to the persistence
// layer. Source is already normalized and the shape deliberately omits runtime
// correlation identities and all raw tool content.
type TokenSourceObservation struct {
	Timestamp         int64
	AgentID           string
	Runtime           string
	Model             string
	Source            TokenSource
	InputTokens       uint64
	CachedInputTokens uint64
	OutputTokens      uint64
	CallCount         uint64
}

type tokenSourceRecorder func(TokenSourceObservation)

// TokenSourceAuditObservation is the agent-owned diagnostic handoff. Unlike
// TokenSourceObservation it may carry the raw shell command, so it must only be
// written to the bounded local audit corpus and never to TSDB labels or APIs.
type TokenSourceAuditObservation struct {
	Timestamp   time.Time
	Event       string
	Runtime     string
	ToolName    string
	RawCommand  string
	Source      TokenSource
	Decision    string
	Candidates  []tokenSourceCandidate
	Allocations []TokenSourceObservation
}

type tokenSourceAuditRecorder func(TokenSourceAuditObservation)

type tokenSourceConsumption struct {
	source     TokenSource
	decision   string
	candidates []tokenSourceCandidate
}

type tokenSourceCandidate struct {
	source       TokenSource
	inputTokens  uint64
	outputTokens uint64
}

type pendingTokenSource struct {
	callID       string
	source       TokenSource
	completed    bool
	consumed     bool
	sequence     uint64
	inputTokens  uint64
	outputTokens uint64
}

type tokenSourceSessionState struct {
	sources  []pendingTokenSource
	lastSeen uint64
}

type tokenSourceState struct {
	sessions map[string]*tokenSourceSessionState
	sequence uint64
}

func (ap *Process) observeTokenSourceToolCallEstimate(
	sessionID string,
	callID string,
	source TokenSource,
	inputTokens uint64,
) {
	ap.mu.Lock()
	defer ap.mu.Unlock()

	session := ap.tokenSourceSessionLocked(sessionID, true)
	callID = strings.TrimSpace(callID)
	if callID != "" {
		for idx := range session.sources {
			item := &session.sources[idx]
			if item.callID != callID {
				continue
			}
			item.source = source
			item.inputTokens = inputTokens
			item.sequence = ap.nextTokenSourceSequenceLocked()
			session.lastSeen = item.sequence
			return
		}
	}
	sequence := ap.nextTokenSourceSequenceLocked()
	session.sources = append(session.sources, pendingTokenSource{
		callID:      callID,
		source:      source,
		sequence:    sequence,
		inputTokens: inputTokens,
	})
	session.lastSeen = sequence
	if overflow := len(session.sources) - maxPendingTokenSourcesPerSession; overflow > 0 {
		session.sources = append([]pendingTokenSource(nil), session.sources[overflow:]...)
	}
}

func (ap *Process) observeTokenSourceCall(ev drivercontract.ParsedEvent) {
	source, decision := normalizedTokenSourceDecision(ev.ToolName, ev.ToolInput)
	runtime := ap.tokenSourceRuntime()
	ap.observeTokenSourceToolCallEstimate(
		ev.SessionID,
		ev.ToolCallID,
		source,
		estimateTokenSourceContent(ev.ToolInput, runtime),
	)
	if isShellTokenSource(ev.ToolName) {
		ap.recordTokenSourceAudit(TokenSourceAuditObservation{
			Timestamp:  time.Now().UTC(),
			Event:      "shell_source",
			Runtime:    runtime,
			ToolName:   strings.TrimSpace(ev.ToolName),
			RawCommand: shellCommandFromInput(ev.ToolInput),
			Source:     source,
			Decision:   decision,
		})
	}
}

func (ap *Process) observeTokenSourceToolOutputEstimate(
	sessionID string,
	callID string,
	source TokenSource,
	outputTokens uint64,
) {
	ap.mu.Lock()
	defer ap.mu.Unlock()

	session := ap.tokenSourceSessionLocked(sessionID, false)
	if session == nil {
		return
	}
	callID = strings.TrimSpace(callID)
	match := -1
	if callID != "" {
		for idx := len(session.sources) - 1; idx >= 0; idx-- {
			if session.sources[idx].callID == callID {
				match = idx
				break
			}
		}
	} else {
		for idx := len(session.sources) - 1; idx >= 0; idx-- {
			item := session.sources[idx]
			if item.completed || item.consumed {
				continue
			}
			if source.Kind != tokenSourceKindUnknown && item.source != source {
				continue
			}
			match = idx
			break
		}
	}
	if match < 0 {
		return
	}
	item := &session.sources[match]
	if source.Kind != tokenSourceKindUnknown {
		item.source = source
	}
	item.completed = true
	item.outputTokens = outputTokens
	item.sequence = ap.nextTokenSourceSequenceLocked()
	session.lastSeen = item.sequence
}

func (ap *Process) observeTokenSourceOutput(ev drivercontract.ParsedEvent) {
	ap.observeTokenSourceToolOutputEstimate(
		ev.SessionID,
		ev.ToolCallID,
		normalizedTokenSource(ev.ToolName, nil),
		estimateTokenSourceText(ev.Text, ap.tokenSourceRuntime()),
	)
}

func (ap *Process) tokenSourceRuntime() string {
	ap.mu.Lock()
	defer ap.mu.Unlock()
	if ap.driver == nil {
		return ""
	}
	return strings.TrimSpace(ap.driver.Name())
}

func (ap *Process) consumeTokenSource(evSessionID string, correlationID string) tokenSourceConsumption {
	ap.mu.Lock()
	defer ap.mu.Unlock()

	sessionKey, primary := ap.tokenSourceSessionKeyLocked(evSessionID)
	session := ap.tokenSourceState.sessions[sessionKey]
	if session == nil {
		if primary {
			return tokenSourceConsumption{
				source:   contextTokenSource(),
				decision: "context",
			}
		}
		return tokenSourceConsumption{
			source:   unknownTokenSource(),
			decision: "unknown_session",
		}
	}
	correlationID = strings.TrimSpace(correlationID)
	if correlationID != "" {
		for idx := range session.sources {
			item := &session.sources[idx]
			if item.callID != correlationID || !item.completed || item.consumed {
				continue
			}
			item.consumed = true
			if item.source.Kind == tokenSourceKindUnknown {
				return tokenSourceConsumption{
					source:     unknownTokenSource(),
					decision:   "unknown_source",
					candidates: []tokenSourceCandidate{pendingTokenSourceCandidate(*item)},
				}
			}
			return tokenSourceConsumption{
				source:     item.source,
				decision:   "exact",
				candidates: []tokenSourceCandidate{pendingTokenSourceCandidate(*item)},
			}
		}
		// An explicit runtime correlation is authoritative. A miss must not
		// consume unrelated completed calls merely because they share a session.
		return tokenSourceConsumption{
			source:   unknownTokenSource(),
			decision: "correlation_miss",
		}
	}

	completed := make([]int, 0, len(session.sources))
	for idx := range session.sources {
		item := session.sources[idx]
		if item.completed && !item.consumed {
			completed = append(completed, idx)
		}
	}
	switch len(completed) {
	case 0:
		if primary {
			return tokenSourceConsumption{
				source:   contextTokenSource(),
				decision: "context",
			}
		}
		return tokenSourceConsumption{
			source:   unknownTokenSource(),
			decision: "unknown_session",
		}
	case 1:
		item := &session.sources[completed[0]]
		item.consumed = true
		if item.source.Kind == tokenSourceKindUnknown {
			return tokenSourceConsumption{
				source:     unknownTokenSource(),
				decision:   "unknown_source",
				candidates: []tokenSourceCandidate{pendingTokenSourceCandidate(*item)},
			}
		}
		return tokenSourceConsumption{
			source:     item.source,
			decision:   "associated",
			candidates: []tokenSourceCandidate{pendingTokenSourceCandidate(*item)},
		}
	default:
		candidates := make([]tokenSourceCandidate, 0, len(completed))
		for _, idx := range completed {
			session.sources[idx].consumed = true
			candidates = append(candidates, pendingTokenSourceCandidate(session.sources[idx]))
		}
		return tokenSourceConsumption{
			source:     multipleTokenSource(),
			decision:   "aggregate",
			candidates: candidates,
		}
	}
}

func pendingTokenSourceCandidate(item pendingTokenSource) tokenSourceCandidate {
	return tokenSourceCandidate{
		source:       item.source,
		inputTokens:  item.inputTokens,
		outputTokens: item.outputTokens,
	}
}

func (ap *Process) recordTokenSourceAudit(observation TokenSourceAuditObservation) {
	ap.mu.Lock()
	recorder := ap.tokenSourceAuditRecorder
	if observation.Runtime == "" && ap.driver != nil {
		observation.Runtime = strings.TrimSpace(ap.driver.Name())
	}
	ap.mu.Unlock()
	if recorder == nil {
		return
	}
	if observation.Timestamp.IsZero() {
		observation.Timestamp = time.Now().UTC()
	}
	recorder(observation)
}

func (ap *Process) recordTokenSourceConsumptionAudit(
	runtime string,
	consumption tokenSourceConsumption,
	allocations []TokenSourceObservation,
) {
	ap.recordTokenSourceAudit(TokenSourceAuditObservation{
		Timestamp:   time.Now().UTC(),
		Event:       "usage_source",
		Runtime:     strings.TrimSpace(runtime),
		Source:      consumption.source,
		Decision:    consumption.decision,
		Candidates:  append([]tokenSourceCandidate(nil), consumption.candidates...),
		Allocations: append([]TokenSourceObservation(nil), allocations...),
	})
}

func (ap *Process) resetTokenSource(sessionID string, all bool) {
	ap.mu.Lock()
	defer ap.mu.Unlock()
	if all {
		ap.tokenSourceState = tokenSourceState{}
		return
	}
	if len(ap.tokenSourceState.sessions) == 0 {
		return
	}
	sessionKey, _ := ap.tokenSourceSessionKeyLocked(sessionID)
	delete(ap.tokenSourceState.sessions, sessionKey)
}

func (ap *Process) resetIncompleteTokenSource(sessionID string) {
	ap.mu.Lock()
	defer ap.mu.Unlock()
	sessionKey, _ := ap.tokenSourceSessionKeyLocked(sessionID)
	session := ap.tokenSourceState.sessions[sessionKey]
	if session == nil {
		return
	}
	kept := session.sources[:0]
	for _, item := range session.sources {
		if item.completed && !item.consumed {
			kept = append(kept, item)
		}
	}
	session.sources = kept
	if len(session.sources) == 0 {
		delete(ap.tokenSourceState.sessions, sessionKey)
	}
}

func (ap *Process) prepareTokenSourceSessionInit(sessionID string) {
	ap.mu.Lock()
	defer ap.mu.Unlock()

	nextSessionID := strings.TrimSpace(sessionID)
	currentSessionID := strings.TrimSpace(ap.SessionID)
	if nextSessionID == "" || nextSessionID == currentSessionID {
		return
	}
	if currentSessionID != "" {
		nextSession := ap.tokenSourceState.sessions[nextSessionID]
		sequence := ap.tokenSourceState.sequence
		ap.tokenSourceState = tokenSourceState{sequence: sequence}
		if nextSession != nil {
			ap.tokenSourceState.sessions = map[string]*tokenSourceSessionState{
				nextSessionID: nextSession,
			}
		}
		return
	}
	if len(ap.tokenSourceState.sessions) == 0 {
		return
	}
	placeholder := ap.tokenSourceState.sessions["\x00primary"]
	if placeholder == nil {
		return
	}
	delete(ap.tokenSourceState.sessions, "\x00primary")
	if existing := ap.tokenSourceState.sessions[nextSessionID]; existing != nil {
		existing.sources = append(placeholder.sources, existing.sources...)
		if overflow := len(existing.sources) - maxPendingTokenSourcesPerSession; overflow > 0 {
			existing.sources = append([]pendingTokenSource(nil), existing.sources[overflow:]...)
		}
		if placeholder.lastSeen > existing.lastSeen {
			existing.lastSeen = placeholder.lastSeen
		}
		return
	}
	ap.tokenSourceState.sessions[nextSessionID] = placeholder
}

func (ap *Process) pendingTokenSourceCount(sessionID string) int {
	ap.mu.Lock()
	defer ap.mu.Unlock()
	sessionKey, _ := ap.tokenSourceSessionKeyLocked(sessionID)
	session := ap.tokenSourceState.sessions[sessionKey]
	if session == nil {
		return 0
	}
	return len(session.sources)
}

func (ap *Process) tokenSourceSessionCount() int {
	ap.mu.Lock()
	defer ap.mu.Unlock()
	return len(ap.tokenSourceState.sessions)
}

func (ap *Process) tokenSourceSessionLocked(sessionID string, create bool) *tokenSourceSessionState {
	sessionKey, _ := ap.tokenSourceSessionKeyLocked(sessionID)
	if ap.tokenSourceState.sessions == nil {
		if !create {
			return nil
		}
		ap.tokenSourceState.sessions = make(map[string]*tokenSourceSessionState)
	}
	if session := ap.tokenSourceState.sessions[sessionKey]; session != nil {
		return session
	}
	if !create {
		return nil
	}
	if len(ap.tokenSourceState.sessions) >= maxTokenSourceSessions {
		var oldestKey string
		var oldestSequence uint64
		for key, session := range ap.tokenSourceState.sessions {
			if oldestKey == "" || session.lastSeen < oldestSequence {
				oldestKey = key
				oldestSequence = session.lastSeen
			}
		}
		delete(ap.tokenSourceState.sessions, oldestKey)
	}
	session := &tokenSourceSessionState{}
	ap.tokenSourceState.sessions[sessionKey] = session
	return session
}

func (ap *Process) tokenSourceSessionKeyLocked(sessionID string) (string, bool) {
	sessionID = strings.TrimSpace(sessionID)
	primarySessionID := strings.TrimSpace(ap.SessionID)
	if sessionID == "" {
		if primarySessionID == "" {
			return "\x00primary", true
		}
		return primarySessionID, true
	}
	return sessionID, primarySessionID != "" && sessionID == primarySessionID
}

func (ap *Process) nextTokenSourceSequenceLocked() uint64 {
	ap.tokenSourceState.sequence++
	return ap.tokenSourceState.sequence
}

func normalizedTokenSource(toolName string, toolInput any) TokenSource {
	source, _ := normalizedTokenSourceDecision(toolName, toolInput)
	return source
}

func normalizedTokenSourceDecision(toolName string, toolInput any) (TokenSource, string) {
	toolName = strings.TrimSpace(toolName)
	if isShellTokenSource(toolName) {
		return normalizedShellTokenSourceDecision(toolInput)
	}
	if server, tool, ok := splitMCPTokenSource(toolName); ok {
		source := validatedTokenSource(tokenSourceKindTool, server+"."+tool)
		if source.Kind == tokenSourceKindUnknown {
			return source, "invalid_tool"
		}
		return source, "mcp_tool"
	}
	if strings.HasPrefix(strings.ToLower(toolName), "mcp_") {
		return unknownTokenSource(), "invalid_tool"
	}
	source := validatedTokenSource(tokenSourceKindTool, strings.ToLower(toolName))
	if source.Kind == tokenSourceKindUnknown {
		return source, "invalid_tool"
	}
	return source, "tool"
}

func directTokenSource() TokenSource {
	return TokenSource{Kind: tokenSourceKindDirect, Name: tokenSourceKindDirect}
}

func contextTokenSource() TokenSource {
	return TokenSource{Kind: tokenSourceKindContext, Name: tokenSourceKindContext}
}

func multipleTokenSource() TokenSource {
	return TokenSource{Kind: tokenSourceKindMultiple, Name: "multiple_tools"}
}

func unknownTokenSource() TokenSource {
	return TokenSource{Kind: tokenSourceKindUnknown, Name: tokenSourceKindUnknown}
}

func validatedTokenSource(kind string, name string) TokenSource {
	name = strings.ToLower(strings.TrimSpace(name))
	if name == "" || len(name) > maxTokenSourceNameBytes {
		return unknownTokenSource()
	}
	for _, r := range name {
		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') ||
			r == '.' || r == ':' || r == '_' || r == '-' {
			continue
		}
		return unknownTokenSource()
	}
	return TokenSource{Kind: kind, Name: name}
}

func isShellTokenSource(name string) bool {
	name = strings.ToLower(strings.TrimSpace(name))
	if idx := strings.LastIndexAny(name, "./"); idx >= 0 {
		name = name[idx+1:]
	}
	switch name {
	case "shell", "bash", "exec_command":
		return true
	default:
		return false
	}
}

func splitMCPTokenSource(name string) (string, string, bool) {
	name = strings.ToLower(strings.TrimSpace(name))
	if strings.HasPrefix(name, "mcp__") {
		body := strings.TrimPrefix(name, "mcp__")
		idx := strings.LastIndex(body, "__")
		if idx <= 0 || idx == len(body)-2 {
			return "", "", false
		}
		server, tool := strings.TrimSpace(body[:idx]), strings.TrimSpace(body[idx+2:])
		return server, tool, server != "" && tool != ""
	}
	if !strings.HasPrefix(name, "mcp_") {
		return "", "", false
	}
	body := strings.TrimPrefix(name, "mcp_")
	idx := strings.IndexByte(body, '_')
	if idx <= 0 || idx == len(body)-1 {
		return "", "", false
	}
	return strings.TrimSpace(body[:idx]), strings.TrimSpace(body[idx+1:]), true
}

func normalizedShellTokenSource(input any) TokenSource {
	source, _ := normalizedShellTokenSourceDecision(input)
	return source
}

func normalizedShellTokenSourceDecision(input any) (TokenSource, string) {
	command := shellCommandFromInput(input)
	fields := conservativeShellFields(command)
	executable, rest, decision := firstShellExecutable(fields)
	if executable == "" {
		return unknownTokenSource(), "missing_command"
	}
	if executable == "agentteam" && len(rest) >= 3 && strings.EqualFold(rest[0], "exec") {
		domain := strings.ToLower(strings.TrimSpace(rest[1]))
		action := strings.ToLower(strings.TrimSpace(rest[2]))
		source := validatedTokenSource(tokenSourceKindCommand, domain+":"+action)
		if source.Kind != tokenSourceKindUnknown {
			return source, "agentteam_exec"
		}
		return unknownTokenSource(), "invalid_executable"
	}
	source := validatedTokenSource(tokenSourceKindCommand, executable)
	if source.Kind == tokenSourceKindUnknown {
		return source, "invalid_executable"
	}
	return source, decision
}

func shellCommandFromInput(input any) string {
	switch value := input.(type) {
	case string:
		trimmed := strings.TrimSpace(value)
		if trimmed == "" {
			return ""
		}
		var object map[string]any
		if strings.HasPrefix(trimmed, "{") && json.Unmarshal([]byte(trimmed), &object) == nil {
			return shellCommandFromMap(object)
		}
		return trimmed
	case []byte:
		return shellCommandFromInput(string(value))
	case json.RawMessage:
		return shellCommandFromInput(string(value))
	case map[string]any:
		return shellCommandFromMap(value)
	case map[string]string:
		for _, key := range []string{"command", "cmd"} {
			if command := strings.TrimSpace(value[key]); command != "" {
				return command
			}
		}
	}
	return ""
}

func shellCommandFromMap(input map[string]any) string {
	for _, key := range []string{"command", "cmd"} {
		if command, ok := input[key].(string); ok {
			if command = strings.TrimSpace(command); command != "" {
				return command
			}
		}
	}
	return ""
}

func conservativeShellFields(command string) []string {
	fields := make([]string, 0, 8)
	var current strings.Builder
	var quote rune
	escaped := false
	flush := func() {
		if current.Len() == 0 {
			return
		}
		fields = append(fields, current.String())
		current.Reset()
	}
	for _, r := range strings.TrimSpace(command) {
		if escaped {
			current.WriteRune(r)
			escaped = false
			continue
		}
		if r == '\\' && quote != '\'' {
			escaped = true
			continue
		}
		if quote != 0 {
			if r == quote {
				quote = 0
				continue
			}
			current.WriteRune(r)
			continue
		}
		if r == '\'' || r == '"' {
			quote = r
			continue
		}
		if unicode.IsSpace(r) {
			flush()
			continue
		}
		if strings.ContainsRune("|;&<>", r) {
			flush()
			break
		}
		current.WriteRune(r)
	}
	if quote != 0 {
		return nil
	}
	if escaped {
		current.WriteRune('\\')
	}
	flush()
	return fields
}

func firstShellExecutable(fields []string) (string, []string, string) {
	for len(fields) > 0 {
		token := strings.TrimSpace(fields[0])
		fields = fields[1:]
		if token == "" || isShellAssignment(token) {
			continue
		}
		base := strings.ToLower(filepath.Base(token))
		switch base {
		case "command":
			remaining, lookup, valid := shellCommandOperands(fields)
			if lookup || !valid {
				decision := "command_builtin"
				if lookup {
					decision = "command_lookup"
				}
				return "command", remaining, decision
			}
			fields = remaining
			continue
		case "builtin", "nohup":
			continue
		case "env":
			for len(fields) > 0 && (strings.HasPrefix(fields[0], "-") || isShellAssignment(fields[0])) {
				fields = fields[1:]
			}
			continue
		case "sudo":
			fields = skipSudoOptions(fields)
			continue
		case "sh", "bash", "zsh":
			if len(fields) > 1 && strings.HasPrefix(fields[0], "-") && strings.Contains(fields[0][1:], "c") {
				return firstShellExecutable(conservativeShellFields(fields[1]))
			}
		}
		return base, fields, "executable"
	}
	return "", nil, "missing_command"
}

// shellCommandOperands separates the POSIX command builtin's own flags from
// the command it may execute. Lookup flags (-v/-V, including combinations such
// as -pv) inspect a name but do not execute it, so they stay attributed to the
// builtin instead of the inspected operand. Unknown flags fail closed to the
// builtin and can be refined later from the audit corpus.
func shellCommandOperands(fields []string) (remaining []string, lookup bool, valid bool) {
	remaining = fields
	for len(remaining) > 0 {
		option := remaining[0]
		if option == "--" {
			return remaining[1:], lookup, true
		}
		if option == "-" || !strings.HasPrefix(option, "-") {
			return remaining, lookup, true
		}
		for _, flag := range option[1:] {
			switch flag {
			case 'p':
			case 'v', 'V':
				lookup = true
			default:
				return remaining, false, false
			}
		}
		remaining = remaining[1:]
	}
	return remaining, lookup, true
}

func isShellAssignment(value string) bool {
	key, _, ok := strings.Cut(value, "=")
	if !ok || key == "" {
		return false
	}
	for idx, r := range key {
		if idx == 0 && !(unicode.IsLetter(r) || r == '_') {
			return false
		}
		if idx > 0 && !(unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_') {
			return false
		}
	}
	return true
}

func skipSudoOptions(fields []string) []string {
	for len(fields) > 0 && strings.HasPrefix(fields[0], "-") {
		option := fields[0]
		fields = fields[1:]
		switch option {
		case "-u", "-g", "-h", "-p", "-r", "-t", "-c",
			"--user", "--group", "--host", "--prompt", "--role", "--type", "--command-timeout":
			if len(fields) > 0 {
				fields = fields[1:]
			}
		default:
			if strings.HasPrefix(option, "--user=") || strings.HasPrefix(option, "--group=") ||
				strings.HasPrefix(option, "--host=") || strings.HasPrefix(option, "--prompt=") {
				continue
			}
		}
	}
	return fields
}