package session

import (
	"antidrift/internal/evidence"
	"antidrift/internal/store"
	"time"
)

// bucketKey identifies a time bucket; Title is the scrubbed title.
type bucketKey struct{ Class, Title string }

// EvidenceStats is the in-memory accounting for the current session only.
type EvidenceStats struct {
	SessionID    string
	StartedUnix  int64
	Buckets      map[bucketKey]time.Duration // total time per bucket (unchanged)
	OnTask       map[bucketKey]time.Duration // on-task portion per bucket
	OffTask      map[bucketKey]time.Duration // off-task portion per bucket
	unclassified time.Duration               // idle/pending time; aggregate only
	SwitchCount  int
	Current      evidence.WindowSnapshot
	lastFocusAt  time.Time
	lastKey      bucketKey
	hasLast      bool
}

// applyEvent advances stats by one observation: it credits the prior segment to
// the prior bucket, counts a context switch on key change, and records the new
// current window. Used by both live tracking and crash replay. Caller holds mu.
func (c *Controller) applyEvent(now time.Time, snap evidence.WindowSnapshot) {
	if c.stats.hasLast {
		c.creditLocked(c.stats.lastKey, now.Sub(c.stats.lastFocusAt))
	}
	newKey := keyFor(snap)
	if c.stats.hasLast && newKey != c.stats.lastKey {
		c.stats.SwitchCount++
	}
	c.stats.lastKey = newKey
	c.stats.lastFocusAt = now
	c.stats.hasLast = true
	c.stats.Current = snap
}

// creditLocked credits duration d to bucket k: always to the running total, and
// to the on/off/unclassified split per the live drift status — which, at every
// credit site, is the classification of the segment being credited (applyEvent
// runs before evaluateDriftLocked reclassifies; the end-of-session flush runs
// before the state transition). idle/pending route to unclassified: honest, never
// falsely on-task. Caller holds mu.
func (c *Controller) creditLocked(k bucketKey, d time.Duration) {
	c.stats.Buckets[k] += d
	switch c.driftStatus {
	case driftOnTask:
		c.stats.OnTask[k] += d
	case driftDrifting:
		c.stats.OffTask[k] += d
	default: // driftIdle, driftPending
		c.stats.unclassified += d
	}
}

// replayStats rebuilds in-memory stats from the raw session log after a crash.
func (c *Controller) replayStats(sessionID string) {
	events, err := store.ReplaySession(c.sessionsDir, sessionID)
	if err != nil || len(events) == 0 {
		c.stats = &EvidenceStats{
			SessionID:   sessionID,
			StartedUnix: c.clock().Unix(),
			Buckets:     map[bucketKey]time.Duration{},
			OnTask:      map[bucketKey]time.Duration{},
			OffTask:     map[bucketKey]time.Duration{},
		}
		return
	}
	c.stats = &EvidenceStats{
		SessionID:   sessionID,
		StartedUnix: events[0].AtUnixMillis / 1000,
		Buckets:     map[bucketKey]time.Duration{},
		OnTask:      map[bucketKey]time.Duration{},
		OffTask:     map[bucketKey]time.Duration{},
	}
	// Drift status is not persisted, so replayed pre-crash segments credit to
	// unclassified (driftStatus is driftIdle here). Totals in Buckets are exact;
	// only the on/off split degrades — honest, never falsely on-task.
	for _, e := range events {
		c.applyEvent(time.UnixMilli(e.AtUnixMillis), snapFromEvent(e))
	}
}

func keyFor(snap evidence.WindowSnapshot) bucketKey {
	if !snap.Health.Available {
		return bucketKey{Class: "", Title: unavailableTitle}
	}
	return bucketKey{Class: snap.Class, Title: evidence.ScrubTitle(snap.Title)}
}

func focusEvent(now time.Time, snap evidence.WindowSnapshot) store.FocusEvent {
	return store.FocusEvent{
		AtUnixMillis: now.UnixMilli(),
		Class:        snap.Class,
		Title:        snap.Title,
		Available:    snap.Health.Available,
		Reason:       snap.Health.Reason,
	}
}

func snapFromEvent(e store.FocusEvent) evidence.WindowSnapshot {
	return evidence.WindowSnapshot{
		Title:  e.Title,
		Class:  e.Class,
		Health: evidence.EvidenceHealth{Available: e.Available, Reason: e.Reason},
	}
}