#!/usr/bin/env python3
"""
JARVIS Agent for Windows — system monitor that reports metrics to JARVIS HUD.
Install via PowerShell (as Admin): irm https://jarvis.orbishosting.com/agent/install-windows.ps1 | iex
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 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.0"

# ── 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, "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)
            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...")
                time.sleep(60)

    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 True:
        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}")

        time.sleep(heartbeat_every)


if __name__ == "__main__":
    main()