#!/usr/bin/env python3
"""Sync MikroTik live PPPoE+DHCP state -> Gaiia MacAddressIpAddressAssignation.

Bypasses the broken Gaiia AWS webhook workflow that's supposed to keep
Gaiia's IP-assignment table in sync with reality. Calls the public-but-
hidden mutations:
    upsertMacAddressIpAddressAssignations(input)
    deactivateMacAddressIpAddressAssignations(input)
both of which are reachable via raw GraphQL even though they don't appear
in Mutation.fields introspection.

Defaults to dry-run. Use --apply to actually push changes.

Usage:
    GAIIA_KEY=... ./scripts/gaiia_ip_sync.py                  # dry-run
    GAIIA_KEY=... ./scripts/gaiia_ip_sync.py --apply          # do it
    GAIIA_KEY=... ./scripts/gaiia_ip_sync.py --routers 982    # one tower
    GAIIA_KEY=... ./scripts/gaiia_ip_sync.py --no-deactivate  # only upserts
"""
from __future__ import annotations

import argparse
import os
import re
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor
from typing import Any

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(REPO, "gaiia", "src"))
from gaiia import GaiiaClient, GaiiaAPIError  # noqa: E402

TOWER_ROUTERS = ["verona", "climax", "culleoka", "newhope", "lowry", "982", "494", "core"]
MTT = os.path.join(REPO, "mikrotik-tool", "mikrotik-tool")

PAYING = {"ACTIVE", "TRIAL"}
CUSTOMER_DHCP_KEYWORDS = ("cpe", "cgnat", "customer", "hotspot")


# ---------- helpers ----------

def norm_mac(s: str) -> str:
    return re.sub(r"[^0-9a-f]", "", (s or "").lower())


def fmt_mac(s: str) -> str:
    n = norm_mac(s)
    return ":".join(n[i:i+2] for i in range(0, 12, 2)).upper() if len(n) == 12 else s


def smush(s: str) -> str:
    return re.sub(r"[^a-z0-9]", "", (s or "").lower())


def parse_kv_line(line: str) -> dict[str, str]:
    kv = {}
    for part in re.split(r"  +", line.strip()):
        if "=" in part:
            k, v = part.split("=", 1)
            kv[k] = v
    return kv


# ---------- MikroTik fetch ----------

def fetch_ppp_active(router: str) -> list[dict[str, Any]]:
    p = subprocess.run([MTT, "api", router, "/ppp/active/print"],
                       capture_output=True, text=True, timeout=90,
                       cwd=os.path.join(REPO, "mikrotik-tool"))
    if p.returncode != 0:
        print(f"!! {router} ppp active: {p.stderr.strip()[:200]}", file=sys.stderr)
        return []
    out = []
    for line in p.stdout.splitlines():
        if not line.startswith(".id="): continue
        kv = parse_kv_line(line)
        if kv.get("service") != "pppoe": continue
        out.append({
            "source": "pppoe",
            "router": router,
            "ip": kv.get("address", ""),
            "mac": norm_mac(kv.get("caller-id", "")),
            "mac_raw": kv.get("caller-id", ""),
            "username": kv.get("name", ""),
            "host": "",
            "uptime": kv.get("uptime", ""),
        })
    return out


def fetch_dhcp_leases(router: str) -> list[dict[str, Any]]:
    p = subprocess.run([MTT, "api", router, "/ip/dhcp-server/lease/print"],
                       capture_output=True, text=True, timeout=90,
                       cwd=os.path.join(REPO, "mikrotik-tool"))
    if p.returncode != 0:
        print(f"!! {router} dhcp leases: {p.stderr.strip()[:200]}", file=sys.stderr)
        return []
    out = []
    for line in p.stdout.splitlines():
        if not line.startswith(".id="): continue
        kv = parse_kv_line(line)
        if kv.get("status") != "bound": continue
        server = (kv.get("active-server") or kv.get("server") or "").lower()
        if not any(k in server for k in CUSTOMER_DHCP_KEYWORDS):
            continue
        mac = kv.get("active-mac-address") or kv.get("mac-address") or ""
        if not mac:
            continue
        out.append({
            "source": "dhcp",
            "router": router,
            "ip": kv.get("active-address") or kv.get("address") or "",
            "mac": norm_mac(mac),
            "mac_raw": mac,
            "username": "",
            "host": kv.get("host-name", ""),
            "server": server,
        })
    return out


