package session

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"ollie/cmd/olliesrv/internal/agent"
	"ollie/cmd/olliesrv/internal/backend"
)

// --- Persistence ---

// PersistedSession is the on-disk format for a saved session.
type PersistedSession struct {
	ID         string           `json:"id"`
	Name       string           `json:"name"`
	CWD        string           `json:"cwd"`
	Remote     string           `json:"remote,omitempty"`
	Yolo       bool             `json:"yolo,omitempty"`
	Paused     bool             `json:"paused,omitempty"`
	Workflow   string           `json:"workflow,omitempty"`
	Variant   string           `json:"variant,omitempty"`
	Goal       string           `json:"goal,omitempty"`
	GoalStatus string           `json:"goalStatus,omitempty"`
	Agents     []PersistedAgent `json:"agents"`
}

// PersistedAgent is per-agent state within a session.
type PersistedAgent struct {
	ID       string            `json:"id"`
	Name     string            `json:"name"`
	ParentID string            `json:"parent,omitempty"`
	Profile  string            `json:"profile"`
	Backend  string            `json:"backend"`
	Model    string            `json:"model"`
	Peers    []string          `json:"peers,omitempty"`
	Messages []backend.Message `json:"messages"`
	// Usage tracking
	TotalInputTokens         int     `json:"totalInputTokens,omitempty"`
	TotalCachedInputTokens   int     `json:"totalCachedInputTokens,omitempty"`
	TotalCacheCreationTokens int     `json:"totalCacheCreationTokens,omitempty"`
	TotalOutputTokens        int     `json:"totalOutputTokens,omitempty"`
	TotalRequests            int     `json:"totalRequests,omitempty"`
	Estimated                bool    `json:"estimated,omitempty"`
	LastTurnCostUSD          float64 `json:"lastTurnCostUSD,omitempty"`
	SessionCostUSD           float64 `json:"sessionCostUSD,omitempty"`
}

// ActiveSessionsDir returns the directory for persisted sessions.
func ActiveSessionsDir() string {
	return filepath.Join(pkgSessionsDir, "active")
}

// PersistSession saves a session (with all agents) to disk.
func PersistSession(name string) error {
	sess := Lookup(name)
	if sess == nil {
		return nil
	}
	agents := sess.Agents()
	if len(agents) == 0 {
		return nil
	}

	goalText, goalStatus := sess.Goal()
	ps := PersistedSession{
		ID:         sess.ID,
		Name:       sess.Name(),
		CWD:        sess.Cwd(),
		Remote:     sess.Remote,
		Yolo:       sess.Yolo,
		Paused:     sess.IsPaused(),
		Workflow:   sess.Workflow(),
		Variant:    sess.Variant(),
		Goal:       goalText,
		GoalStatus: goalStatus,
	}

	for _, ag := range agents {
		pa := PersistedAgent{
			ID:       ag.ID(),
			Name:     ag.Name(),
			ParentID: ag.ParentID(),
			Profile:  ag.Profile(),
			Backend:  ag.BackendName(),
			Model:    ag.ModelName(),
			Peers:    ag.Peers(),
			Messages: backend.SanitizeMessages(ag.Messages()),
		}
		if usage := ag.Usage(); usage != nil {
			pa.TotalInputTokens = usage.TotalInputTokens
			pa.TotalCachedInputTokens = usage.TotalCachedInputTokens
			pa.TotalCacheCreationTokens = usage.TotalCacheCreationTokens
			pa.TotalOutputTokens = usage.TotalOutputTokens
			pa.TotalRequests = usage.TotalRequests
			pa.Estimated = usage.Estimated
			pa.LastTurnCostUSD = usage.LastTurnCostUSD
			pa.SessionCostUSD = usage.SessionCostUSD
		}
		ps.Agents = append(ps.Agents, pa)
	}

	dir := ActiveSessionsDir()
	os.MkdirAll(dir, 0700)
	path := filepath.Join(dir, sess.ID+".json")

	data, err := json.MarshalIndent(ps, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(path, data, 0600)
}

// RemovePersistedSession removes a persisted session file.
func RemovePersistedSession(id string) {
	path := filepath.Join(ActiveSessionsDir(), id+".json")
	os.Remove(path)
}

// RestoredAgent holds restored agent info for fs layer.
type RestoredAgent struct {
	ID       string
	Remote   string
	Messages []backend.Message
}

// RestoredSession holds a restored session and metadata for fs layer.
type RestoredSession struct {
	Session   *Session
	Agents    []RestoredAgent
	ToolCount int
}

// RestoreAllSessions restores all persisted sessions.
func RestoreAllSessions() []*RestoredSession {
	dir := ActiveSessionsDir()
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil
	}

	var files []string
	for _, e := range entries {
		if strings.HasSuffix(e.Name(), ".json") {
			files = append(files, filepath.Join(dir, e.Name()))
		}
	}
	if len(files) == 0 {
		return nil
	}

	var rmu sync.Mutex
	var results []*RestoredSession
	var wg sync.WaitGroup
	for _, path := range files {
		wg.Add(1)
		go func(path string) {
			defer wg.Done()
			rs, err := restoreSessionFromFile(path)
			if err != nil {
				pkgLog.Error("restore session %s: %v; removing", path, err)
				os.Remove(path)
				return
			}
			rmu.Lock()
			results = append(results, rs)
			rmu.Unlock()
		}(path)
	}
	wg.Wait()
	return results
}

