#!/usr/bin/env python3
"""
JARVIS Agent for Windows — system monitor that reports metrics to JARVIS HUD.
Runs as a Windows Service (Win 8.1+) via pywin32.

Install (run PowerShell as Admin):
    irm https://jarvis.orbishosting.com/agent/install-windows.ps1 | iex

Service management:
    python jarvis-agent-windows.py --startup auto install   # register service
    python jarvis-agent-windows.py start                    # start
    python jarvis-agent-windows.py stop                     # stop
    python jarvis-agent-windows.py remove                   # uninstall
    python jarvis-agent-windows.py debug                    # run in console (for testing)

Config: C:\ProgramData\jarvis-agent\config.json
Logs:   C:\ProgramData\jarvis-agent\jarvis-agent.log
"""

import json
import os
import platform
import socket
import subprocess
import sys
import threading
import time
import urllib.request
import urllib.error
import hashlib
from datetime import datetime
from pathlib import Path

INSTALL_DIR   = Path(r"C:\ProgramData\jarvis-agent")
CONFIG_PATH   = INSTALL_DIR / "config.json"
STATE_PATH    = INSTALL_DIR / "state.json"
LOG_PATH      = INSTALL_DIR / "jarvis-agent.log"
AGENT_VERSION = "3.1"

# Set by the service wrapper so self_update knows to stop instead of exec
_is_service      = False
_stop_event      = threading.Event()
_update_restart  = False  # True when stopping for self-update; triggers SCM restart

# ── Logging ────────────────────────────────────────────────────────────────────

def log(msg: str):
    ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    line = f"[{ts}] {msg}"
    print(line, flush=True)
    try:
        with open(LOG_PATH, "a", encoding="utf-8") as f:
            f.write(line + "\n")
    except Exception:
        pass

# ── Config ─────────────────────────────────────────────────────────────────────

def load_config() -> dict:
    if not CONFIG_PATH.exists():
        print(f"[ERROR] Config not found at {CONFIG_PATH}. Run the installer first.")
        sys.exit(1)
    with open(CONFIG_PATH, encoding="utf-8") as f:
        return json.load(f)

def load_state() -> dict:
    if STATE_PATH.exists():
        with open(STATE_PATH, encoding="utf-8") as f:
            return json.load(f)
    return {}

def save_state(state: dict):
    INSTALL_DIR.mkdir(parents=True, exist_ok=True)
    with open(STATE_PATH, "w", encoding="utf-8") as f:
        json.dump(state, f, indent=2)

# ── HTTP ───────────────────────────────────────────────────────────────────────

import ssl as _ssl

def _make_ssl_ctx(verify: bool):
    if not verify:
        ctx = _ssl.create_default_context()
        ctx.check_hostname = False
        ctx.verify_mode = _ssl.CERT_NONE
        return ctx
    return None

_host_header: str = ""

def api_post(url: str, payload: dict, headers: dict = {}, timeout: int = 15,
             ssl_verify: bool = True) -> dict:
    body = json.dumps(payload).encode()
    req  = urllib.request.Request(url, data=body, method="POST")
    req.add_header("Content-Type", "application/json")
    req.add_header("User-Agent", "JARVIS-Agent-Windows/3.0")
    if _host_header:
        req.add_header("Host", _host_header)
    for k, v in headers.items():
        req.add_header(k, v)
    try:
        ctx = _make_ssl_ctx(ssl_verify)
        with urllib.request.urlopen(req, timeout=timeout, context=ctx) as resp:
            return json.loads(resp.read().decode())
    except urllib.error.HTTPError as e:
        return {"error": f"HTTP {e.code}: {e.read().decode()[:200]}"}
    except Exception as e:
        return {"error": str(e)}

def api_get(url: str, headers: dict = {}, timeout: int = 10,
            ssl_verify: bool = True) -> dict:
    req = urllib.request.Request(url)
    req.add_header("User-Agent", "JARVIS-Agent-Windows/3.0")
    if _host_header:
        req.add_header("Host", _host_header)
    for k, v in headers.items():
        req.add_header(k, v)
    try:
        ctx = _make_ssl_ctx(ssl_verify)
        with urllib.request.urlopen(req, timeout=timeout, context=ctx) as resp:
            return json.loads(resp.read().decode())
    except Exception as e:
        return {"error": str(e)}

# ── PowerShell helper ──────────────────────────────────────────────────────────

