//go:build windows

package client

import (
	"encoding/json"
	"os"
	"path/filepath"
	"strconv"
	"strings"
)

// buildPostureScript injects the agent's own service name into the probe script
// so the service allowlist includes the running binary without hard-coding it.
func buildPostureScript() string {
	selfName := selfServiceName()
	return buildPostureScriptWithSelf(selfName)
}

// selfServiceName returns the base name of the running executable (no .exe).
// This is what the agent registers as a Windows service when installed.
func selfServiceName() string {
	exe, err := os.Executable()
	if err != nil || exe == "" {
		return "AetherForge"
	}
	base := filepath.Base(exe)
	return strings.TrimSuffix(base, ".exe")
}

func buildPostureScriptWithSelf(selfSvc string) string {
	// Sanitize: only allow safe service-name characters
	var safe strings.Builder
	for _, c := range selfSvc {
		if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '-' || c == '_' {
			safe.WriteRune(c)
		}
	}
	if safe.Len() == 0 {
		safe.WriteString("AetherForge")
	}
	return strings.Replace(postureProbeTemplate, "__SELF_SVC__", safe.String(), 1)
}

// postureProbeTemplate is the base PowerShell probe.
// __SELF_SVC__ is replaced at runtime with the agent's own service name.
const postureProbeTemplate = `
$ErrorActionPreference = 'SilentlyContinue'
$p = [ordered]@{}

# ── Defender ──────────────────────────────────────────────────────────────────
try {
  $mp = Get-MpComputerStatus
  if ($mp) {
    $p.defender_enabled = [bool]$mp.AntivirusEnabled
    $p.defender_rtp     = [bool]$mp.RealTimeProtectionEnabled
  }
} catch {}

# ── AV products via WMI / CIM (catches third-party AV) ───────────────────────
try {
  $avList = @(Get-CimInstance -Namespace root\SecurityCenter2 -ClassName AntiVirusProduct |
              Select-Object -ExpandProperty displayName)
  if ($avList.Count -gt 0) { $p.av_products = $avList }
} catch {
  try {
    $avList = @(Get-WmiObject -Namespace root\SecurityCenter2 -Class AntiVirusProduct |
                Select-Object -ExpandProperty displayName)
    if ($avList.Count -gt 0) { $p.av_products = $avList }
  } catch {}
}

# ── Firewall — per profile (T1686.003) ───────────────────────────────────────
try {
  foreach ($profile in (Get-NetFirewallProfile)) {
    switch ($profile.Name) {
      'Domain'  { $p.firewall_domain  = [bool]$profile.Enabled }
      'Private' { $p.firewall_private = [bool]$profile.Enabled }
      'Public'  { $p.firewall_public  = [bool]$profile.Enabled }
    }
  }
} catch {}

# ── SSH service ───────────────────────────────────────────────────────────────
try {
  $svc = Get-Service -Name sshd
  $p.ssh_listening = ($svc.Status -eq 'Running')
} catch { $p.ssh_listening = $false }

# ── Patch age — last installed hotfix ─────────────────────────────────────────
try {
  $hf = Get-HotFix | Sort-Object InstalledOn -Descending | Select-Object -First 1
  if ($hf -and $hf.InstalledOn) {
    $d    = [datetime]$hf.InstalledOn
    $days = [int]((Get-Date) - $d).TotalDays
    if ($days -lt 0) { $days = 0 }
    $p.last_patch      = $d.ToString('yyyy-MM-dd')
    $p.last_patch_days = $days
    $p.patch_recent    = ($days -le 30)
  }
} catch {}

# ── Pending updates — WUA COM, capped at 15 s ────────────────────────────────
# Runs in a background job so a stuck WU stack cannot block the agent.
try {
  $job = Start-Job -ScriptBlock {
    $sess     = New-Object -ComObject Microsoft.Update.Session
    $searcher = $sess.CreateUpdateSearcher()
    $result   = $searcher.Search("IsInstalled=0 and IsHidden=0 and Type='Software'")
    $result.Updates.Count
  }
  $done = Wait-Job $job -Timeout 15
  if ($done -and $done.State -eq 'Completed') {
    $cnt = [int](Receive-Job $job)
    $p.pending_updates = $cnt
  } else {
    $p.pending_updates = -1   # timed out or failed
    Remove-Job $job -Force
  }
} catch { $p.pending_updates = -1 }

# ── Reboot pending — registry tripwires ──────────────────────────────────────
try {
  $rp = $false
  # Windows Update requested reboot
  if (Test-Path 'HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\WindowsUpdate\Auto Update\RebootRequired') {
    $rp = $true
  }
  # CBS / Feature Update pending
  if (Test-Path 'HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Component Based Servicing\RebootPending') {
    $rp = $true
  }
  # Session Manager pending file rename (common post-patch signal)
  $pfro = (Get-ItemProperty -Path 'HKLM:\SYSTEM\CurrentControlSet\Control\Session Manager' -Name PendingFileRenameOperations -ErrorAction SilentlyContinue)
  if ($pfro -and $pfro.PendingFileRenameOperations) { $rp = $true }
  $p.reboot_pending = $rp
} catch { $p.reboot_pending = $false }

# ── Elevation ─────────────────────────────────────────────────────────────────
try {
  $id = [Security.Principal.WindowsIdentity]::GetCurrent()
  $pr = New-Object Security.Principal.WindowsPrincipal($id)
  $p.agent_elevated = $pr.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
} catch { $p.agent_elevated = $false }

$p.agent_service_ok = $true

# ── T1007 Service Discovery — fixed allowlist only ────────────────────────────
$svcNames = @('sshd','OpenSSH SSH Server','cloudflared','wuauserv','WinDefend','__SELF_SVC__')
$svcs = @()
foreach ($n in $svcNames) {
  try {
    $s = Get-CimInstance -ClassName Win32_Service -Filter "Name='$n'" -ErrorAction SilentlyContinue
    if (-not $s) {
      # Try matching by DisplayName too
      $s = Get-CimInstance -ClassName Win32_Service -Filter "DisplayName='$n'" -ErrorAction SilentlyContinue
    }
    if ($s) {
      $st = if ($s.State -eq 'Running') { 'running' } else { 'stopped' }
      $sm = switch ($s.StartMode) {
        'Auto'     { 'auto' }
        'Manual'   { 'manual' }
        'Disabled' { 'disabled' }
        default    { 'unknown' }
      }
      $svcs += [ordered]@{ name = $s.Name; display_name = $s.DisplayName; status = $st; start_type = $sm }
    } else {
      $svcs += [ordered]@{ name = $n; status = 'not_found'; start_type = 'unknown' }
    }
  } catch {
    $svcs += [ordered]@{ name = $n; status = 'not_found'; start_type = 'unknown' }
  }
}
$p.services = $svcs
$p | ConvertTo-Json -Depth 4 -Compress
`

