package main

import (
	"encoding/json"
	"os"
	"path/filepath"
	"sync"
	"time"
)

// TrackedSession holds a session's info alongside state-transition tracking.
type TrackedSession struct {
	Info      SessionInfo
	PrevState string
}

// SessionRegistry maintains an in-memory index of active sessions.
// It tracks WaitingSince transitions: when a session moves to "Needs Input",
// the timestamp is recorded; when it leaves that state, it is cleared.
type SessionRegistry struct {
	mu       sync.RWMutex
	sessions map[string]*TrackedSession
}

func NewRegistry() *SessionRegistry {
	return &SessionRegistry{
		sessions: make(map[string]*TrackedSession),
	}
}

// Update adds or refreshes a session in the registry.
// WaitingSince is set when transitioning to "Needs Input" and cleared otherwise.
func (r *SessionRegistry) Update(info SessionInfo) {
	r.mu.Lock()
	defer r.mu.Unlock()

	existing, ok := r.sessions[info.SessionID]
	if !ok {
		existing = &TrackedSession{}
		r.sessions[info.SessionID] = existing
	}

	isWaiting := info.State == "Needs Input"
	wasWaiting := existing.PrevState == "Needs Input"

	if isWaiting && !wasWaiting {
		now := time.Now()
		info.WaitingSince = &now
	} else if isWaiting && wasWaiting {
		// Keep existing timestamp
		info.WaitingSince = existing.Info.WaitingSince
	}
	// If not waiting, WaitingSince stays nil (default)

	existing.Info = info
	existing.PrevState = info.State
}

// List returns a snapshot of all tracked sessions.
func (r *SessionRegistry) List() []SessionInfo {
	r.mu.RLock()
	defer r.mu.RUnlock()

	result := make([]SessionInfo, 0, len(r.sessions))
	for _, ts := range r.sessions {
		result = append(result, ts.Info)
	}
	return result
}

// RemoveStale removes sessions whose SessionID is not in activeIDs.
func (r *SessionRegistry) RemoveStale(activeIDs map[string]bool) {
	r.mu.Lock()
	defer r.mu.Unlock()

	for id := range r.sessions {
		if !activeIDs[id] {
			delete(r.sessions, id)
		}
	}
}

// LabelStore persists user-assigned labels in a JSON file.
type LabelStore struct {
	mu     sync.RWMutex
	labels map[string]string
	path   string
}

// NewLabelStore loads labels from path. Returns an empty store if the file does not exist.
func NewLabelStore(path string) (*LabelStore, error) {
	ls := &LabelStore{
		labels: make(map[string]string),
		path:   path,
	}

	data, err := os.ReadFile(path)
	if err != nil {
		if os.IsNotExist(err) {
			return ls, nil
		}
		return nil, err
	}

	if err := json.Unmarshal(data, &ls.labels); err != nil {
		return nil, err
	}

	return ls, nil
}

// Set assigns a label to a session and persists to disk.
func (ls *LabelStore) Set(sessionID, label string) error {
	ls.mu.Lock()
	defer ls.mu.Unlock()

	ls.labels[sessionID] = label
	return ls.save()
}

// Get returns the label for a session, or empty string if not found.
func (ls *LabelStore) Get(sessionID string) string {
	ls.mu.RLock()
	defer ls.mu.RUnlock()

	return ls.labels[sessionID]
}

func (ls *LabelStore) save() error {
	data, err := json.MarshalIndent(ls.labels, "", "  ")
	if err != nil {
		return err
	}

	dir := filepath.Dir(ls.path)
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}

	return os.WriteFile(ls.path, data, 0o644)
}