package tools

import (
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

// WorktreeEntry describes a single active worktree managed by xbot.
type WorktreeEntry struct {
	SessionKey  string // "cli:/path/to/repo:debug" or "agent:role/instance"
	Role        string // "primary" | "peer" | "child"
	RepoPath    string // absolute path to the main git repo
	WorktreeDir string // absolute path to the worktree (empty for primary)
	Branch      string // branch name
	CreatedAt   time.Time
	Status      string // "working" | "merge-ready" | "done"
}

// WorktreeRegistry is a process-level registry of active worktrees.
// It is the single source of truth for peer discovery and is shared
// between WorktreeTool (writer) and DynamicContextInjector (reader).
type WorktreeRegistry struct {
	mu       sync.RWMutex
	byRepo   map[string][]*WorktreeEntry // repoPath → entries
	bySess   map[string]*WorktreeEntry   // sessionKey → entry
	loaded   map[string]bool             // repoPath → whether persisted data has been loaded
	loadedMu sync.Mutex                  // protects loaded
}

// GlobalWorktreeRegistry is the singleton registry used by all components.
var GlobalWorktreeRegistry = &WorktreeRegistry{
	byRepo: make(map[string][]*WorktreeEntry),
	bySess: make(map[string]*WorktreeEntry),
	loaded: make(map[string]bool),
}

// ensureLoaded lazily loads persisted registry data for a repo.
func (r *WorktreeRegistry) ensureLoaded(repoPath string) {
	r.loadedMu.Lock()
	if r.loaded[repoPath] {
		r.loadedMu.Unlock()
		return
	}
	r.loaded[repoPath] = true
	r.loadedMu.Unlock()

	r.mu.Lock()
	r.loadRepoLocked(repoPath)
	r.mu.Unlock()
}

// Register adds an entry to the registry. Returns error if sessionKey already exists.
func (r *WorktreeRegistry) Register(entry *WorktreeEntry) error {
	r.mu.Lock()
	defer r.mu.Unlock()

	r.loadRepoLocked(entry.RepoPath)

	if _, exists := r.bySess[entry.SessionKey]; exists {
		return fmt.Errorf("worktree: session %q already registered", entry.SessionKey)
	}

	r.bySess[entry.SessionKey] = entry
	r.byRepo[entry.RepoPath] = append(r.byRepo[entry.RepoPath], entry)
	r.saveRepoLocked(entry.RepoPath)
	return nil
}

// Deregister removes an entry and cleans up empty repo buckets.
// For entries with physical worktrees, use CleanupSession instead which also
// removes the git worktree and branch.
func (r *WorktreeRegistry) Deregister(sessionKey string) {
	r.mu.Lock()
	defer r.mu.Unlock()

	entry, exists := r.bySess[sessionKey]
	if !exists {
		return
	}
	delete(r.bySess, sessionKey)

	entries := r.byRepo[entry.RepoPath]
	for i, e := range entries {
		if e.SessionKey == sessionKey {
			r.byRepo[entry.RepoPath] = append(entries[:i], entries[i+1:]...)
			break
		}
	}
	if len(r.byRepo[entry.RepoPath]) == 0 {
		delete(r.byRepo, entry.RepoPath)
	}
	r.saveRepoLocked(entry.RepoPath)
}

// CleanupSession removes a session's worktree registration and, if the session
// has a physical worktree (WorktreeDir != ""), deletes the git worktree and branch.
// This is the correct method to call when a session is deleted/closed.
func (r *WorktreeRegistry) CleanupSession(sessionKey string) {
	entry := r.GetBySession(sessionKey)
	if entry == nil {
		return
	}
	// Remove physical worktree + branch if present
	if entry.WorktreeDir != "" && entry.RepoPath != "" && entry.Branch != "" {
		_ = removeWorktree(entry.RepoPath, entry.WorktreeDir, entry.Branch)
	}
	r.Deregister(sessionKey)
}

// GetPeers returns all entries for a repo, excluding the given sessionKey.
func (r *WorktreeRegistry) GetPeers(repoPath, excludeSessionKey string) []*WorktreeEntry {
	r.ensureLoaded(repoPath)
	r.mu.RLock()
	defer r.mu.RUnlock()

	var peers []*WorktreeEntry
	for _, e := range r.byRepo[repoPath] {
		if e.SessionKey != excludeSessionKey {
			peers = append(peers, cloneEntry(e))
		}
	}
	return peers
}

// GetBySession returns the entry for a given session, or nil.
func (r *WorktreeRegistry) GetBySession(sessionKey string) *WorktreeEntry {
	r.mu.RLock()
	defer r.mu.RUnlock()
	e, ok := r.bySess[sessionKey]
	if !ok {
		return nil
	}
	return cloneEntry(e)
}

// GetPrimary returns the primary entry for a repo, or nil.
func (r *WorktreeRegistry) GetPrimary(repoPath string) *WorktreeEntry {
	r.ensureLoaded(repoPath)
	r.mu.RLock()
	defer r.mu.RUnlock()
	for _, e := range r.byRepo[repoPath] {
		if e.Role == "primary" {
			return cloneEntry(e)
		}
	}
	return nil
}

// HasPeers returns true if there are other active entries in the repo.
func (r *WorktreeRegistry) HasPeers(repoPath, excludeSessionKey string) bool {
	r.ensureLoaded(repoPath)
	r.mu.RLock()
	defer r.mu.RUnlock()
	for _, e := range r.byRepo[repoPath] {
		if e.SessionKey != excludeSessionKey {
			return true
		}
	}
	return false
}

// ListRepo returns all entries for a repo (including the caller).
func (r *WorktreeRegistry) ListRepo(repoPath string) []*WorktreeEntry {
	r.ensureLoaded(repoPath)
	r.mu.RLock()
	defer r.mu.RUnlock()
	result := make([]*WorktreeEntry, len(r.byRepo[repoPath]))
	for i, e := range r.byRepo[repoPath] {
		result[i] = cloneEntry(e)
	}
	return result
}

// UpdateStatus updates the status of an entry.
func (r *WorktreeRegistry) UpdateStatus(sessionKey, status string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if e, ok := r.bySess[sessionKey]; ok {
		e.Status = status
		r.saveRepoLocked(e.RepoPath)
	}
}

// UpdateCWD updates the WorktreeDir for a session (called from SetCurrentDir).
// If the session is registered, updates its WorktreeDir and persists.
func (r *WorktreeRegistry) UpdateCWD(sessionKey, dir string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	e, ok := r.bySess[sessionKey]
	if !ok {
		return // not a worktree session, nothing to persist
	}
	e.WorktreeDir = dir
	r.saveRepoLocked(e.RepoPath)
}

// GetCWD returns the persisted WorktreeDir for a session, or "".
func (r *WorktreeRegistry) GetCWD(sessionKey string) string {
	r.ensureLoadedBySession(sessionKey)
	r.mu.RLock()
	defer r.mu.RUnlock()
	if e, ok := r.bySess[sessionKey]; ok && e.WorktreeDir != "" {
		return e.WorktreeDir
	}
	return ""
}

// RegisterPeer registers a session for peer awareness without creating a worktree.
// Used when auto_worktree is disabled. The first session in a repo is registered
// as "primary"; subsequent sessions are registered as "peer" (sharing the main
// workspace, no file isolation).
//
// Entries created by RegisterPeer are NOT persisted to disk — they are runtime-only
// peer awareness data that becomes stale across process restarts.
func (r *WorktreeRegistry) RegisterPeer(sessionKey, workDir string) {
	if r.GetBySession(sessionKey) != nil {
		return // already registered
	}
	repoPath, err := GitRepoRoot(workDir)
	if err != nil {
		return // not a git repo
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	r.loadRepoLocked(repoPath)
	if _, exists := r.bySess[sessionKey]; exists {
		return
	}
	// Determine role: first session → primary, others → peer.
	// Must check inline (not via GetPrimary) because we already hold mu.Lock.
	role := "primary"
	for _, e := range r.byRepo[repoPath] {
		if e.Role == "primary" {
			role = "peer"
			break
		}
	}
	entry := &WorktreeEntry{
		SessionKey:  sessionKey,
		Role:        role,
		RepoPath:    repoPath,
		WorktreeDir: "",
		Branch:      "",
		Status:      "working",
	}
	r.bySess[sessionKey] = entry
	r.byRepo[repoPath] = append(r.byRepo[repoPath], entry)
	// Do NOT saveRepoLocked: peer-awareness entries are runtime-only.
	// Only entries with real worktrees (WorktreeDir != "") need persistence.
}

// ensureLoadedBySession tries to load persisted data for the repo that
// contains this session. Uses the sessionKey format "channel:chatID" to
// derive a possible repo path.
func (r *WorktreeRegistry) ensureLoadedBySession(sessionKey string) {
	// Try to find the repo path from already-loaded entries or the session key.
	r.mu.RLock()
	e, ok := r.bySess[sessionKey]
	r.mu.RUnlock()
	if ok {
		r.ensureLoaded(e.RepoPath)
		return
	}
	// Not loaded yet — will load on first Register/GetPrimary/etc.
}

func cloneEntry(e *WorktreeEntry) *WorktreeEntry {
	c := *e
	return &c
}

// --- Persistence ---

type registryFile struct {
	Entries []*WorktreeEntry `json:"entries"`
}

func registryPath(repoPath string) string {
	return filepath.Join(worktreeBaseDir(repoPath), "registry.json")
}

// loadRepo loads persisted entries for a repo from disk. Caller must hold r.mu.
func (r *WorktreeRegistry) loadRepoLocked(repoPath string) {
	path := registryPath(repoPath)
	data, err := os.ReadFile(path)
	if err != nil {
		return
	}
	var rf registryFile
	if err := json.Unmarshal(data, &rf); err != nil {
		return
	}
	if len(rf.Entries) == 0 {
		return
	}

	for _, e := range rf.Entries {
		if _, exists := r.bySess[e.SessionKey]; exists {
			continue
		}
		if e.WorktreeDir != "" {
			if _, err := os.Stat(e.WorktreeDir); os.IsNotExist(err) {
				continue // orphaned worktree dir gone
			}
		} else {
			// Entries without worktrees are runtime-only peer awareness data
			// from a previous process. Skip them — they become stale on restart.
			continue
		}
		r.bySess[e.SessionKey] = e
		r.byRepo[e.RepoPath] = append(r.byRepo[e.RepoPath], e)
	}
}

// saveRepoLocked persists entries for a repo to disk. Caller must hold r.mu.
func (r *WorktreeRegistry) saveRepoLocked(repoPath string) {
	entries := r.byRepo[repoPath]
	if len(entries) == 0 {
		return
	}

	rf := registryFile{Entries: entries}
	data, err := json.MarshalIndent(rf, "", "  ")
	if err != nil {
		return
	}

	path := registryPath(repoPath)
	dir := filepath.Dir(path)
	os.MkdirAll(dir, 0755)

	tmpPath := path + ".tmp"
	if err := os.WriteFile(tmpPath, data, 0644); err != nil {
		return
	}
	os.Rename(tmpPath, path)
}

// --- Worktree helper functions ---

// gitEnvBlocklist contains GIT_* variables that must be stripped from
// subprocess environments so that git commands operate on the intended
// directory instead of the parent repo (e.g. during pre-commit hooks).
var gitEnvBlocklist = []string{
	"GIT_DIR",
	"GIT_WORK_TREE",
	"GIT_INDEX_FILE",
	"GIT_OBJECT_DIRECTORY",
	"GIT_ALTERNATE_OBJECT_DIRECTORIES",
}

// cleanGitEnv returns os.Environ() with git plumbing variables removed.
func cleanGitEnv() []string {
	var env []string
	for _, e := range os.Environ() {
		strip := false
		for _, prefix := range gitEnvBlocklist {
			if strings.HasPrefix(e, prefix+"=") {
				strip = true
				break
			}
		}
		if !strip {
			env = append(env, e)
		}
	}
	return env
}

// GitRepoRoot returns the absolute root of the git repo containing dir.
// Works correctly in both regular repos and git worktrees (uses git-common-dir).
func GitRepoRoot(dir string) (string, error) {
	// Check if this is a git worktree by reading the .git file
	gitFile := filepath.Join(dir, ".git")
	content, err := os.ReadFile(gitFile)
	if err == nil && strings.HasPrefix(string(content), "gitdir:") {
		// It's a worktree — the .git file points to the main repo's metadata.
		// Example: "gitdir: /main/repo/.git/worktrees/name\n"
		gitDirLine := strings.TrimPrefix(string(content), "gitdir:")
		gitDirLine = strings.TrimSpace(gitDirLine)
		// .git/worktrees/name → .git is two levels up
		worktreeMetaDir := filepath.Dir(gitDirLine) // → /main/repo/.git/worktrees
		gitDir := filepath.Dir(worktreeMetaDir)     // → /main/repo/.git
		return filepath.Dir(gitDir), nil            // → /main/repo
	}

	// Regular directory: use git rev-parse
	cmd := exec.Command("git", "-C", dir, "rev-parse", "--show-toplevel")
	cmd.Env = cleanGitEnv()
	out, err := cmd.Output()
	if err != nil {
		return "", fmt.Errorf("not a git repository: %w", err)
	}
	return strings.TrimSpace(string(out)), nil
}

// gitIsDirty returns true if the repo at dir has uncommitted changes.
func gitIsDirty(dir string) (bool, error) {
	cmd := exec.Command("git", "-C", dir, "status", "--porcelain")
	out, err := cmd.Output()
	if err != nil {
		return false, err
	}
	return len(strings.TrimSpace(string(out))) > 0, nil
}

// worktreeBaseDir returns the base directory for all xbot worktrees.
// Worktrees are placed outside the main repo: {repo}/../.xbot-worktrees/
func worktreeBaseDir(repoPath string) string {
	return filepath.Join(filepath.Dir(repoPath), ".xbot-worktrees")
}

// generateBranchName creates a unique branch name for an agent.
func generateBranchName(role, instance, taskHint string) string {
	sanitize := func(s string) string {
		s = strings.ToLower(s)
		s = strings.Map(func(r rune) rune {
			if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '-' {
				return r
			}
			return '-'
		}, s)
		return strings.Trim(s, "-")
	}

	task := sanitize(taskHint)
	if len(task) > 30 {
		task = task[:30]
	}
	if task == "" {
		task = time.Now().Format("20060102-150405")
	}

	return fmt.Sprintf("agent/%s/%s/%s", sanitize(role), sanitize(instance), task)
}

// resolveRemoteMainBranch fetches the latest refs from ALL remotes and returns
// the remote-tracking ref for the default branch of the most up-to-date remote
// (e.g. "origin/master"). When multiple remotes exist, picks the one whose
// default branch has the newest commit timestamp.
// Returns "" when the repo has no remote or detection fails — callers should
// fall back to using local HEAD.
func resolveRemoteMainBranch(repoPath string) string {
	// Step 1: list all remotes
	remoteCmd := exec.Command("git", "-C", repoPath, "remote")
	remoteCmd.Env = cleanGitEnv()
	out, err := remoteCmd.Output()
	if err != nil || len(strings.TrimSpace(string(out))) == 0 {
		return "" // no remote configured
	}
	remotes := strings.Fields(strings.TrimSpace(string(out)))

	// Step 2: fetch all remotes (best-effort, don't fail on network errors)
	for _, remote := range remotes {
		fetchCmd := exec.Command("git", "-C", repoPath, "fetch", remote)
		fetchCmd.Env = cleanGitEnv()
		_ = fetchCmd.Run() // ignore errors — use cached refs if offline
	}

	// Step 3: for each remote, detect its default branch and collect
	// (remote, branch, commitDate) tuples. Pick the one with the newest commit.
	// Among equal dates, prefer remote named "origin" as tiebreaker.
	type remoteCandidate struct {
		ref        string // e.g. "origin/master"
		remote     string // e.g. "origin"
		commitDate string // ISO date for comparison
	}
	var best remoteCandidate

	for _, remote := range remotes {
		branch := detectRemoteDefaultBranch(repoPath, remote)
		if branch == "" {
			continue
		}
		ref := remote + "/" + branch

		// Get committer date of this remote ref
		dateCmd := exec.Command("git", "-C", repoPath, "log", "-1", "--format=%cI", ref)
		dateCmd.Env = cleanGitEnv()
		dateOut, dateErr := dateCmd.Output()
		if dateErr != nil {
			continue
		}
		date := strings.TrimSpace(string(dateOut))

		pick := false
		if best.commitDate == "" {
			pick = true
		} else if date > best.commitDate {
			pick = true
		} else if date == best.commitDate && remote == "origin" && best.remote != "origin" {
			pick = true // tiebreaker: prefer "origin"
		}

		if pick {
			best = remoteCandidate{ref: ref, remote: remote, commitDate: date}
		}
	}

	return best.ref
}

// detectRemoteDefaultBranch resolves the default branch name for a single remote.
// Uses symbolic-ref first, then falls back to probing common names.
func detectRemoteDefaultBranch(repoPath, remote string) string {
	// Try symbolic-ref: refs/remotes/origin/HEAD → refs/remotes/origin/master
	symRef := "refs/remotes/" + remote + "/HEAD"
	symCmd := exec.Command("git", "-C", repoPath, "symbolic-ref", symRef)
	symCmd.Env = cleanGitEnv()
	symOut, err := symCmd.Output()
	if err == nil {
		ref := strings.TrimSpace(string(symOut))
		prefix := "refs/remotes/" + remote + "/"
		if branch := strings.TrimPrefix(ref, prefix); branch != ref && branch != "" {
			return branch
		}
	}

	// Fallback: try common branch names against cached remote refs
	for _, candidate := range []string{"main", "master"} {
		revCmd := exec.Command("git", "-C", repoPath, "rev-parse", "--verify",
			"refs/remotes/"+remote+"/"+candidate)
		revCmd.Env = cleanGitEnv()
		if err := revCmd.Run(); err == nil {
			return candidate
		}
	}

	return ""
}

// createWorktree creates a git worktree and returns its path and branch.
// The worktree is always based on the latest remote main branch (fetched
// from origin before creation). Falls back to local HEAD when no remote
// is configured.
// Caller must hold WorktreeRegistry.mu or ensure serialization externally.
func createWorktree(repoPath, branch string) (worktreePath string, err error) {
	baseDir := worktreeBaseDir(repoPath)
	dirName := strings.TrimPrefix(branch, "agent/")
	dirName = strings.ReplaceAll(dirName, "/", "-")
	worktreePath = filepath.Join(baseDir, dirName)

	// Ensure base directory exists
	if err := os.MkdirAll(baseDir, 0755); err != nil {
		return "", fmt.Errorf("create worktree base dir: %w", err)
	}

	// Resolve start point: prefer latest remote main branch over local HEAD.
	// This ensures every new worktree starts from the freshest upstream state
	// rather than a stale local checkout.
	startPoint := "HEAD"
	if remoteRef := resolveRemoteMainBranch(repoPath); remoteRef != "" {
		startPoint = remoteRef
	}

	// Step 1: git worktree add --detach <path> <startPoint>
	cmd := exec.Command("git", "-C", repoPath, "worktree", "add",
		"--detach", worktreePath, startPoint)
	cmd.Env = cleanGitEnv()
	out, err := cmd.CombinedOutput()
	if err != nil {
		return "", fmt.Errorf("git worktree add: %s: %w", strings.TrimSpace(string(out)), err)
	}

	// Step 2: create branch in the new worktree
	brCmd := exec.Command("git", "-C", worktreePath, "checkout", "-b", branch)
	brCmd.Env = cleanGitEnv()
	brOut, brErr := brCmd.CombinedOutput()
	if brErr != nil {
		// Rollback: remove the worktree we just created
		removeWorktree(repoPath, worktreePath, "")
		return "", fmt.Errorf("git checkout -b in worktree: %s: %w", strings.TrimSpace(string(brOut)), brErr)
	}

	return worktreePath, nil
}

// removeWorktree removes a git worktree and deletes its branch.
func removeWorktree(repoPath, worktreePath, branch string) error {
	// Remove worktree (--force handles dirty state)
	cmd := exec.Command("git", "-C", repoPath, "worktree", "remove",
		worktreePath, "--force")
	cmd.Env = cleanGitEnv()
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("git worktree remove: %s: %w", strings.TrimSpace(string(out)), err)
	}

	// Delete branch
	delCmd := exec.Command("git", "-C", repoPath, "branch", "-d", branch)
	delCmd.Env = cleanGitEnv()
	delOut, delErr := delCmd.CombinedOutput()
	if delErr != nil {
		// Non-fatal: branch may already be merged/deleted
		_ = delOut
	}

	return nil
}

// pruneOrphanWorktrees cleans up worktree metadata for directories that no longer exist.
// AutoDetectAndInit creates an isolated git worktree for the current session.
// Called automatically at session start when auto_worktree is enabled.
//
// Every session gets its own worktree — no primary concept. All agents are equal peers.
// Returns (entry, created): entry is non-nil on success; created is true only when a
// new worktree was physically created (false when returning an existing entry from disk).
func AutoDetectAndInit(workDir, sessionKey string) (*WorktreeEntry, bool) {
	return autoDetectAndInitInto(workDir, sessionKey, GlobalWorktreeRegistry)
}

// autoDetectAndInitInto is the testable core that accepts a custom registry.
func autoDetectAndInitInto(workDir, sessionKey string, reg *WorktreeRegistry) (*WorktreeEntry, bool) {
	// Check if in a git repo
	repoPath, err := GitRepoRoot(workDir)
	if err != nil {
		return nil, false // not a git repo
	}

	// Ensure persisted data is loaded so we detect existing sessions
	// from a previous process (restart recovery).
	reg.ensureLoaded(repoPath)

	// Already registered?
	if entry := reg.GetBySession(sessionKey); entry != nil {
		return entry, false
	}

	// All sessions get a worktree — no primary concept.
	// Use short session name (after last ":") to keep branch names readable.
	shortName := sessionKey
	if idx := strings.LastIndex(sessionKey, ":"); idx >= 0 {
		shortName = sessionKey[idx+1:]
	}
	// Further shorten: strip common prefixes like "Agent-"
	shortName = strings.TrimPrefix(shortName, "Agent-")
	branch := fmt.Sprintf("wt/%s/%s", shortName, time.Now().Format("20060102-150405"))
	branch = strings.ReplaceAll(branch, ":", "-")

	// Serialize worktree creation
	reg.mu.Lock()
	worktreePath, err := createWorktree(repoPath, branch)
	reg.mu.Unlock()
	if err != nil {
		return nil, false
	}

	entry := &WorktreeEntry{
		SessionKey:  sessionKey,
		Role:        "peer",
		RepoPath:    repoPath,
		WorktreeDir: worktreePath,
		Branch:      branch,
		Status:      "working",
	}
	if err := reg.Register(entry); err != nil {
		removeWorktree(repoPath, worktreePath, branch)
		// Safety net: Register loaded from disk and found existing entry.
		// Return it for idempotent behavior across restarts.
		if existing := reg.GetBySession(sessionKey); existing != nil {
			return existing, false
		}
		return nil, false
	}
	return entry, true
}
