package runtimeproc

import (
	"context"
	"log/slog"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	DefaultScanInterval = 5 * time.Minute
	DefaultGracePeriod  = 2 * time.Minute
)

// ActiveWorkDirsFunc returns dispatcher-known live runtime workdirs.
// Matching entries are always protected from cleanup.
type ActiveWorkDirsFunc func() map[string]struct{}

// Janitor scans for runtime process trees that AgentTeam likely owns but no active
// dispatcher session still tracks. It requires ownership evidence before killing
// so user-launched processes with similar binary names are not reaped.
type Janitor struct {
	registry       *Registry
	manager        Manager
	agentRootDirs  []string
	activeWorkDirs ActiveWorkDirsFunc
	now            func() time.Time
	scanInterval   time.Duration
	gracePeriod    time.Duration

	mu         sync.Mutex
	suspicious map[string]time.Time
}

// JanitorOptions injects process enumeration, persisted registry state, active
// dispatcher state, and clocks so production code and tests use the same policy.
type JanitorOptions struct {
	Registry       *Registry
	Manager        Manager
	AgentRootDirs  []string
	ActiveWorkDirs ActiveWorkDirsFunc
	ScanInterval   time.Duration
	GracePeriod    time.Duration
	Now            func() time.Time
}

// ScanResult summarizes one sweep for tests and observability.
type ScanResult struct {
	ScannedProcesses int
	Suspicious       int
	Killed           []KilledProcess
}

// KilledProcess describes one process tree terminated by a Janitor scan and
// the ownership evidence that made the kill safe.
type KilledProcess struct {
	PID     int
	WorkDir string
	Reason  string
}

// NewJanitor creates a conservative runtime process janitor. Nil dependencies
// are filled with platform defaults, and invalid intervals/paths are normalized
// so production and tests exercise the same ownership policy.
func NewJanitor(opts JanitorOptions) *Janitor {
	manager := opts.Manager
	if manager == nil {
		manager = NewPlatformManager()
	}
	now := opts.Now
	if now == nil {
		now = time.Now
	}
	interval := opts.ScanInterval
	if interval <= 0 {
		interval = DefaultScanInterval
	}
	grace := opts.GracePeriod
	if grace <= 0 {
		grace = DefaultGracePeriod
	}
	roots := make([]string, 0, len(opts.AgentRootDirs))
	for _, root := range opts.AgentRootDirs {
		if cleaned := cleanPath(strings.TrimSpace(root)); cleaned != "" {
			roots = append(roots, cleaned)
		}
	}
	return &Janitor{
		registry:       opts.Registry,
		manager:        manager,
		agentRootDirs:  roots,
		activeWorkDirs: opts.ActiveWorkDirs,
		now:            now,
		scanInterval:   interval,
		gracePeriod:    grace,
		suspicious:     make(map[string]time.Time),
	}
}

// Start launches the periodic orphan-runtime scan agentteam until ctx is cancelled.
// Use ScanOnce directly in tests or callers that need synchronous evidence.
func (j *Janitor) Start(ctx context.Context) {
	if j == nil || j.manager == nil {
		return
	}
	go j.agentteam(ctx)
}

func (j *Janitor) agentteam(ctx context.Context) {
	// Delay the first periodic sweep so startup and runtime detection can settle.
	timer := time.NewTimer(j.scanInterval)
	defer timer.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-timer.C:
			if result, err := j.ScanOnce(ctx); err != nil {
				slog.Warn("runtime janitor scan failed", "error", err)
			} else if len(result.Killed) > 0 || result.Suspicious > 0 {
				slog.Info("runtime janitor scan complete", "scanned", result.ScannedProcesses, "suspicious", result.Suspicious, "killed", len(result.Killed))
			}
			timer.Reset(j.scanInterval)
		}
	}
}

