#!/usr/bin/env python3
"""shelly-mcp — MCP server for managing Shellys and the VLAN1→VLAN30 migration (TP17).

Transport: StreamableHTTP (endpoint /mcp), matching the repo's other MCP servers.
Also serves a browser-viewable status page at `/` and JSON health at `/health`,
so it works as a local dev server (launch.json / preview_start).

Config via env (see .env.example):
  SHELLY_IOT_SSID   target IOT SSID           (default: dirtyvader-iot)
  SHELLY_IOT_PSK    target IOT PSK            (required for set_sta)
  SHELLY_SWEEP      default discovery range   (default: 10.40.41.180-250)
  SHELLY_MCP_HOST   bind host                 (default: 127.0.0.1)
  SHELLY_MCP_PORT   bind port                 (default: 3030)

Run:  python3 server.py         (loads ./.env if present)
"""
import json
import os
import sys
import urllib.request

# reuse the read-only recon logic
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import discover  # probe(), load_inventory(), expand()

from mcp.server.fastmcp import FastMCP
from starlette.requests import Request
from starlette.responses import HTMLResponse, JSONResponse

# ---- tiny .env loader (no external deps) ----
_envfile = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".env")
if os.path.exists(_envfile):
    for _line in open(_envfile):
        _line = _line.strip()
        if _line and not _line.startswith("#") and "=" in _line:
            _k, _v = _line.split("=", 1)
            os.environ.setdefault(_k.strip(), _v.strip().strip('"').strip("'"))

IOT_SSID = os.environ.get("SHELLY_IOT_SSID", "dirtyvader-iot")
IOT_PSK = os.environ.get("SHELLY_IOT_PSK", "")
SWEEP = os.environ.get("SHELLY_SWEEP", "10.40.41.180-250")
HOST = os.environ.get("SHELLY_MCP_HOST", "127.0.0.1")
PORT = int(os.environ.get("SHELLY_MCP_PORT", "3030"))

mcp = FastMCP("shelly-mcp", host=HOST, port=PORT, stateless_http=True,
              instructions="Manage Shellys and drive the VLAN1→VLAN30 (dirtyvader-iot) "
                           "migration. Read tools are safe; set_sta reconfigures live "
                           "Wi-Fi and requires confirm=true.")


