package session

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"ollie/agent"
	"ollie/backend"
	coresession "ollie/session"
	"ollie/toolsrv"
	"ollie/tools"
	"olliesrv/fs"
	"olliesrv/prompts"
)

// Kill terminates a session and removes it from the tree.
func Kill(rs *rootState, id string) {
	rs.mu.Lock()
	sess := rs.sessions[id]
	delete(rs.sessions, id)
	rs.mu.Unlock()
	if sess != nil {
		sess.Cancel()
		sess.Core.Close()
		removePersistedAgent(rs, id)
		rs.cfg.Log.Info("killed session %s", id)
		if rs.cfg.OnSessionKilled != nil {
			rs.cfg.OnSessionKilled(id)
		}
	}
}

// KillFromRoot is the public API for callers with *fs.Tree.
func KillFromRoot(root *fs.Tree, id string) {
	Kill(state(root), id)
}

// Rename renames a session.
func Rename(rs *rootState, old, new string) error {
	rs.mu.Lock()
	sess, ok := rs.sessions[old]
	if !ok {
		rs.mu.Unlock()
		return fmt.Errorf("session not found: %s", old)
	}
	if _, exists := rs.sessions[new]; exists {
		rs.mu.Unlock()
		return fmt.Errorf("session already exists: %s", new)
	}
	if sess.Core.Agent().IsRunning() {
		rs.mu.Unlock()
		return fmt.Errorf("cannot rename while agent is running")
	}

	if err := sess.Core.SetSessionID(new); err != nil {
		rs.mu.Unlock()
		return err
	}

	sess.id = new
	rs.sessions[new] = sess
	delete(rs.sessions, old)
	rs.mu.Unlock()

	sess.AppendLog([]byte(fmt.Sprintf(":: session renamed: %s -> %s\n", old, new)))
	rs.cfg.Log.Info("renamed session %s -> %s", old, new)
	if rs.cfg.OnSessionRenamed != nil {
		rs.cfg.OnSessionRenamed(old, new)
	}
	return nil
}

// RenameFromRoot is the public API for callers with *fs.Tree.
func RenameFromRoot(root *fs.Tree, old, new string) error {
	return Rename(state(root), old, new)
}

// Shutdown interrupts all sessions, waits for idle, persists, and closes.
func Shutdown(root *fs.Tree) {
	rs := state(root)
	InterruptAll(root)
	waitIdle(rs, 100*time.Millisecond, 5*time.Second)
	saveAllSessions(rs)
	rs.mu.Lock()
	ids := make([]string, 0, len(rs.sessions))
	for id := range rs.sessions {
		ids = append(ids, id)
	}
	rs.mu.Unlock()
	for _, id := range ids {
		rs.mu.Lock()
		sess := rs.sessions[id]
		delete(rs.sessions, id)
		rs.mu.Unlock()
		if sess != nil {
			sess.Cancel()
			sess.Core.Close()
			rs.cfg.Log.Info("shutdown session %s", id)
		}
	}
}

func waitIdle(rs *rootState, poll, timeout time.Duration) {
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		allIdle := true
		rs.mu.RLock()
		for _, sess := range rs.sessions {
			if sess.Core.Agent().State() != "idle" {
				allIdle = false
				break
			}
		}
		rs.mu.RUnlock()
		if allIdle {
			return
		}
		time.Sleep(poll)
	}
	rs.cfg.Log.Warn("shutdown: timed out waiting for sessions to become idle, saving anyway")
}

// --- Persistence ---

func activeSessionsDir(rs *rootState) string {
	return filepath.Join(rs.cfg.SessionsDir, "active")
}

func persistSession(rs *rootState, id string) {
	rs.mu.RLock()
	sess, ok := rs.sessions[id]
	rs.mu.RUnlock()
	if !ok {
		return
	}
	dir := activeSessionsDir(rs)
	os.MkdirAll(dir, 0700)
	path := filepath.Join(dir, id+".json")
	if err := sess.Core.SaveSession(path); err != nil {
		rs.cfg.Log.Error("persist session %s: %v", id, err)
	}
}

func removePersistedAgent(rs *rootState, id string) {
	path := filepath.Join(activeSessionsDir(rs), id+".json")
	os.Remove(path)
}

func saveAllSessions(rs *rootState) {
	rs.mu.RLock()
	ids := make([]string, 0, len(rs.sessions))
	for id := range rs.sessions {
		ids = append(ids, id)
	}
	rs.mu.RUnlock()
	for _, id := range ids {
		persistSession(rs, id)
	}
}

func restoreAllSessions(rs *rootState) {
	dir := activeSessionsDir(rs)
	entries, err := os.ReadDir(dir)
	if err != nil {
		return
	}

	type loadedSession struct {
		ps   *agent.PersistedAgent
		name string
	}
	var loaded []loadedSession
	for _, e := range entries {
		if !strings.HasSuffix(e.Name(), ".json") {
			continue
		}
		path := filepath.Join(dir, e.Name())
		ps, err := agent.LoadPersistedAgent(path)
		if err != nil {
			rs.cfg.Log.Error("restore session %s: %v", e.Name(), err)
			continue
		}
		loaded = append(loaded, loadedSession{ps: ps, name: e.Name()})
	}

	if len(loaded) == 0 {
		return
	}

	var wg sync.WaitGroup
	for _, ls := range loaded {
		wg.Add(1)
		go func(ls loadedSession) {
			defer wg.Done()
			if err := restoreSession(rs, ls.ps); err != nil {
				rs.cfg.Log.Error("restore session %s: %v", ls.ps.ID, err)
			}
		}(ls)
	}
	wg.Wait()
}