def fetch_router_state(routers: list[str]) -> list[dict[str, Any]]:
    with ThreadPoolExecutor(max_workers=len(routers) * 2) as ex:
        futs = []
        for r in routers:
            futs.append(ex.submit(fetch_ppp_active, r))
            futs.append(ex.submit(fetch_dhcp_leases, r))
        sessions = []
        for f in futs:
            sessions.extend(f.result())
    return sessions


# ---------- Gaiia fetch ----------

def paginate(client: GaiiaClient, q_with: str, q_first: str, path: str):
    cur: str | None = None
    while True:
        d = client.query(q_first) if cur is None else client.query(q_with, {"after": cur})
        node = d
        for part in path.split("."):
            node = node[part]
        for n in (node.get("nodes") or []):
            yield n
        pi = node.get("pageInfo") or {}
        if not pi.get("hasNextPage"): return
        cur = pi.get("endCursor")
        if not cur: return


def fetch_gaiia(client: GaiiaClient):
    """Return (accounts_by_id, mac_to_inventory, ip_to_current_assignment, all_assignments)."""
    accounts: dict[str, dict] = {}
    q_acct_with = """query Q($after:String!){ accounts(first:100,after:$after){
      nodes{ id readableId displayName status{ name }
             billingSubscriptions(first:50){ nodes{ status } } }
      pageInfo{ hasNextPage endCursor } } }"""
    q_acct_first = """query{ accounts(first:100){
      nodes{ id readableId displayName status{ name }
             billingSubscriptions(first:50){ nodes{ status } } }
      pageInfo{ hasNextPage endCursor } } }"""
    for a in paginate(client, q_acct_with, q_acct_first, "accounts"):
        subs = [s["status"] for s in ((a.get("billingSubscriptions") or {}).get("nodes") or [])]
        accounts[a["id"]] = {
            "id": a["id"],
            "readable": a["readableId"],
            "name": a["displayName"],
            "status": a["status"]["name"],
            "paying": any(s in PAYING for s in subs),
            "subs": subs,
        }

    mac_to_inv: dict[str, dict] = {}
    ip_to_assn: dict[str, list[dict]] = {}
    all_assn: list[dict] = []
    q_inv_with = """query Q($after:String!){ inventoryItems(first:100,after:$after){
      nodes{ id ipAddressV4
        fields{ nodes{ data modelField{ name } } }
        assignation{ assigneeType assignee{ __typename ... on Account{ id } } }
        ipAddressAssignments(filter:{isActive:true}, first:50){
          nodes{ id ipAddress macAddress isActive type
                 account{ id readableId displayName } } }
      }
      pageInfo{ hasNextPage endCursor } } }"""
    q_inv_first = q_inv_with.replace("query Q($after:String!)", "query").replace(",after:$after", "")
    for it in paginate(client, q_inv_with, q_inv_first, "inventoryItems"):
        macs = []
        for f in ((it.get("fields") or {}).get("nodes") or []):
            mf = f.get("modelField") or {}
            if mf.get("name", "").lower().startswith("mac") and f.get("data"):
                m = norm_mac(f["data"])
                if len(m) == 12:
                    macs.append(m)
        assn = it.get("assignation") or {}
        assignee = assn.get("assignee") or {}
        acct_id = assignee["id"] if assignee.get("__typename") == "Account" else None
        for m in macs:
            rec = {"item_id": it["id"], "account_id": acct_id}
            if m in mac_to_inv and mac_to_inv[m]["account_id"]:
                continue
            mac_to_inv[m] = rec
        for ipa in (it.get("ipAddressAssignments") or {}).get("nodes") or []:
            if not ipa.get("ipAddress"): continue
            entry = {**ipa, "inventory_item_id": it["id"]}
            ip_to_assn.setdefault(ipa["ipAddress"], []).append(entry)
            all_assn.append(entry)

    return accounts, mac_to_inv, ip_to_assn, all_assn


# ---------- resolution ----------

def build_username_idx(accounts):
    idx: dict[str, list[str]] = {}
    for aid, a in accounts.items():
        k = smush(a["name"])
        if k:
            idx.setdefault(k, []).append(aid)
    return idx