// ScanOnce performs one conservative sweep. Candidates that lose ownership
// proof, are active again, or only match by binary name are forgotten rather than
// killed.
func (j *Janitor) ScanOnce(ctx context.Context) (ScanResult, error) {
	if j == nil || j.manager == nil {
		return ScanResult{}, nil
	}
	processes, err := j.manager.ListProcesses(ctx)
	if err != nil {
		return ScanResult{}, err
	}
	byPID := make(map[int]ProcessInfo, len(processes))
	for _, proc := range processes {
		if proc.PID > 0 {
			proc.CWD = cleanPath(proc.CWD)
			byPID[proc.PID] = proc
		}
	}
	active := normalizeDirSet(nil)
	if j.activeWorkDirs != nil {
		active = normalizeDirSet(j.activeWorkDirs())
	}
	j.pruneSuspicious(byPID)
	entries := entriesByPID(j.registry)
	discoveredWorkDirs := j.discoverAgentWorkDirs()

	result := ScanResult{ScannedProcesses: len(processes)}
	for _, proc := range processes {
		if proc.PID <= 0 || proc.PID == os.Getpid() {
			continue
		}
		workDir, reason, owned := j.ownedWorkDir(proc, entries, byPID, discoveredWorkDirs)
		if !owned || workDir == "" {
			j.forget(proc.PID)
			continue
		}
		if _, ok := active[workDir]; ok {
			j.forget(proc.PID)
			continue
		}
		if !runtimeLikeProcess(proc) && !strings.HasPrefix(reason, "registry") {
			j.forget(proc.PID)
			continue
		}
		result.Suspicious++
		if !j.readyToKill(proc.PID, workDir) {
			continue
		}
		if err := j.manager.KillTree(ctx, proc.PID); err != nil {
			slog.Warn("runtime janitor failed to kill process tree", "pid", proc.PID, "work_dir", workDir, "reason", reason, "error", err)
			continue
		}
		if j.registry != nil {
			j.registry.UnregisterPID(proc.PID)
		}
		j.forget(proc.PID)
		result.Killed = append(result.Killed, KilledProcess{PID: proc.PID, WorkDir: workDir, Reason: reason})
		slog.Warn("runtime janitor killed orphaned runtime process tree", "pid", proc.PID, "work_dir", workDir, "reason", reason)
	}
	sort.Slice(result.Killed, func(i, k int) bool { return result.Killed[i].PID < result.Killed[k].PID })
	return result, nil
}

func entriesByPID(registry *Registry) map[int]Entry {
	entries := make(map[int]Entry)
	if registry == nil {
		return entries
	}
	for _, entry := range registry.Entries() {
		entries[entry.PID] = entry
	}
	return entries
}

// ownedWorkDir returns the AgentTeam agent workdir that proves ownership for proc.
// Registry evidence is accepted only after the current process still looks like
// the registered runtime root, which avoids recycled-PID false positives.
func (j *Janitor) ownedWorkDir(proc ProcessInfo, entries map[int]Entry, byPID map[int]ProcessInfo, discovered []string) (string, string, bool) {
	if entry, ok := entries[proc.PID]; ok && entry.WorkDir != "" && j.registryRootMatchesProcess(proc, entry) {
		return cleanPath(entry.WorkDir), "registry", true
	}
	if entry, ok := j.registryAncestorEntry(proc, entries, byPID); ok && entry.WorkDir != "" {
		return cleanPath(entry.WorkDir), "registry_descendant", true
	}
	if proc.CWD != "" {
		if workDir := j.agentWorkDirForPath(proc.CWD); workDir != "" {
			return workDir, "cwd", true
		}
	}
	cmdline := strings.ToLower(filepath.ToSlash(proc.Cmdline))
	for _, workDir := range discovered {
		if workDir == "" {
			continue
		}
		if strings.Contains(cmdline, strings.ToLower(filepath.ToSlash(workDir))) {
			return workDir, "cmdline", true
		}
	}
	return "", "", false
}

// registryAncestorEntry treats children of a verified registry root as owned.
// The ancestor itself must pass registryRootMatchesProcess first.
func (j *Janitor) registryAncestorEntry(proc ProcessInfo, entries map[int]Entry, byPID map[int]ProcessInfo) (Entry, bool) {
	seen := map[int]struct{}{}
	parentPID := proc.PPID
	for parentPID > 0 {
		parent, ok := byPID[parentPID]
		if !ok {
			break
		}
		if entry, ok := entries[parentPID]; ok && entry.WorkDir != "" && j.registryRootMatchesProcess(parent, entry) {
			return entry, true
		}
		if _, ok := seen[parentPID]; ok {
			break
		}
		seen[parentPID] = struct{}{}
		parentPID = parent.PPID
	}
	return Entry{}, false
}

// registryRootMatchesProcess guards registry PID matches against PID reuse by
// requiring runtime identity plus cwd/cmdline workdir evidence or a close process
// creation timestamp.
func (j *Janitor) registryRootMatchesProcess(proc ProcessInfo, entry Entry) bool {
	if !runtimeLikeProcess(proc) {
		return false
	}
	workDir := cleanPath(entry.WorkDir)
	if proc.CWD != "" {
		return j.agentWorkDirForPath(proc.CWD) == workDir
	}
	cmdline := strings.ToLower(filepath.ToSlash(proc.Cmdline))
	if workDir != "" && strings.Contains(cmdline, strings.ToLower(filepath.ToSlash(workDir))) {
		return true
	}
	if !entry.StartedAt.IsZero() && !proc.StartedAt.IsZero() {
		return sameProcessStart(entry.StartedAt, proc.StartedAt)
	}
	// Without cwd, cmdline workdir, or creation time, a registry PID match is
	// only a stale hint. Refuse to kill rather than risk a recycled-PID false
	// positive.
	return false
}