func restoreSession(rs *rootState, ps *agent.PersistedAgent) error {
	cwd := ps.CWD
	if cwd == "" {
		cwd, _ = os.Getwd()
	}
	agentName := ps.Agent
	if agentName == "" {
		agentName = "default"
	}
	sessID := ps.ID

	cfg := LoadAgentConfig(rs.cfg.AgentsDir, agentName, nil)

	backendName := ps.Backend
	if backendName == "" && cfg != nil && cfg.Backend != "" {
		backendName = cfg.Backend
	}
	be, err := backend.NewWithName(backendName)
	if err != nil {
		return fmt.Errorf("backend: %w", err)
	}
	modelName := ps.Model
	if modelName == "" && cfg != nil && cfg.Model != "" {
		modelName = cfg.Model
	}
	if modelName == "" {
		modelName = os.Getenv("OLLIE_MODEL")
	}
	if modelName != "" {
		be.SetModel(modelName)
	}

	uname := rs.nextUname()
	var newToolServer func() toolsrv.Runner
	var promptEnv []string
	remoteTarget := ps.Remote

	if remoteTarget != "" {
		rsrv, dialErr := toolsrv.RemoteDial(rs.cfg.Ctx, toolsrv.RemoteConfig{
			SSHTarget: remoteTarget,
			CWD:       cwd,
		})
		if dialErr != nil {
			return fmt.Errorf("remote dial: %w", dialErr)
		}
		newToolServer = func() toolsrv.Runner { return rsrv }
		promptEnv = []string{
			"PRIME_CWD=" + cwd,
			"PRIME_PLATFORM=" + rsrv.Info.Platform,
			"PRIME_IS_GIT_REPO=" + fmt.Sprintf("%v", rsrv.Info.IsGitRepo),
		}
	} else {
		var execOpts []toolsrv.Option
		if rs.cfg.Strict {
			execOpts = append(execOpts, toolsrv.WithStrict())
		}
		if rs.cfg.Yolo {
			execOpts = append(execOpts, toolsrv.WithYolo())
		}
		if rs.cfg.ToolRegistry != nil {
			execOpts = append(execOpts, toolsrv.WithToolRegistry(rs.cfg.ToolRegistry, sessID))
		}
		if rs.cfg.SkillsRegistry != nil {
			execOpts = append(execOpts, toolsrv.WithSkillsRegistry(rs.cfg.SkillsRegistry))
			execOpts = append(execOpts, toolsrv.WithBuiltins(tools.Builtins()))
		}
		newToolServer = func() toolsrv.Runner {
			srv, _ := toolsrv.Dial(rs.cfg.Ctx, toolsrv.LocalAddr(cwd), execOpts...)
			return srv
		}
		promptEnv = agent.PromptEnv(cwd)
	}

	env := []string{"OLLIE_SESSION_ID=" + sessID, "OLLIE_UNAME=" + uname}
	env = append(env, promptEnv...)

	var spOverride string
	if cfg != nil {
		spOverride = cfg.SystemPrompt
	}
	sysPrompt := prompts.ResolveSystemPrompt(spOverride)
	envMap := make(map[string]string)
	for _, e := range env {
		if k, v, ok := strings.Cut(e, "="); ok {
			envMap[k] = v
		}
	}
	opModel := prompts.OperationalModel(rs.cfg.Enable9P, rs.cfg.EnableDBus, envMap)
	platform := "linux"
	isGitRepo := false
	for _, e := range promptEnv {
		if k, v, ok := strings.Cut(e, "="); ok {
			switch k {
			case "PRIME_PLATFORM":
				platform = v
			case "PRIME_IS_GIT_REPO":
				isGitRepo = v == "true"
			}
		}
	}
	envBlock := prompts.Environment(cwd, platform, isGitRepo, "")

	toolSrv := newToolServer()
	rt := agent.BuildRuntime(cfg, toolSrv, cwd, env, sysPrompt, opModel, envBlock)

	restoredSession := agent.RestoreHistory(ps)

	var sessPtr *Session
	core := coresession.New(coresession.Config{
		Backend:        be,
		AgentName:      agentName,
		AgentsDir:      rs.cfg.AgentsDir,
		SessionsDir:    rs.cfg.SessionsDir,
		SessionID:      sessID,
		AgentID:        uname,
		CWD:            cwd,
		Remote:         remoteTarget,
		History:        restoredSession,
		Runtime:        rt,
		NewToolServer:  newToolServer,
		PromptEnvExtra: promptEnv,
		BaseLayers:     []string{sysPrompt, opModel, envBlock},
		Log:            rs.cfg.Sink.NewLogger("core"),
		ReadPlanStep: func() string {
			if sessPtr == nil {
				return ""
			}
			sessPtr.mu.RLock()
			data := make([]byte, len(sessPtr.plan))
			copy(data, sessPtr.plan)
			sessPtr.mu.RUnlock()
			return coresession.NextUncheckedStep(data)
		},
	})

	sessionCtx, sessionCancel := context.WithCancel(rs.cfg.Ctx)
	sess := NewSession(sessID, core, sessionCtx, sessionCancel)
	sessPtr = sess
	sess.uname = uname
	sess.remote = remoteTarget

	replayMessagesToLog(sess, ps.Messages)

	rs.mu.Lock()
	rs.sessions[sessID] = sess
	rs.mu.Unlock()

	rs.cfg.Log.Info("restored session %s (backend=%s model=%s agent=%s)", sessID, backendName, modelName, agentName)
	if rs.cfg.OnSessionCreated != nil {
		rs.cfg.OnSessionCreated(sessID, sess)
	}
	return nil
}