func collectPosture() *PostureReport {
	if postureCollector != nil {
		return postureCollector()
	}
	out, err := silentCombinedOutput(
		"powershell", "-NoProfile", "-ExecutionPolicy", "Bypass", "-WindowStyle", "Hidden", "-Command",
		buildPostureScript(),
	)
	if err != nil {
		return fallbackPosture()
	}

	raw := strings.TrimSpace(string(out))
	// PowerShell may emit warnings before the JSON; find the last '{'.
	if idx := strings.LastIndex(raw, "{"); idx > 0 {
		raw = raw[idx:]
	}

	var m map[string]interface{}
	if err := json.Unmarshal([]byte(raw), &m); err != nil {
		return fallbackPosture()
	}

	r := &PostureReport{
		DefenderEnabled: jsonBool(m, "defender_enabled"),
		DefenderRTP:     jsonBool(m, "defender_rtp"),
		AVProducts:      jsonStringSlice(m, "av_products"),
		FirewallDomain:  jsonBool(m, "firewall_domain"),
		FirewallPrivate: jsonBool(m, "firewall_private"),
		FirewallPublic:  jsonBool(m, "firewall_public"),
		SSHListening:    jsonBool(m, "ssh_listening"),
		LastPatch:       jsonString(m, "last_patch"),
		LastPatchDays:   jsonInt(m, "last_patch_days"),
		PatchRecent:     jsonBool(m, "patch_recent"),
		PendingUpdates:  jsonInt(m, "pending_updates"),
		RebootPending:   jsonBool(m, "reboot_pending"),
		AgentElevated:   jsonBool(m, "agent_elevated"),
		Services:        jsonServiceSlice(m, "services"),
		AgentServiceOK:  boolPtr(true),
	}
	r.PostureScore = computePostureScore(r)
	return r
}