def resolve_session(s, mac_to_inv, accounts, name_idx):
    """Return (account_id, inventory_item_id, match_via) or (None, None, reason)."""
    inv = mac_to_inv.get(s["mac"]) if s["mac"] else None
    if inv and inv["account_id"]:
        return inv["account_id"], inv["item_id"], "mac"
    # fallback: fuzzy displayName match on username (PPPoE) or host (DHCP)
    key = smush(s["username"] or s["host"])
    cands = name_idx.get(key, []) if key else []
    paying = [c for c in cands if accounts[c]["paying"]]
    if len(paying) == 1:
        return paying[0], inv["item_id"] if inv else None, "name-paying"
    if len(cands) == 1:
        return cands[0], inv["item_id"] if inv else None, "name"
    if not cands:
        return None, inv["item_id"] if inv else None, ("mac-unassigned" if inv else "unmatched")
    return None, inv["item_id"] if inv else None, "name-ambiguous"


# ---------- diff ----------

def plan(sessions, accounts, mac_to_inv, ip_to_assn):
    name_idx = build_username_idx(accounts)
    upserts = []
    skipped = []  # sessions we can't resolve
    seen_ips = set()

    for s in sessions:
        if ":" in s["ip"] or not s["ip"]:
            continue  # IPv6 / blank skip
        seen_ips.add(s["ip"])
        acct_id, item_id, via = resolve_session(s, mac_to_inv, accounts, name_idx)
        if not acct_id:
            skipped.append((s, via))
            continue
        current = ip_to_assn.get(s["ip"], [])
        # is there an exact-match active assignment already?
        match = None
        for c in current:
            same_mac = norm_mac(c["macAddress"]) == s["mac"]
            same_acct = (c.get("account") or {}).get("id") == acct_id
            same_item = c.get("inventory_item_id") == item_id if item_id else True
            if same_mac and same_acct and same_item:
                match = c
                break
        if match:
            continue  # already correct
        upserts.append({
            "session": s,
            "via": via,
            "account_id": acct_id,
            "inventory_item_id": item_id,
            "current": current,
        })

    # Deactivations: any IP that has an active Gaiia assignment but no live session for it
    deactivations = []
    for ip, assignments in ip_to_assn.items():
        if ip in seen_ips:
            continue
        for a in assignments:
            deactivations.append({"ip": ip, "id": a["id"], "current": a})

    return upserts, deactivations, skipped


# ---------- mutations ----------

def do_upserts(client, upserts):
    if not upserts:
        return
    BATCH = 50
    payload = []
    for u in upserts:
        s = u["session"]
        entry = {
            "ipAddress": s["ip"],
            "type": "IPV4",
            "macAddress": fmt_mac(s["mac"]),
            "accountId": u["account_id"],
            "expired": False,
        }
        if u["inventory_item_id"]:
            entry["inventoryItemId"] = u["inventory_item_id"]
        payload.append(entry)

    mut = """mutation M($i: UpsertMacAddressIpAddressAssignationsInput!) {
      upsertMacAddressIpAddressAssignations(input: $i) {
        assignations { id ipAddress macAddress account { readableId } }
        errors { code message }
      }
    }"""
    total_ok = 0
    for i in range(0, len(payload), BATCH):
        chunk = payload[i:i+BATCH]
        try:
            res = client.mutate(mut, "upsertMacAddressIpAddressAssignations",
                                {"i": {"assignations": chunk}})
            n = len(res.get("assignations") or [])
            total_ok += n
            print(f"  upsert batch {i//BATCH+1}: {n}/{len(chunk)} ok")
        except GaiiaAPIError as e:
            print(f"  upsert batch {i//BATCH+1}: ERROR {str(e)[:300]}", file=sys.stderr)
    print(f"upserts applied: {total_ok}/{len(payload)}")


def do_deactivates(client, deactivations):
    if not deactivations:
        return
    BATCH = 50
    ids = [d["id"] for d in deactivations]
    mut = """mutation M($i: DeactivateMacAddressIpAddressAssignationsInput!) {
      deactivateMacAddressIpAddressAssignations(input: $i) {
        assignations { id ipAddress }
        errors { code message }
      }
    }"""
    total_ok = 0
    for i in range(0, len(ids), BATCH):
        chunk = ids[i:i+BATCH]
        try:
            res = client.mutate(mut, "deactivateMacAddressIpAddressAssignations",
                                {"i": {"ids": chunk}})
            n = len(res.get("assignations") or [])
            total_ok += n
            print(f"  deactivate batch {i//BATCH+1}: {n}/{len(chunk)} ok")
        except GaiiaAPIError as e:
            print(f"  deactivate batch {i//BATCH+1}: ERROR {str(e)[:300]}", file=sys.stderr)
    print(f"deactivations applied: {total_ok}/{len(ids)}")