func restoreSessionFromFile(path string) (*RestoredSession, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	var ps PersistedSession
	if err := json.Unmarshal(data, &ps); err != nil {
		return nil, err
	}
	if ps.ID == "" || len(ps.Agents) == 0 {
		return nil, fmt.Errorf("invalid session file: missing ID or agents")
	}
	return restoreMultiAgentSession(&ps)
}

func restoreMultiAgentSession(ps *PersistedSession) (*RestoredSession, error) {
	ctx, cancel := context.WithCancel(serverCtx)
	sess := NewEmpty(ps.ID, ctx, cancel)
	sess.SetName(ps.Name)
	sess.Remote = ps.Remote
	sess.Yolo = ps.Yolo
	if ps.Paused {
		sess.mu.Lock()
		sess.paused = true
		sess.mu.Unlock()
	}
	if ps.CWD != "" {
		sess.SetCwd(ps.CWD)
	}
	if ps.Workflow != "" {
		sess.SetWorkflow(ps.Workflow)
	}
	if ps.Variant != "" {
		sess.SetVariant(ps.Variant)
	}
	if ps.Goal != "" {
		sess.SetGoal(ps.Goal)
	}
	if ps.GoalStatus != "" {
		sess.SetGoalStatus(ps.GoalStatus)
	}

	var restoredAgents []RestoredAgent

	for _, pa := range ps.Agents {
		// For paused sessions, create a minimal agent without tool server
		if ps.Paused {
			history := agent.RestoreHistoryFromMessages(pa.Messages)
			ag := agent.NewAgent(agent.AgentParams{
				History:   history,
				Runtime:   &agent.Runtime{Preamble: &agent.Preamble{}},
				Profile:   pa.Profile,
				AgentsDir: pkgAgentsDir,
				ID:        pa.ID,
				ParentID:  pa.ParentID,
				GetCwd:    sess.Cwd,
				Log:       pkgSink.NewLogger("agent"),
				SessionID: ps.ID,
				Save:      sess.saveSession,
				Flush:     sess.flushSave,
			})
			ag.SetName(pa.Name)
			ag.RequireMemoryWake()
			if err := sess.AddAgent(ag); err != nil {
				pkgLog.Error("restore agent %s: %v", pa.ID, err)
				ag.Close()
				continue
			}
			restoredAgents = append(restoredAgents, RestoredAgent{
				ID:       pa.ID,
				Remote:   ps.Remote,
				Messages: pa.Messages,
			})
			continue
		}

		// Active session — use shared buildAgent
		history := agent.RestoreHistoryFromMessages(pa.Messages)
		ag, err := buildAgent(sess, AgentParams{
			Profile:  pa.Profile,
			ID:       pa.ID,
			Name:     pa.Name,
			Backend:  pa.Backend,
			Model:    pa.Model,
			Remote:   ps.Remote,
			ParentID: pa.ParentID,
			History:  history,
		})
		if err != nil {
			pkgLog.Error("restore agent %s: %v", pa.ID, err)
			continue
		}
		ag.RequireMemoryWake()
		if err := sess.AddAgent(ag); err != nil {
			pkgLog.Error("restore agent %s: %v", pa.ID, err)
			ag.Close()
			continue
		}
		restoredAgents = append(restoredAgents, RestoredAgent{
			ID:       pa.ID,
			Remote:   ps.Remote,
			Messages: pa.Messages,
		})
	}

	if len(sess.Agents()) == 0 {
		cancel()
		return nil, fmt.Errorf("no agents restored")
	}

	// Restore peer links (each agent's peers were persisted independently).
	for _, pa := range ps.Agents {
		if len(pa.Peers) == 0 {
			continue
		}
		ag := sess.FindAgent(pa.ID)
		if ag == nil {
			continue
		}
		for _, peerName := range pa.Peers {
			ag.AddPeer(peerName)
		}
	}

	sess.Uname = sess.Agents()[0].ID()
	Register(sess.Name(), sess)

	// Start bypass approval loop for restored sessions (unless paused)
	if !ps.Paused {
		sess.StartBypassLoop()
	}

	toolCount := 0
	if conn := sess.ToolsConn(); conn != nil {
		if infos, err := conn.ListTools(); err == nil {
			toolCount = len(infos)
		}
	}
	pkgLog.Info("restored session %s with %d agents, tools=%d, paused=%v", ps.ID, len(restoredAgents), toolCount, ps.Paused)

	return &RestoredSession{
		Session:   sess,
		Agents:    restoredAgents,
		ToolCount: toolCount,
	}, nil
}