package runtimeproc

import (
	"encoding/json"
	"errors"
	"log/slog"
	"os"
	"path/filepath"
	"sort"
	"sync"
	"time"
)

// Entry records a runtime process that AgentTeam started. It is persisted so a
// restarted daemon can clean up runtime trees left behind by a previous crash.
type Entry struct {
	PID       int       `json:"pid"`
	AgentID   string    `json:"agentId"`
	Runtime   string    `json:"runtime"`
	WorkDir   string    `json:"workDir"`
	StartedAt time.Time `json:"startedAt"`
}

// Registry persists runtime root process hints by PID. The janitor still
// verifies active dispatcher state and current process identity before using a
// registry entry to terminate anything.
type Registry struct {
	mu      sync.Mutex
	saveMu  sync.Mutex
	path    string
	entries map[int]Entry
}

// NewRegistry loads a process registry from path. A missing or malformed file
// starts with an empty registry; janitor safety does not depend on registry
// contents alone.
func NewRegistry(path string) *Registry {
	r := &Registry{path: path, entries: make(map[int]Entry)}
	if err := r.load(); err != nil {
		// Registry state is only a safety hint, so load failures should not
		// prevent startup; keep enough telemetry to diagnose lost cleanup hints.
		slog.Warn("runtime process registry load failed; starting empty", "path", path, "error", err)
	}
	return r
}

// Register records or replaces the root process entry for one AgentTeam-started
// runtime tree and persists the registry best-effort.
func (r *Registry) Register(entry Entry) {
	persist := r.StageRegister(entry)
	if persist != nil {
		persist()
	}
}

// StageRegister records an entry in memory and returns the best-effort durable
// persistence step. This split lets a lifecycle owner make the liveness
// mutation under its own short critical section without holding that admission
// lock across filesystem I/O. The returned function is idempotent and may be
// called after a concurrent unregister; save always snapshots the latest map.
func (r *Registry) StageRegister(entry Entry) func() {
	if r == nil || entry.PID <= 0 {
		return nil
	}
	entry.WorkDir = cleanPath(entry.WorkDir)
	r.mu.Lock()
	r.entries[entry.PID] = entry
	r.mu.Unlock()
	var once sync.Once
	return func() {
		once.Do(func() {
			if err := r.save(); err != nil {
				slog.Warn("runtime process registry save failed after register",
					"path", r.path,
					"pid", entry.PID,
					"agent_id", entry.AgentID,
					"runtime", entry.Runtime,
					"error", err,
				)
			}
		})
	}
}

// UnregisterPID removes a runtime root PID from the persisted registry.
func (r *Registry) UnregisterPID(pid int) {
	if r == nil || pid <= 0 {
		return
	}
	r.mu.Lock()
	delete(r.entries, pid)
	r.mu.Unlock()
	if err := r.save(); err != nil {
		slog.Warn("runtime process registry save failed after unregister",
			"path", r.path,
			"pid", pid,
			"error", err,
		)
	}
}

// Entries returns a PID-sorted snapshot of the registered runtime roots.
func (r *Registry) Entries() []Entry {
	if r == nil {
		return nil
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	entries := make([]Entry, 0, len(r.entries))
	for _, entry := range r.entries {
		entries = append(entries, entry)
	}
	sort.Slice(entries, func(i, j int) bool { return entries[i].PID < entries[j].PID })
	return entries
}

// EntryByPID returns the registry entry for pid, if present.
func (r *Registry) EntryByPID(pid int) (Entry, bool) {
	if r == nil || pid <= 0 {
		return Entry{}, false
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	entry, ok := r.entries[pid]
	return entry, ok
}

func (r *Registry) load() error {
	if r == nil || r.path == "" {
		return nil
	}
	data, err := os.ReadFile(r.path)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return nil
		}
		return err
	}
	var entries []Entry
	if err := json.Unmarshal(data, &entries); err != nil {
		return err
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	for _, entry := range entries {
		if entry.PID <= 0 || entry.WorkDir == "" {
			continue
		}
		entry.WorkDir = cleanPath(entry.WorkDir)
		r.entries[entry.PID] = entry
	}
	return nil
}

func (r *Registry) save() error {
	if r == nil || r.path == "" {
		return nil
	}
	// Serialize snapshot+write as one operation. Without saveMu, an older
	// register snapshot can finish writing after a newer unregister snapshot and
	// resurrect a stale PID in the persisted file even though entries is empty.
	r.saveMu.Lock()
	defer r.saveMu.Unlock()
	r.mu.Lock()
	entries := make([]Entry, 0, len(r.entries))
	for _, entry := range r.entries {
		entries = append(entries, entry)
	}
	r.mu.Unlock()
	sort.Slice(entries, func(i, j int) bool { return entries[i].PID < entries[j].PID })
	data, err := json.MarshalIndent(entries, "", "  ")
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(r.path), 0o755); err != nil {
		return err
	}
	return os.WriteFile(r.path, append(data, '\n'), 0o644)
}

func cleanPath(path string) string {
	if path == "" {
		return ""
	}
	if abs, err := filepath.Abs(path); err == nil {
		path = abs
	}
	if eval, err := filepath.EvalSymlinks(path); err == nil {
		path = eval
	}
	return filepath.Clean(path)
}