func fallbackPosture() *PostureReport {
	ssh := probeSSH()
	r := &PostureReport{SSHListening: &ssh, AgentServiceOK: boolPtr(true)}
	r.PostureScore = computePostureScore(r)
	return r
}

// ── JSON field helpers ────────────────────────────────────────────────────────

func jsonBool(m map[string]interface{}, key string) *bool {
	v, ok := m[key]
	if !ok {
		return nil
	}
	switch t := v.(type) {
	case bool:
		return &t
	case string:
		b := strings.EqualFold(t, "true") || t == "1"
		return &b
	}
	return nil
}

func jsonInt(m map[string]interface{}, key string) *int {
	v, ok := m[key]
	if !ok {
		return nil
	}
	switch t := v.(type) {
	case float64:
		n := int(t)
		return &n
	case int:
		return &t
	case string:
		if n, err := strconv.Atoi(t); err == nil {
			return &n
		}
	}
	return nil
}

func jsonString(m map[string]interface{}, key string) *string {
	v, ok := m[key]
	if !ok {
		return nil
	}
	if s, ok := v.(string); ok && s != "" {
		return &s
	}
	return nil
}

func jsonStringSlice(m map[string]interface{}, key string) []string {
	v, ok := m[key]
	if !ok {
		return nil
	}
	switch t := v.(type) {
	case string:
		if t != "" {
			return []string{t}
		}
	case []interface{}:
		var out []string
		for _, item := range t {
			if s, ok := item.(string); ok && s != "" {
				out = append(out, s)
			}
		}
		return out
	}
	return nil
}

// jsonServiceSlice parses the services array from the PS output.
// PS ConvertTo-Json emits either a single object or []interface{}.
func jsonServiceSlice(m map[string]interface{}, key string) []ServiceStatus {
	v, ok := m[key]
	if !ok {
		return nil
	}

	parseOne := func(raw interface{}) (ServiceStatus, bool) {
		obj, ok := raw.(map[string]interface{})
		if !ok {
			return ServiceStatus{}, false
		}
		s := ServiceStatus{Status: "not_found", StartType: "unknown"}
		if n, ok := obj["name"].(string); ok {
			s.Name = n
		}
		if d, ok := obj["display_name"].(string); ok {
			s.DisplayName = d
		}
		if st, ok := obj["status"].(string); ok {
			s.Status = st
		}
		if sm, ok := obj["start_type"].(string); ok {
			s.StartType = sm
		}
		return s, s.Name != ""
	}

	switch t := v.(type) {
	case map[string]interface{}:
		if s, ok := parseOne(t); ok {
			return []ServiceStatus{s}
		}
	case []interface{}:
		var out []ServiceStatus
		seen := map[string]bool{}
		for _, item := range t {
			if s, ok := parseOne(item); ok && !seen[s.Name] {
				seen[s.Name] = true
				out = append(out, s)
			}
		}
		return out
	}
	return nil
}