package dispatcher

import (
	"errors"
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent"
	"github.com/tidbcloud/AgentTeam/desktop/internal/identitynotice"
	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
	"github.com/tidbcloud/AgentTeam/desktop/internal/workspace"
)

var errAgentIdentityUnchanged = errors.New("agent identity unchanged")

// UpdateAgentIdentity handles an agent:update_identity message. It updates
// the persisted agent config, rewrites MEMORY.md, appends a historical
// checkpoint note, and saves an identity notice if the agent is running or
// has an active session.
func (d *Dispatcher) UpdateAgentIdentity(serverAgentID, name, displayName string) {
	agentID := strings.TrimSpace(serverAgentID)
	if agentID == "" {
		return
	}

	nextName, nextDisplayName := normalizeAgentIdentityNames(name, displayName)
	if nextName == "" {
		return
	}

	workDir := sessionWorkDir(agentID)
	var (
		currentCfg *types.AgentStartConfig
		oldName    string
	)
	currentCfg, oldName, identityChanged, err := persistAgentIdentityStartConfig(workDir, agentID, nextName, nextDisplayName)
	if err != nil {
		slog.Warn("persist session resume state for identity update", "agent_id", agentID, "error", err)
		if currentCfg == nil {
			currentCfg = &types.AgentStartConfig{Name: nextName, DisplayName: nextDisplayName}
		}
		if oldName == "" {
			oldName = agentID
		}
	}

	newEffectiveName := effectiveIdentityName(currentCfg, agentID)
	if !identityChanged {
		d.RefreshRunningAgentConfig(runningAgentConfigForIdentityUpdate(agentID, nextName, nextDisplayName))
		return
	}
	if err := rewriteAgentMemoryTitle(workDir, newEffectiveName); err != nil {
		slog.Warn("rewrite MEMORY.md title", "agent_id", agentID, "error", err)
	}
	if err := appendIdentityCheckpointNote(workDir, oldName, newEffectiveName); err != nil {
		slog.Warn("append identity checkpoint note", "agent_id", agentID, "error", err)
	}
	d.RefreshRunningAgentConfig(runningAgentConfigForIdentityUpdate(agentID, nextName, nextDisplayName))
	// Inline LLM configuration uses the agent name as the provider label passed
	// on the runtime command line. Updating the session snapshot cannot mutate
	// an already-started child process, so retire it and let the next delivery
	// launch with the authoritative identity.
	if identityChanged && d.interruptIdentityStaleSession(agentID) {
		return
	}

	needsNotice := d.agentNeedsIdentityNotice(agentID, currentCfg)
	if needsNotice && identityChanged {
		if err := identitynotice.Save(workDir, identitynotice.Pending{
			CurrentName: newEffectiveName,
			OldName:     oldName,
		}); err != nil {
			slog.Warn("persist identity notice", "agent_id", agentID, "error", err)
		}
		return
	}
	if err := identitynotice.Clear(workDir); err != nil {
		slog.Warn("clear identity notice", "agent_id", agentID, "error", err)
	}
}

func normalizeAgentIdentityNames(name, displayName string) (string, string) {
	normalizedName := strings.TrimSpace(name)
	normalizedDisplayName := strings.TrimSpace(displayName)
	if normalizedDisplayName == "" {
		normalizedDisplayName = normalizedName
	}
	if normalizedName == "" {
		normalizedName = normalizedDisplayName
	}
	return normalizedName, normalizedDisplayName
}

func (d *Dispatcher) interruptIdentityStaleSession(agentID string) bool {
	if d == nil {
		return false
	}
	targets := d.sessionManager.Each(func(sess *interactiveSession) bool {
		if sess == nil {
			return false
		}
		sess.mu.Lock()
		defer sess.mu.Unlock()
		return sess.process != nil && (strings.TrimSpace(sess.agentID) == agentID ||
			strings.TrimSpace(sess.serverAgentID) == agentID)
	})
	interrupted := false
	for _, sess := range targets {
		interrupted = d.interruptSessionForRestart(sess) || interrupted
	}
	return interrupted
}