def _ps(script: str, timeout: int = 8) -> str:
    try:
        r = subprocess.run(
            ["powershell", "-NoProfile", "-NonInteractive", "-Command", script],
            capture_output=True, text=True, timeout=timeout
        )
        return r.stdout.strip()
    except Exception:
        return ""

# ── Metrics ────────────────────────────────────────────────────────────────────

def get_local_ip() -> str:
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
        s.close()
        return ip
    except Exception:
        return "unknown"

def get_cpu_percent() -> float:
    try:
        out = _ps("(Get-CimInstance Win32_Processor | Measure-Object -Property LoadPercentage -Average).Average")
        return round(float(out), 1)
    except Exception:
        return 0.0

def get_memory() -> dict:
    try:
        out = _ps("$o=Get-CimInstance Win32_OperatingSystem; [PSCustomObject]@{total=$o.TotalVisibleMemorySize;free=$o.FreePhysicalMemory}|ConvertTo-Json")
        d = json.loads(out)
        total_kb = int(d.get("total", 0))
        free_kb  = int(d.get("free", 0))
        used_kb  = total_kb - free_kb
        if total_kb == 0:
            return {}
        return {
            "total_mb": round(total_kb / 1024, 1),
            "used_mb":  round(used_kb  / 1024, 1),
            "free_mb":  round(free_kb  / 1024, 1),
            "percent":  round(used_kb / total_kb * 100, 1),
        }
    except Exception:
        return {}

def get_disk() -> list:
    try:
        out = _ps("Get-PSDrive -PSProvider FileSystem | Where-Object{$_.Used -ne $null} | Select-Object Name,@{n='used';e={[math]::Round($_.Used/1GB,2)}},@{n='free';e={[math]::Round($_.Free/1GB,2)}} | ConvertTo-Json")
        if not out:
            return []
        items = json.loads(out)
        if isinstance(items, dict):
            items = [items]
        disks = []
        for d in items:
            used  = float(d.get("used", 0))
            free  = float(d.get("free", 0))
            total = used + free
            pct   = round(used / total * 100, 1) if total else 0
            disks.append({
                "mount":   d.get("Name", "?") + ":\\",
                "size":    f"{round(total, 1)}G",
                "used":    f"{used}G",
                "avail":   f"{free}G",
                "percent": str(int(pct)),
            })
        return disks
    except Exception:
        return []