# ---------- main ----------

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--apply", action="store_true", help="actually call mutations (default: dry-run)")
    ap.add_argument("--routers", help="comma-separated subset")
    ap.add_argument("--no-deactivate", action="store_true", help="don't deactivate stale Gaiia assignments")
    ap.add_argument("--limit", type=int, default=0, help="cap upserts (for testing)")
    args = ap.parse_args()

    if not os.environ.get("GAIIA_KEY"):
        sys.exit("GAIIA_KEY not set")

    routers = args.routers.split(",") if args.routers else TOWER_ROUTERS
    print(f"Fetching live state from {len(routers)} routers...", file=sys.stderr)
    sessions = fetch_router_state(routers)
    print(f"  {len(sessions)} customer sessions "
          f"({sum(1 for s in sessions if s['source']=='pppoe')} pppoe, "
          f"{sum(1 for s in sessions if s['source']=='dhcp')} dhcp)", file=sys.stderr)

    print("Fetching Gaiia state...", file=sys.stderr)
    with GaiiaClient(timezone="America/Chicago") as g:
        accounts, mac_to_inv, ip_to_assn, all_assn = fetch_gaiia(g)
    paying = sum(1 for a in accounts.values() if a["paying"])
    print(f"  {len(accounts)} accounts ({paying} paying), "
          f"{len(mac_to_inv)} MACs in inventory, "
          f"{len(all_assn)} active IP assignments", file=sys.stderr)

    upserts, deactivations, skipped = plan(sessions, accounts, mac_to_inv, ip_to_assn)

    if args.limit:
        upserts = upserts[:args.limit]
        deactivations = deactivations[:args.limit]

    print()
    print(f"=== Plan ===")
    print(f"  {len(upserts)} upserts (Gaiia is wrong / missing)")
    print(f"  {len(deactivations)} deactivations (Gaiia has stale active assignment)")
    print(f"  {len(skipped)} sessions we couldn't resolve to an account (skipped)")
    if upserts:
        print("\n--- UPSERTS ---")
        for u in upserts[:200]:
            s = u["session"]
            acct = accounts.get(u["account_id"], {})
            cur = ""
            if u["current"]:
                c = u["current"][0]
                ca = (c.get("account") or {}).get("displayName") or "<none>"
                cur = f"   was: ip={c['ipAddress']} mac={c['macAddress']} acct={ca!r}"
            print(f"  {s['router']:9} {s['ip']:18} {fmt_mac(s['mac']):17} "
                  f"src={s['source']:5} user={(s['username'] or s['host'])!r:25} -> "
                  f"#{acct.get('readable')} {acct.get('name')!r} via={u['via']}")
            if cur: print(cur)
        if len(upserts) > 200:
            print(f"  ... and {len(upserts)-200} more")
    if deactivations:
        print("\n--- DEACTIVATIONS ---")
        for d in deactivations[:200]:
            c = d["current"]
            ca = (c.get("account") or {}).get("displayName") or "<none>"
            print(f"  ip={d['ip']} mac={c['macAddress']} acct={ca!r}")
        if len(deactivations) > 200:
            print(f"  ... and {len(deactivations)-200} more")
    if skipped:
        print("\n--- SKIPPED (couldn't resolve) ---")
        for s, why in skipped[:50]:
            print(f"  {s['router']:9} {s['ip']:18} {fmt_mac(s['mac']):17} "
                  f"src={s['source']:5} user={(s['username'] or s['host'])!r:25} reason={why}")
        if len(skipped) > 50:
            print(f"  ... and {len(skipped)-50} more")

    if not args.apply:
        print("\n(dry-run; pass --apply to execute)")
        return

    print("\n=== Applying ===")
    with GaiiaClient(timezone="America/Chicago") as g:
        do_upserts(g, upserts)
        if not args.no_deactivate:
            do_deactivates(g, deactivations)


if __name__ == "__main__":
    main()