func persistAgentIdentityStartConfig(workDir, agentID, name, displayName string) (*types.AgentStartConfig, string, bool, error) {
	var (
		currentCfg *types.AgentStartConfig
		oldName    string
		changed    bool
	)
	err := updateSessionResumeState(workDir, func(state *sessionResumeState) error {
		currentCfg = cloneAgentStartConfig(state.LastStartConfig)
		oldName = effectiveIdentityName(currentCfg, agentID)
		if currentCfg == nil {
			currentCfg = &types.AgentStartConfig{}
		}
		nextName := strings.TrimSpace(name)
		nextDisplayName := strings.TrimSpace(displayName)
		changed = strings.TrimSpace(currentCfg.Name) != nextName ||
			strings.TrimSpace(currentCfg.DisplayName) != nextDisplayName
		if !changed {
			return errAgentIdentityUnchanged
		}
		currentCfg.Name = nextName
		currentCfg.DisplayName = nextDisplayName
		state.LastStartConfig = cloneAgentStartConfig(currentCfg)
		return nil
	})
	if errors.Is(err, errAgentIdentityUnchanged) {
		err = nil
	}
	return currentCfg, oldName, changed, err
}

func runningAgentConfigForIdentityUpdate(agentID, name, displayName string) agent.Config {
	if cfg, ok := profile.Default().Agents().Get(agentID); ok {
		if trimmed := strings.TrimSpace(displayName); trimmed != "" {
			cfg.Name = trimmed
		} else if trimmed := strings.TrimSpace(name); trimmed != "" {
			cfg.Name = trimmed
		}
		return cfg
	}
	nextName := strings.TrimSpace(displayName)
	if nextName == "" {
		nextName = strings.TrimSpace(name)
	}
	return agent.Config{
		ID:   strings.TrimSpace(agentID),
		Name: nextName,
	}
}

// agentNeedsIdentityNotice returns true if the agent is currently running
// or has a saved session, meaning it should receive an identity notice on
// the next prompt.
func (d *Dispatcher) agentNeedsIdentityNotice(serverAgentID string, cfg *types.AgentStartConfig) bool {
	if _, ok := d.findProcess(serverAgentID); ok {
		return true
	}
	if cfg != nil && strings.TrimSpace(cfg.SessionID) != "" {
		return true
	}
	return false
}

func effectiveIdentityName(cfg *types.AgentStartConfig, agentID string) string {
	if cfg != nil {
		if trimmed := strings.TrimSpace(cfg.DisplayName); trimmed != "" {
			return trimmed
		}
		if trimmed := strings.TrimSpace(cfg.Name); trimmed != "" {
			return trimmed
		}
	}
	return strings.TrimSpace(agentID)
}

func rewriteAgentMemoryTitle(workDir string, title string) error {
	title = strings.TrimSpace(title)
	if title == "" {
		return nil
	}
	title = workspace.BoundedMemoryTitle(title)

	memPath := filepath.Join(workDir, "MEMORY.md")
	content, err := os.ReadFile(memPath)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return err
	}

	updated := replaceFirstMarkdownH1(string(content), title)
	if updated == string(content) {
		return nil
	}
	return os.WriteFile(memPath, []byte(updated), 0644)
}

func replaceFirstMarkdownH1(content string, title string) string {
	lines := strings.Split(content, "\n")
	for i, line := range lines {
		trimmed := strings.TrimSpace(line)
		if trimmed == "" {
			continue
		}
		if strings.HasPrefix(trimmed, "# ") {
			lines[i] = "# " + title
			return strings.Join(lines, "\n")
		}
		return content
	}
	return content
}

func appendIdentityCheckpointNote(workDir string, oldName string, newName string) error {
	oldName = strings.TrimSpace(oldName)
	newName = strings.TrimSpace(newName)
	summary := fmt.Sprintf("Agent identity synced to %s.", newName)
	if oldName != "" && oldName != newName {
		summary = fmt.Sprintf("Agent identity synced from %s to %s.", oldName, newName)
	}
	return workspace.AppendArchiveCheckpoint(workDir, "identity", summary+" Pending identity notice will remain until a fresh start clears it.", time.Now().UTC())
}