def get_uptime() -> dict:
    try:
        out = _ps("(Get-Date) - (gcim Win32_OperatingSystem).LastBootUpTime | Select-Object -ExpandProperty TotalSeconds")
        secs    = float(out)
        days    = int(secs // 86400)
        hours   = int((secs % 86400) // 3600)
        minutes = int((secs % 3600) // 60)
        return {"seconds": int(secs), "days": days, "hours": hours, "minutes": minutes,
                "human": f"{days}d {hours}h {minutes}m"}
    except Exception:
        return {}

def get_services(cfg: dict) -> list:
    watch = cfg.get("watch_services", ["WinDefend", "Spooler"])
    statuses = []
    for svc in watch:
        try:
            out = _ps(f"(Get-Service -Name '{svc}' -ErrorAction SilentlyContinue).Status")
            status = "active" if out.lower() == "running" else (out.lower() or "unknown")
            statuses.append({"service": svc, "status": status})
        except Exception:
            statuses.append({"service": svc, "status": "unknown"})
    return statuses

def detect_capabilities(cfg: dict) -> list:
    caps = ["metrics", "commands"]
    import shutil
    if Path(r"C:\Program Files\Docker\Docker\Docker Desktop.exe").exists():
        caps.append("docker")
    # Screenshot via PIL or PowerShell .NET (always available on Windows)
    caps.append("screenshot")
    caps.append("sysinfo")
    return caps

def collect_metrics(cfg: dict) -> dict:
    return {
        "hostname":    cfg.get("hostname", socket.gethostname()),
        "cpu_percent": get_cpu_percent(),
        "memory":      get_memory(),
        "disk":        get_disk(),
        "uptime":      get_uptime(),
        "load":        [0, 0, 0],
        "services":    get_services(cfg),
        "platform":    "Windows",
        "os_version":  platform.version(),
        "timestamp":   datetime.utcnow().isoformat() + "Z",
    }

# ── Registration ───────────────────────────────────────────────────────────────

def register(cfg: dict, state: dict) -> str:
    hostname     = cfg.get("hostname", socket.gethostname().lower())
    agent_type   = cfg.get("agent_type", "windows")
    ip           = get_local_ip()
    capabilities = detect_capabilities(cfg)
    agent_id     = cfg.get("agent_id", f"{hostname}_windows")
    ssl_verify   = bool(cfg.get("ssl_verify", True))

    log(f"Registering as '{agent_id}' ({agent_type}) from {ip}...")

    result = api_post(
        f"{cfg['jarvis_url']}/api/agent/register",
        {"hostname": hostname, "version": AGENT_VERSION, "agent_type": agent_type, "ip_address": ip,
         "capabilities": capabilities, "agent_id": agent_id},
        headers={"X-Registration-Key": cfg["registration_key"]},
        ssl_verify=ssl_verify,
    )

    if "error" in result:
        log(f"[ERROR] Registration failed: {result['error']}")
        return ""

    api_key = result.get("api_key", "")
    if api_key:
        state["api_key"]  = api_key
        state["agent_id"] = result.get("agent_id", agent_id)
        save_state(state)
        log(f"Registered. agent_id={state['agent_id']}")
    return api_key

# ── Screenshot ─────────────────────────────────────────────────────────────────

def _take_screenshot(cmd_data: dict) -> dict:
    import base64, tempfile
    tmp = str(INSTALL_DIR / "screenshot_tmp.png")
    method = "unknown"

    # Try PIL/Pillow first
    try:
        from PIL import ImageGrab
        img = ImageGrab.grab(all_screens=True)
        img.save(tmp, "PNG")
        method = "pil"
    except ImportError:
        pass
    except Exception:
        pass

    # Fallback: PowerShell .NET screenshot (no extra packages needed)
    if method == "unknown":
        ps_script = (
            "Add-Type -AssemblyName System.Windows.Forms,System.Drawing; "
            "$s=[System.Windows.Forms.SystemInformation]::VirtualScreen; "
            "$bmp=New-Object System.Drawing.Bitmap($s.Width,$s.Height); "
            "$g=[System.Drawing.Graphics]::FromImage($bmp); "
            "$g.CopyFromScreen($s.Location,[System.Drawing.Point]::Empty,$s.Size); "
            f"$bmp.Save('{tmp}'); $g.Dispose(); $bmp.Dispose()"
        )
        try:
            r = subprocess.run(
                ["powershell", "-NoProfile", "-NonInteractive", "-Command", ps_script],
                capture_output=True, timeout=15
            )
            if r.returncode == 0 and os.path.exists(tmp):
                method = "powershell"
        except Exception:
            pass

    if method == "unknown" or not os.path.exists(tmp):
        return _sysinfo_snapshot()

    try:
        with open(tmp, "rb") as f:
            raw = f.read()
        b64 = base64.b64encode(raw).decode()
        try:
            os.unlink(tmp)
        except Exception:
            pass
        return {
            "success":    True,
            "method":     method,
            "image_b64":  b64,
            "image_mime": "image/png",
            "file_size":  len(raw),
            "hostname":   socket.gethostname(),
        }
    except Exception as e:
        return {"success": False, "error": str(e), "method": method}

# ── Sysinfo ────────────────────────────────────────────────────────────────────

def _sysinfo_snapshot() -> dict:
    data = {"success": True, "hostname": socket.gethostname(),
            "snapshot_type": "text", "screenshot_available": False, "platform": "Windows"}
    try:
        mem = get_memory()
        data["mem_total_mb"] = int(mem.get("total_mb", 0))
        data["mem_used_mb"]  = int(mem.get("used_mb", 0))
        data["mem_avail_mb"] = int(mem.get("free_mb", 0))
    except Exception:
        pass
    try:
        data["cpu_percent"] = get_cpu_percent()
    except Exception:
        pass
    try:
        data["disk"] = get_disk()
    except Exception:
        pass
    try:
        ut = get_uptime()
        data["uptime"] = ut.get("human", "")
    except Exception:
        pass
    try:
        out = _ps("Get-Process | Sort-Object CPU -Descending | Select-Object -First 8 Name,CPU,WorkingSet | ConvertTo-Json")
        data["top_procs"] = json.loads(out) if out else []
    except Exception:
        pass
    try:
        out = _ps("netstat -an | findstr LISTENING | head -20")
        data["listening_ports"] = out[:800]
    except Exception:
        pass
    return data

# ── Self-update ────────────────────────────────────────────────────────────────

def self_update(cfg: dict) -> bool:
    jarvis_url = cfg.get("jarvis_url", "").rstrip("/")
    default_update_url = f"{jarvis_url}/agent/jarvis-agent-windows.py" if jarvis_url else ""
    update_url = cfg.get("update_url", default_update_url)
    if not update_url:
        return False
    script_path = os.path.abspath(__file__)
    ssl_verify = bool(cfg.get("ssl_verify", True))
    try:
        # Download expected hash
        hash_url = update_url + ".sha256"
        req_hash = urllib.request.Request(hash_url)
        req_hash.add_header("User-Agent", "JARVIS-Agent-Windows/3.0")
        if _host_header:
            req_hash.add_header("Host", _host_header)
        expected_hash = None
        try:
            ctx = _make_ssl_ctx(ssl_verify)
            with urllib.request.urlopen(req_hash, timeout=10, context=ctx) as resp:
                expected_hash = resp.read().decode().strip().split()[0]
        except Exception:
            pass

        # Download new script
        req = urllib.request.Request(update_url)
        req.add_header("User-Agent", "JARVIS-Agent-Windows/3.0")
        if _host_header:
            req.add_header("Host", _host_header)
        ctx = _make_ssl_ctx(ssl_verify)
        with urllib.request.urlopen(req, timeout=30, context=ctx) as resp:
            new_content = resp.read()

        # Verify hash
        if expected_hash:
            actual_hash = hashlib.sha256(new_content).hexdigest()
            if actual_hash != expected_hash:
                log(f"Update hash mismatch (expected {expected_hash[:16]}… got {actual_hash[:16]}…) — aborting")
                return False

        with open(script_path, "rb") as f:
            current = f.read()
        if new_content != current:
            log(f"Update verified — replacing {script_path} and restarting...")
            with open(script_path, "wb") as f:
                f.write(new_content)
            if _is_service:
                global _update_restart
                _update_restart = True
                log("Running as service — stopping for SCM-managed restart after update.")
                _stop_event.set()
            else:
                os.execv(sys.executable, [sys.executable] + sys.argv)
            return True
        return False
    except Exception as e:
        log(f"Self-update check failed: {e}")
        return False

# ── Command execution ──────────────────────────────────────────────────────────

def execute_command(cmd: dict, cfg: dict) -> dict:
    cmd_type = cmd.get("command_type", "")
    cmd_data = cmd.get("command_data", {})
    try:
        if cmd_type == "ping":
            host = cmd_data.get("host", "8.8.8.8")
            r = subprocess.run(["ping", "-n", "3", host], capture_output=True, text=True, timeout=15)
            return {"success": r.returncode == 0, "output": r.stdout}

        elif cmd_type == "update":
            log("[CMD] Self-update requested")
            updated = self_update(cfg)
            return {"success": True, "updated": updated}

        elif cmd_type == "shell":
            # Guard reads LOCAL config — never trust server-supplied flags
            _cfg = load_config()
            if not _cfg.get("allow_shell_commands", False):
                return {"success": False, "error": "Shell commands not enabled in agent config"}
            cmd_str = cmd_data.get("command", "")
            r = subprocess.run(["powershell", "-NoProfile", "-NonInteractive", "-Command", cmd_str],
                               capture_output=True, text=True, timeout=30)
            return {"success": True, "stdout": r.stdout[:2000], "stderr": r.stderr[:500]}

        elif cmd_type == "restart_service":
            svc = cmd_data.get("service", "")
            if not svc:
                return {"success": False, "error": "No service specified"}
            out = _ps(f"Restart-Service -Name '{svc}' -Force -ErrorAction Stop; 'ok'")
            return {"success": "ok" in out.lower(), "output": out}

        elif cmd_type == "get_logs":
            svc   = cmd_data.get("service", "")
            lines = min(int(cmd_data.get("lines", 50)), 200)
            if not svc:
                return {"success": False, "error": "No service specified"}
            out = _ps(f"Get-EventLog -LogName Application -Source '{svc}' -Newest {lines} -ErrorAction SilentlyContinue | Format-List TimeGenerated,EntryType,Message | Out-String")
            return {"success": True, "output": out[:3000]}

        elif cmd_type == "screenshot":
            return _take_screenshot(cmd_data)

        elif cmd_type == "sysinfo":
            return _sysinfo_snapshot()

        else:
            return {"success": False, "error": f"Unknown command type: {cmd_type}"}

    except subprocess.TimeoutExpired:
        return {"success": False, "error": "Command timed out"}
    except Exception as e:
        return {"success": False, "error": str(e)}

# ── Main loop ──────────────────────────────────────────────────────────────────

def main():
    global _host_header

    cfg   = load_config()
    state = load_state()

    jarvis_url      = cfg["jarvis_url"].rstrip("/")
    ssl_verify      = bool(cfg.get("ssl_verify", True))
    _host_header    = cfg.get("host_header", "")
    poll_interval   = int(cfg.get("poll_interval", 30))
    heartbeat_every = int(cfg.get("heartbeat_every", 10))
    update_interval = int(cfg.get("update_check_hours", 24)) * 3600

    # Always re-register on startup to refresh capabilities, version, IP
    api_key = state.get("api_key", "")
    registered_key = register(cfg, state)
    if registered_key:
        api_key = registered_key
    elif not api_key:
        while not api_key:
            api_key = register(cfg, state)
            if not api_key:
                log("[ERROR] Could not register with JARVIS. Retrying in 60s...")
                if _stop_event.wait(60):
                    return

    headers         = {"X-Agent-Key": api_key}
    last_metrics    = 0
    last_update_chk = 0

    log(f"Agent v{AGENT_VERSION} (Windows) running. Polling {jarvis_url} every {heartbeat_every}s.")

    while not _stop_event.is_set():
        now = time.time()
        try:
            hb = api_post(f"{jarvis_url}/api/agent/heartbeat", {}, headers, ssl_verify=ssl_verify)
            if "error" in hb:
                log(f"[WARN] Heartbeat failed: {hb['error']}")
            else:
                for cmd in hb.get("commands", []):
                    log(f"[CMD] Executing: {cmd['command_type']}")
                    result = execute_command(cmd, cfg)
                    api_post(f"{jarvis_url}/api/agent/command_result",
                             {"command_id": cmd["id"], "success": result.get("success", False), "result": result},
                             headers, ssl_verify=ssl_verify)

            # Self-update check
            if now - last_update_chk >= update_interval:
                last_update_chk = now
                self_update(cfg)

            # Push metrics
            if now - last_metrics >= poll_interval:
                metrics = collect_metrics(cfg)
                api_post(f"{jarvis_url}/api/agent/metrics",
                         {"type": "system", "data": metrics}, headers, ssl_verify=ssl_verify)
                last_metrics = now

        except Exception as e:
            log(f"[ERROR] Loop error: {e}")

        if _stop_event.wait(heartbeat_every):
            break  # service stop was requested

    log("Agent stopped.")


# ── Windows Service wrapper ────────────────────────────────────────────────────

try:
    import win32serviceutil
    import win32service
    import win32event
    import servicemanager
    _HAS_WIN32 = True
except ImportError:
    _HAS_WIN32 = False

if _HAS_WIN32:
    class JarvisAgentService(win32serviceutil.ServiceFramework):
        _svc_name_         = "JARVISAgent"
        _svc_display_name_ = "JARVIS AI Agent"
        _svc_description_  = "JARVIS system monitoring and AI agent — reports metrics to JARVIS HUD"

        def __init__(self, args):
            win32serviceutil.ServiceFramework.__init__(self, args)
            self._svc_stop_event = win32event.CreateEvent(None, 0, 0, None)

        def SvcStop(self):
            self.ReportServiceStatus(win32service.SERVICE_STOP_PENDING)
            _stop_event.set()
            win32event.SetEvent(self._svc_stop_event)

        def SvcDoRun(self):
            global _is_service
            _is_service = True
            servicemanager.LogMsg(
                servicemanager.EVENTLOG_INFORMATION_TYPE,
                servicemanager.PYS_SERVICE_STARTED,
                (self._svc_name_, ""),
            )
            main()
            if _update_restart:
                # Non-zero exit triggers SCM failure recovery → automatic restart
                sys.exit(1)


if __name__ == "__main__":
    if _HAS_WIN32:
        if len(sys.argv) == 1:
            # No args — SCM is starting us as a service
            servicemanager.Initialize()
            servicemanager.PrepareToHostSingle(JarvisAgentService)
            servicemanager.StartServiceCtrlDispatcher()
        else:
            # install / start / stop / remove / debug
            win32serviceutil.HandleCommandLine(JarvisAgentService)
    else:
        # pywin32 not available — run directly (useful for testing)
        main()