package dispatcher

import (
	"context"
	"log/slog"
	"path/filepath"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/agent"
	"github.com/tidbcloud/AgentTeam/desktop/internal/agent/runtimeproc"
	"github.com/tidbcloud/AgentTeam/desktop/internal/profile"
)

const runtimeProcessRegistryFile = "runtime-processes.json"

func newRuntimeProcessRegistry() *runtimeproc.Registry {
	return runtimeproc.NewRegistry(filepath.Join(profile.ConfigDir(), runtimeProcessRegistryFile))
}

func newRuntimeProcessJanitor(d *Dispatcher, registry *runtimeproc.Registry) *runtimeproc.Janitor {
	return runtimeproc.NewJanitor(runtimeproc.JanitorOptions{
		Registry:       registry,
		AgentRootDirs:  []string{filepath.Join(profile.ConfigDir(), "agents")},
		ActiveWorkDirs: d.activeRuntimeWorkDirs,
	})
}

func (d *runtimeProcessService) startRuntimeProcessJanitor(ctx context.Context) {
	if d == nil || d.runtimeProcessJanitor == nil {
		return
	}
	d.runtimeJanitorOnce.Do(func() {
		d.runtimeProcessJanitor.Start(ctx)
	})
}

func (d *runtimeProcessService) registerRuntimeProcess(sess *interactiveSession, process *agent.Process) {
	if d == nil || sess == nil || process == nil {
		return
	}
	if d.runtimeProcessRegistry == nil {
		return
	}
	pid := process.PID()
	if pid <= 0 {
		return
	}
	workDir := process.WorkDir()
	if strings.TrimSpace(workDir) == "" {
		return
	}
	runtimeName := strings.TrimSpace(string(sess.runtime.Type))
	if runtimeName == "" {
		runtimeName = process.DriverName()
	}
	startedAt := sess.startedAt
	if startedAt.IsZero() {
		startedAt = time.Now()
	}
	d.runtimeProcessRegistry.Register(runtimeproc.Entry{
		PID:       pid,
		AgentID:   process.AgentID,
		Runtime:   runtimeName,
		WorkDir:   workDir,
		StartedAt: startedAt,
	})
	slog.Debug("registered runtime process", "agent_id", process.AgentID, "pid", pid, "runtime", runtimeName, "work_dir", workDir)
}

func (d *runtimeProcessService) unregisterRuntimeProcess(process *agent.Process) {
	if d == nil || d.runtimeProcessRegistry == nil || process == nil {
		return
	}
	pid := process.PID()
	if pid <= 0 {
		return
	}
	d.runtimeProcessRegistry.UnregisterPID(pid)
	slog.Debug("unregistered runtime process", "agent_id", process.AgentID, "pid", pid)
}

// RecordActiveTokenAnalysisContext flushes the body-free launch attribution
// retained by currently live root Processes. Each Process writes its launch
// ledger at most once.
func (d *runtimeProcessService) RecordActiveTokenAnalysisContext() {
	for process := range d.activeTokenAnalysisProcesses() {
		process.InvalidateTokenAnalysisExposure()
		process.RecordRetainedPromptCompileMetrics()
	}
}

// ActiveTokenAnalysisLaunches snapshots body-free launch identities that can
// still produce usage after a report is generated.
func (d *runtimeProcessService) ActiveTokenAnalysisLaunches() []string {
	processes := d.activeTokenAnalysisProcesses()
	launches := make([]string, 0, len(processes))
	for process := range processes {
		if launchID := process.ActiveTokenAnalysisLaunchID(); launchID != "" {
			launches = append(launches, launchID)
		}
	}
	return launches
}

func (d *runtimeProcessService) activeTokenAnalysisProcesses() map[*agent.Process]struct{} {
	processes := make(map[*agent.Process]struct{})
	if d == nil {
		return processes
	}
	for _, sess := range d.sessionManager.Snapshot(true) {
		if sess == nil {
			continue
		}
		sess.mu.Lock()
		process := sess.process
		sess.mu.Unlock()
		if process != nil {
			processes[process] = struct{}{}
		}
	}
	return processes
}

func (d *runtimeProcessService) activeRuntimeWorkDirs() map[string]struct{} {
	active := make(map[string]struct{})
	if d == nil {
		return active
	}

	sessions := d.sessionManager.Snapshot(true)

	for _, sess := range sessions {
		if sess == nil {
			continue
		}
		sess.mu.Lock()
		process := sess.process
		workDir := sess.workDir
		sess.mu.Unlock()
		if process == nil {
			continue
		}
		if processWorkDir := process.WorkDir(); strings.TrimSpace(processWorkDir) != "" {
			workDir = processWorkDir
		}
		if strings.TrimSpace(workDir) != "" {
			active[workDir] = struct{}{}
		}
	}

	return active
}