# ---------------- helpers ----------------
def _rpc(ip: str, method: str, params: dict | None = None, timeout: float = 6.0) -> dict:
    """Gen2+ Shelly RPC over HTTP POST /rpc."""
    body = json.dumps({"id": 1, "method": method, "params": params or {}}).encode()
    req = urllib.request.Request(f"http://{ip}/rpc", data=body,
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.load(r)


def _sweep(spec: str | None = None) -> list[dict]:
    inv = discover.load_inventory()
    ips = discover.expand(spec or SWEEP)
    found = []
    from concurrent.futures import ThreadPoolExecutor
    with ThreadPoolExecutor(max_workers=32) as ex:
        for res in ex.map(discover.probe, ips):
            if res:
                m = inv.get(res["mac"])
                res["inv_name"] = m["name"] if m else None
                res["inv_id"] = m["shelly_id"] if m else None
                res["inv_ip"] = m["inv_ip"] if m else None
                res["ip_drifted"] = bool(m and m["inv_ip"] and m["inv_ip"] != res["ip"])
                found.append(res)
    found.sort(key=lambda x: (x["inv_id"] or "zzz"))
    return found


# ---------------- MCP tools ----------------
@mcp.tool()
def discover_shellys(sweep: str | None = None) -> dict:
    """Sweep the LAN for Shellys (read-only) and match them to inventory by MAC.

    sweep: IPv4 range like '10.40.41.180-250' or a single IP. Defaults to SHELLY_SWEEP.
    Returns reachable devices (ip, gen, model, fw, mac, inventory name, ip_drift).
    """
    found = _sweep(sweep)
    return {"sweep": sweep or SWEEP, "reachable": len(found), "devices": found}


@mcp.tool()
def get_device(ip: str) -> dict:
    """Read a single Shelly's identity + current Wi-Fi STA status (read-only)."""
    ident = discover.probe(ip)
    if not ident:
        return {"ip": ip, "reachable": False}
    out = {"ip": ip, "reachable": True, "identity": ident}
    try:
        out["wifi_status"] = _rpc(ip, "WiFi.GetStatus")
        out["wifi_config"] = _rpc(ip, "WiFi.GetConfig")
    except Exception as e:
        out["rpc_error"] = f"{type(e).__name__}: {e}"
    return out


@mcp.tool()
def set_sta(ip: str, confirm: bool = False, ssid: str | None = None) -> dict:
    """MIGRATION ACTION — point a Shelly's Wi-Fi STA at the IOT SSID, then reboot.

    Reconfigures a LIVE device: it drops off VLAN1 and rejoins on VLAN30 via DHCP.
    A wrong PSK leaves it in AP-mode (physical reset needed). Therefore:
      - confirm MUST be true to actually write,
      - SHELLY_IOT_PSK must be set in the environment.
    With confirm=false this is a dry-run (returns the exact config that would be sent).
    """
    target_ssid = ssid or IOT_SSID
    ident = discover.probe(ip)
    plan = {"ip": ip, "target_ssid": target_ssid,
            "device": ident, "psk_configured": bool(IOT_PSK)}
    if not ident:
        return {**plan, "ok": False, "error": "device not reachable"}
    if not IOT_PSK:
        return {**plan, "ok": False, "error": "SHELLY_IOT_PSK not set in environment"}
    if not confirm:
        return {**plan, "ok": False, "dry_run": True,
                "would_send": {"method": "WiFi.SetConfig",
                               "params": {"config": {"sta": {"ssid": target_ssid,
                                                             "pass": "***", "enable": True,
                                                             "ipv4mode": "dhcp"}}}}}
    try:
        cfg = {"config": {"sta": {"ssid": target_ssid, "pass": IOT_PSK,
                                  "enable": True, "ipv4mode": "dhcp"}}}
        res = _rpc(ip, "WiFi.SetConfig", cfg)
        reboot = _rpc(ip, "Shelly.Reboot")
        return {**plan, "ok": True, "set_result": res, "reboot": reboot,
                "note": "device rebooting; expect it on VLAN30 (10.40.30.x) shortly"}
    except Exception as e:
        return {**plan, "ok": False, "error": f"{type(e).__name__}: {e}"}


@mcp.tool()
def verify_on_iot(ip: str) -> dict:
    """Verify a device now sits on the IOT SSID / VLAN30 (post-migration check).

    ip = the device's NEW (10.40.30.x) address. Confirms reachability + SSID.
    """
    ident = discover.probe(ip)
    if not ident:
        return {"ip": ip, "on_iot": False, "reachable": False}
    out = {"ip": ip, "reachable": True, "identity": ident}
    try:
        st = _rpc(ip, "WiFi.GetStatus")
        ssid = (st.get("result") or {}).get("ssid")
        out["ssid"] = ssid
        out["on_iot"] = (ssid == IOT_SSID) and ip.startswith("10.40.30.")
    except Exception as e:
        out["rpc_error"] = f"{type(e).__name__}: {e}"
        out["on_iot"] = False
    return out


@mcp.tool()
def reservation_plan(sweep: str | None = None) -> dict:
    """Build the VLAN30 DHCP-reservation plan (MAC → suggested .30.x, keeping last octet).

    Suggests preserving each Shelly's current last octet under 10.40.30.0/24.
    Output feeds MikroMCP manage_dhcp_lease. Read-only.
    """
    found = _sweep(sweep)
    rows = []
    for f in found:
        last = f["ip"].split(".")[-1]
        rows.append({"name": f.get("inv_name") or f.get("inv_id") or "?",
                     "mac": f["mac"], "current_ip": f["ip"],
                     "suggested_v30_ip": f"10.40.30.{last}", "gen": f["gen"]})
    return {"count": len(rows), "server": "dhcp-v30", "reservations": rows}


# ---------------- browser-viewable routes ----------------
# NOTE: sync route fns — Starlette runs them in a threadpool, so the blocking
# network sweep never freezes the async event loop (which also serves /mcp).
@mcp.custom_route("/health", methods=["GET"])
def health(_req: Request) -> JSONResponse:
    return JSONResponse({"status": "ok", "server": "shelly-mcp",
                         "iot_ssid": IOT_SSID, "psk_configured": bool(IOT_PSK),
                         "mcp_endpoint": "/mcp"})


@mcp.custom_route("/", methods=["GET"])
def index(_req: Request) -> HTMLResponse:
    try:
        found = _sweep()
    except Exception as e:
        found = []
        err = f"{type(e).__name__}: {e}"
    else:
        err = None
    rows = "".join(
        f"<tr><td>{f.get('inv_id') or '?'}</td><td>{f['ip']}</td>"
        f"<td>g{f['gen']}</td><td>{f['model']}</td><td>{f['fw']}</td>"
        f"<td>{f['mac']}</td><td>{f.get('inv_name') or ''}</td>"
        f"<td>{'⚠︎ '+(f.get('inv_ip') or '') if f.get('ip_drifted') else ''}</td></tr>"
        for f in found)
    html = f"""<!doctype html><meta charset=utf-8>
<title>shelly-mcp</title>
<style>body{{font:14px/1.5 system-ui;margin:2rem;color:#222}}
h1{{font-size:1.3rem}} table{{border-collapse:collapse;margin-top:1rem}}
td,th{{border:1px solid #ccc;padding:.25rem .6rem;text-align:left}}
th{{background:#f4f4f4}} .m{{color:#666}}</style>
<h1>shelly-mcp <span class=m>· TP17 VLAN1→VLAN30</span></h1>
<p class=m>MCP endpoint <code>/mcp</code> · target SSID <b>{IOT_SSID}</b> ·
PSK {'set' if IOT_PSK else '<b>NOT set</b>'} · sweep <code>{SWEEP}</code></p>
{'<p style=color:#b00>'+err+'</p>' if err else ''}
<p>{len(found)} Shellys reachable:</p>
<table><tr><th>id<th>ip<th>gen<th>model<th>fw<th>mac<th>name<th>drift</tr>{rows}</table>
"""
    return HTMLResponse(html)


if __name__ == "__main__":
    mcp.run(transport="streamable-http")