// sameProcessStart allows small clock/precision skew between persisted runtime
// start time and platform process creation time.
func sameProcessStart(registered time.Time, observed time.Time) bool {
	registered = registered.UTC()
	observed = observed.UTC()
	delta := observed.Sub(registered)
	if delta < 0 {
		delta = -delta
	}
	return delta <= 2*time.Minute
}

func (j *Janitor) discoverAgentWorkDirs() []string {
	var dirs []string
	for _, root := range j.agentRootDirs {
		entries, err := os.ReadDir(root)
		if err != nil {
			continue
		}
		for _, entry := range entries {
			if !entry.IsDir() {
				continue
			}
			dirs = append(dirs, cleanPath(filepath.Join(root, entry.Name())))
		}
	}
	sort.Strings(dirs)
	return dirs
}

func (j *Janitor) agentWorkDirForPath(path string) string {
	path = cleanPath(path)
	for _, root := range j.agentRootDirs {
		rel, err := filepath.Rel(root, path)
		if err != nil || rel == "." || strings.HasPrefix(rel, "..") {
			continue
		}
		parts := strings.Split(filepath.Clean(rel), string(filepath.Separator))
		if len(parts) == 0 || parts[0] == "." || parts[0] == "" {
			continue
		}
		if len(parts) >= 3 && parts[1] == "sessions" && parts[2] != "" {
			return cleanPath(filepath.Join(root, parts[0], parts[1], parts[2]))
		}
		return cleanPath(filepath.Join(root, parts[0]))
	}
	return ""
}

// readyToKill requires a candidate to survive one grace window, preventing
// cleanup while dispatcher state and registry writes are still settling.
func (j *Janitor) readyToKill(pid int, workDir string) bool {
	key := suspicionKey(pid, workDir)
	now := j.now()
	j.mu.Lock()
	defer j.mu.Unlock()
	firstSeen, ok := j.suspicious[key]
	if !ok {
		j.suspicious[key] = now
		return false
	}
	return !now.Before(firstSeen.Add(j.gracePeriod))
}

func (j *Janitor) forget(pid int) {
	j.mu.Lock()
	defer j.mu.Unlock()
	prefix := suspicionKey(pid, "")
	for key := range j.suspicious {
		if strings.HasPrefix(key, prefix) {
			delete(j.suspicious, key)
		}
	}
}

// pruneSuspicious prevents the suspicion cache from retaining exited PIDs
// forever, which also reduces risk if an old PID is later reused.
func (j *Janitor) pruneSuspicious(live map[int]ProcessInfo) {
	j.mu.Lock()
	defer j.mu.Unlock()
	for key := range j.suspicious {
		pidText, _, _ := strings.Cut(key, "|")
		pid, err := strconv.Atoi(pidText)
		if err != nil {
			delete(j.suspicious, key)
			continue
		}
		if _, ok := live[pid]; !ok {
			delete(j.suspicious, key)
		}
	}
}

func suspicionKey(pid int, workDir string) string {
	return strings.Join([]string{strconvItoa(pid), cleanPath(workDir)}, "|")
}

func normalizeDirSet(input map[string]struct{}) map[string]struct{} {
	out := make(map[string]struct{}, len(input))
	for dir := range input {
		if cleaned := cleanPath(dir); cleaned != "" {
			out[cleaned] = struct{}{}
		}
	}
	return out
}

// runtimeLikeProcess is a secondary filter after ownership proof. It includes
// common shim hosts because OpenCode/Codex/Claude may run through node/bun/npm.
func runtimeLikeProcess(proc ProcessInfo) bool {
	text := strings.ToLower(proc.Exe + "\x00" + proc.Cmdline)
	for _, token := range []string{"opencode", "claude", "codex", "kimi"} {
		if strings.Contains(text, token) {
			return true
		}
	}
	base := strings.ToLower(filepath.Base(proc.Exe))
	switch base {
	case "node", "node.exe", "bun", "bun.exe", "npm", "npm.cmd", "npx", "npx.cmd":
		return true
	}
	if runtime.GOOS == "windows" && (base == "cmd.exe" || base == "powershell.exe" || base == "pwsh.exe") {
		return strings.Contains(text, "opencode") || strings.Contains(text, "claude") || strings.Contains(text, "codex") || strings.Contains(text, "kimi")
	}
	return false
}

func strconvItoa(v int) string {
	return strconv.FormatInt(int64(v), 10)
}