#!/usr/bin/env python3
"""
JARVIS Agent for Windows — system monitor that reports metrics to JARVIS HUD.
Install via PowerShell: iwr 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 uuid
from datetime import datetime, timezone
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 = "2.2"

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

_log_file = None

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-sig") 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/1.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/1.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)}

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

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

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"

_last_cpu_counters = None

def get_cpu_percent() -> float:
    global _last_cpu_counters
    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", "wuauserv", "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"]
    if Path(r"C:\Program Files\Docker\Docker\Docker Desktop.exe").exists():
        caps.append("docker")
    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",
        "timestamp":   datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
    }

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

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

    log(f"[JARVIS] Registering as '{agent_id}' 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"[JARVIS] Registered. agent_id={state['agent_id']}")
    return api_key

# ── 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")
            return {"success": True, "message": "Windows agent self-update not implemented"}

        elif cmd_type == "shell":
            if not cmd_data.get("allowed", False):
                return {"success": False, "error": "Shell commands not enabled"}
            cmd_str = cmd_data.get("command", "")
            r = subprocess.run(["powershell", "-NoProfile", "-Command", cmd_str],
                               capture_output=True, text=True, timeout=30)
            return {"success": True, "stdout": r.stdout[:2000], "stderr": r.stderr[:500]}

        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))

    api_key = state.get("api_key", "")
    if 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)
            main()
            return

    headers      = {"X-Agent-Key": api_key}
    last_metrics = 0
    log(f"[JARVIS] Agent v{AGENT_VERSION} (Windows) running. Connecting to {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)

            if now - last_metrics >= poll_interval:
                metrics = collect_metrics(cfg)
                r = api_post(f"{jarvis_url}/api/agent/metrics",
                             {"system": metrics}, headers, ssl_verify=ssl_verify)
                if "error" not in r:
                    last_metrics = now

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

        time.sleep(heartbeat_every)

if __name__ == "__main__":
    main()