#!/usr/bin/env python3
"""
nb_sync_one_device.py

Update a single NetBox device from Cloud by hostname.

Behavior
- Liveness gate via /v1/devices/{cloud_id}/liveness.
  * If online: keep ORIGINAL behavior (Cloud is SoT, including IP).
  * If offline: query Subsystem for the device and use Subsystem IP as fallback,
    then continue with normal NetBox updates.
- Updates NetBox fields from Cloud:
  custom_fields.fw_version ← firmwareVersion
  custom_fields.nodeName   ← nodeName
  custom_fields.sectorName ← sectorName
  custom_fields.smallCellName ← smallCellName
  serial ← serialNumber
- IP handling:
  If chosen IP (Cloud when online, Subsystem when offline) is valid (not None/""/"0.0.0.0"):
    ensure eth0, ensure/create IP, MOVE from other device if needed, set primary_ip4,
    then PRUNE all other IPs on this device (default).
  If chosen IP is placeholder/invalid:
    skip IP changes and do not prune.
- --chatty logs step-by-step to stderr; stdout remains one-line OK/FAIL.
- Logs custom field upgrade_cmd as `NB: upgrade_cmd=<value>` when --chatty.

Exit codes:
  0 = success
  1 = not found / missing data
  3 = network/HTTP error
  4 = NetBox update error
"""

import sys
import json
import argparse
from typing import Optional, Union, List, Dict, Any

import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

# ------------ HARD-CODED CONFIG (per Pavel) ------------
NB_URL = "http://netbox.gt-tiso.ikeja.co.za"
NB_TOKEN = "7648e4f5ee370cda7834682e61b47c2ee8e95623"

CLOUD_API_BASE = "https://cloud.ikeja.co.za/v1/devices"
CLOUD_BEARER = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJlbWFpbCI6InBhdmVsLmxAOGRldmljZXMuY29tIiwic3ViIjoyMiwiaWF0IjoxNzg5MTk2NTQ2LCJleHAiOjE3OTE3ODg1NDZ9.fQlcMWTL4uTKewd-AjlaHOdXSMJ2JOw2RfnKKF3weWk"


# NEW: Subsystem fallback (only used when Cloud liveness isConnected=false)
SUBSYSTEM_BASE = "https://subsystem.ikeja.co.za"
SUBSYSTEM_TOKEN = "HJL+&XRCoeHwh5?13@gxvg86qD#kQfgc"
SUBSYSTEM_OUTDOOR_ENDPOINT = "/customers/wave-devices/get-outdoor-devices"

REQ_TIMEOUT = 30
CHATTY = False

# ------------ HTTP utilities ------------
def _new_session() -> requests.Session:
    s = requests.Session()
    retries = Retry(
        total=3, connect=3, read=3, status=3,
        backoff_factor=0.5,
        status_forcelist=(429, 500, 502, 503, 504),
        allowed_methods=("GET", "POST", "PATCH", "PUT", "DELETE"),
        respect_retry_after_header=True,
    )
    adapter = HTTPAdapter(max_retries=retries, pool_connections=8, pool_maxsize=16)
    s.mount("http://", adapter)
    s.mount("https://", adapter)
    s.headers["Accept"] = "application/json"
    return s

S_NB = _new_session()
S_CL = _new_session()
S_SUB = _new_session()

# ------------ Logging / status helpers ------------
def log(msg: str):
    if CHATTY:
        print(msg, file=sys.stderr)

def die(code: int, msg: str):
    print(msg, file=sys.stdout if code == 0 else sys.stderr)
    raise SystemExit(code)

def _is_placeholder_ip(v: Any) -> bool:
    return v in (None, "", "0.0.0.0")

def _clean_ip(v: Any) -> Any:
    if isinstance(v, str):
        return v.strip()
    return v

# ------------ NetBox API ------------
class NetBox:
    def __init__(self, base: str, token: str):
        self.base = base.rstrip("/")
        self.token = token

    def _h(self):
        return {"Authorization": f"Token {self.token}", "Content-Type": "application/json"}

    def _url(self, path: str) -> str:
        return f"{self.base}{path}"

    def get_device_by_name(self, name: str) -> Optional[dict]:
        log(f"NB: lookup device by name {name}")
        r = S_NB.get(self._url("/api/dcim/devices/"), headers=self._h(),
                     params={"name": name}, timeout=REQ_TIMEOUT)
        if r.status_code != 200:
            die(3, f"FAIL NetBox GET devices name={name} HTTP={r.status_code}")
        res = r.json().get("results") or []
        return res[0] if res else None

    def get_device(self, dev_id: int) -> dict:
        log(f"NB: fetch device id={dev_id}")
        r = S_NB.get(self._url(f"/api/dcim/devices/{dev_id}/"), headers=self._h(), timeout=REQ_TIMEOUT)
        if r.status_code != 200:
            die(3, f"FAIL NetBox GET device id={dev_id} HTTP={r.status_code}")
        return r.json()

    def patch_device(self, dev_id: int, patch: dict) -> None:
        if not patch:
            log("NB: no device patch needed")
            return
        log(f"NB: patch device id={dev_id} keys={list(patch.keys())}")
        r = S_NB.patch(self._url(f"/api/dcim/devices/{dev_id}/"), headers=self._h(),
                       data=json.dumps(patch), timeout=REQ_TIMEOUT)
        if not (200 <= r.status_code < 300):
            die(4, f"FAIL patch device dev={dev_id} HTTP={r.status_code} body={r.text[:200]}")

    def get_iface_id(self, dev_id: int, name: str) -> Optional[int]:
        r = S_NB.get(self._url("/api/dcim/interfaces/"), headers=self._h(),
                     params={"device_id": dev_id, "name": name}, timeout=REQ_TIMEOUT)
        if r.status_code != 200:
            die(3, f"FAIL NetBox GET interfaces device_id={dev_id} HTTP={r.status_code}")
        res = r.json().get("results") or []
        return res[0]["id"] if res else None

    def ensure_eth0(self, dev_id: int) -> int:
        ifid = self.get_iface_id(dev_id, "eth0")
        log(f"NB: ensure eth0 (current id={ifid})")
        if ifid:
            return ifid
        payload = {"device": dev_id, "name": "eth0", "type": "1000base-t"}
        r = S_NB.post(self._url("/api/dcim/interfaces/"), headers=self._h(),
                      data=json.dumps(payload), timeout=REQ_TIMEOUT)
        if r.status_code == 201:
            return r.json()["id"]
        if r.status_code == 400:
            ifid = self.get_iface_id(dev_id, "eth0")
            if ifid:
                return ifid
        die(4, f"FAIL create eth0 HTTP={r.status_code} body={r.text[:200]}")

    def get_ip_by_addr(self, addr: str) -> Optional[dict]:
        r = S_NB.get(self._url("/api/ipam/ip-addresses/"), headers=self._h(),
                     params={"address": f"{addr}/32"}, timeout=REQ_TIMEOUT)
        if r.status_code != 200:
            die(3, f"FAIL NetBox GET ip-addresses addr={addr} HTTP={r.status_code}")
        res = r.json().get("results") or []
        return res[0] if res else None

    def create_ip_for_iface(self, addr: str, iface_id: int) -> int:
        payload = {
            "address": f"{addr}/32",
            "status": "active",
            "assigned_object_type": "dcim.interface",
            "assigned_object_id": iface_id,
        }
        r = S_NB.post(self._url("/api/ipam/ip-addresses/"), headers=self._h(),
                      data=json.dumps(payload), timeout=REQ_TIMEOUT)
        if r.status_code == 201:
            return r.json()["id"]
        die(4, f"FAIL create IP {addr} HTTP={r.status_code} body={r.text[:200]}")

    def assign_ip_to_iface(self, ip_id: int, iface_id: int) -> bool:
        log(f"NB: assign IP id={ip_id} -> iface={iface_id}")
        r = S_NB.patch(self._url(f"/api/ipam/ip-addresses/{ip_id}/"), headers=self._h(),
                       data=json.dumps({"assigned_object_type": "dcim.interface", "assigned_object_id": iface_id}),
                       timeout=REQ_TIMEOUT)
        return 200 <= r.status_code < 300

    def device_set_primary_ip4(self, dev_id: int, ip_id: int) -> None:
        log(f"NB: set primary_ip4 dev={dev_id} -> ip_id={ip_id}")
        r = S_NB.patch(self._url(f"/api/dcim/devices/{dev_id}/"), headers=self._h(),
                       data=json.dumps({"primary_ip4": ip_id}), timeout=REQ_TIMEOUT)
        if not (200 <= r.status_code < 300):
            die(4, f"FAIL set primary_ip4 dev={dev_id} ip_id={ip_id} HTTP={r.status_code} body={r.text[:200]}")

    def get_device_id_of_interface(self, iface_id: int) -> Optional[int]:
        r = S_NB.get(self._url(f"/api/dcim/interfaces/{iface_id}/"), headers=self._h(), timeout=REQ_TIMEOUT)
        if r.status_code == 200:
            return (r.json().get("device") or {}).get("id")
        return None

    def list_device_ips(self, dev_id: int) -> List[Dict]:
        r = S_NB.get(self._url("/api/ipam/ip-addresses/"), headers=self._h(),
                     params={"device_id": dev_id, "limit": 1000}, timeout=REQ_TIMEOUT)
        if r.status_code == 200:
            return r.json().get("results") or []
        return []

    def delete_ip(self, ip_id: int) -> None:
        log(f"NB: delete IP id={ip_id}")
        r = S_NB.delete(self._url(f"/api/ipam/ip-addresses/{ip_id}/"), headers=self._h(), timeout=REQ_TIMEOUT)
        if not (200 <= r.status_code < 300 or r.status_code == 204):
            die(4, f"FAIL delete IP id={ip_id} HTTP={r.status_code} body={r.text[:200]}")

# ------------ Cloud API ------------
class Cloud:
    def __init__(self, base: str, bearer: str):
        self.base = base.rstrip("/")
        self.bearer = bearer

    def _h(self):
        return {"Authorization": f"Bearer {self.bearer}", "Accept": "application/json"}

    def device_detail(self, cloud_id: Union[str, int]) -> dict:
        url = f"{self.base}/{cloud_id}"
        log(f"CL: fetch detail cloud_id={cloud_id}")
        r = S_CL.get(url, headers=self._h(), timeout=REQ_TIMEOUT)
        if r.status_code != 200:
            die(3, f"FAIL Cloud GET {cloud_id} HTTP={r.status_code}")
        return r.json()

    def device_liveness(self, cloud_id: Union[str, int]) -> dict:
        url = f"{self.base}/{cloud_id}/liveness"
        log(f"CL: fetch liveness cloud_id={cloud_id}")
        r = S_CL.get(url, headers=self._h(), timeout=REQ_TIMEOUT)
        if r.status_code != 200:
            die(3, f"FAIL Cloud GET {cloud_id}/liveness HTTP={r.status_code}")
        return r.json()

# ------------ Subsystem fallback ------------
class Subsystem:
    def __init__(self, base: str, token: str):
        self.base = base.rstrip("/")
        self.token = token

    def _h(self):
        return {"token": self.token, "Accept": "application/json"}

    def fetch_outdoor_devices(self) -> List[Dict[str, Any]]:
        url = f"{self.base}{SUBSYSTEM_OUTDOOR_ENDPOINT}"
        log("SUBSYS: fetch outdoor devices list")
        r = S_SUB.post(url, headers=self._h(), timeout=REQ_TIMEOUT)
        if r.status_code != 200:
            die(3, f"FAIL Subsystem POST get-outdoor-devices HTTP={r.status_code}")
        data = r.json()
        return data.get("device_list") or []

    def find_by_name(self, name: str) -> Optional[Dict[str, Any]]:
        devs = self.fetch_outdoor_devices()
        want = name.lower()
        for rec in devs:
            n = rec.get("name")
            if isinstance(n, str) and n.lower() == want:
                return rec
        return None

# ------------ Core ------------
def _primary_ip4_text(dev_json: dict) -> Optional[str]:
    p = dev_json.get("primary_ip4") or {}
    addr = p.get("address")
    if isinstance(addr, str) and addr.endswith("/32"):
        return addr[:-3]
    return addr

def run(hostname: str) -> None:
    nb = NetBox(NB_URL, NB_TOKEN)
    cl = Cloud(CLOUD_API_BASE, CLOUD_BEARER)
    sub = Subsystem(SUBSYSTEM_BASE, SUBSYSTEM_TOKEN)

    dev = nb.get_device_by_name(hostname)
    if not dev:
        die(1, f"FAIL {hostname} not found in NetBox")

    dev_id = dev["id"]
    dev_full = nb.get_device(dev_id)
    cf = dev_full.get("custom_fields") or {}

    upgrade_cmd = cf.get("upgrade_cmd")
    log(f"NB: upgrade_cmd={upgrade_cmd}")

    cloud_id = cf.get("cloud_id")
    if cloud_id in (None, "", "null"):
        die(1, f"FAIL {hostname} has no custom_fields.cloud_id in NetBox")

    # 1) Cloud liveness gate (ONLINE path stays as-is)
    live = cl.device_liveness(cloud_id)
    cloud_connected = bool(live.get("isConnected"))
    if cloud_connected:
        log(f"CL: liveness isConnected=true for cloud_id={cloud_id}; using Cloud ipAddress as source of truth")
    else:
        log(f"CL: liveness isConnected=false for cloud_id={cloud_id}; falling back to Subsystem for IP")

    current_nb_ip = _primary_ip4_text(dev_full)

    # 2) Cloud detail (still used for fw/node/sector/small/serial)
    d = cl.device_detail(cloud_id)
    fw = d.get("firmwareVersion") or d.get("version")
    cloud_ip = _clean_ip(d.get("ipAddress"))
    node = d.get("nodeName")
    sector = d.get("sectorName")
    small = d.get("smallCellName")
    serial = d.get("serialNumber")

    if not fw:
        die(1, f"FAIL {hostname} cloud_id={cloud_id}: missing firmwareVersion")

    # Choose IP source
    chosen_ip = cloud_ip
    ip_source = "cloud"

    if not cloud_connected:
        sub_rec = sub.find_by_name(hostname)
        if not sub_rec:
            # No subsystem record -> keep previous behavior: fail because offline and no fallback source
            die(1, f"FAIL {hostname} cloud_id={cloud_id} device offline and not found in Subsystem")
        sub_ip = _clean_ip(sub_rec.get("ip"))
        chosen_ip = sub_ip
        ip_source = "subsystem"
        log(f"SUBSYS: {hostname} ip={sub_ip} (fallback)")

    # 3) NetBox patch (idempotent)
    cf_patch = {}
    if cf.get("fw_version") != fw:
        cf_patch["fw_version"] = fw
        log(f"CF: fw_version -> {fw}")
    if node and cf.get("nodeName") != node:
        cf_patch["nodeName"] = node
        log(f"CF: nodeName -> {node}")
    if sector and cf.get("sectorName") != sector:
        cf_patch["sectorName"] = sector
        log(f"CF: sectorName -> {sector}")
    if small and cf.get("smallCellName") != small:
        cf_patch["smallCellName"] = small
        log(f"CF: smallCellName -> {small}")

    dev_patch = {}
    if serial and (dev_full.get("serial") != serial):
        dev_patch["serial"] = serial
        log(f"DEV: serial -> {serial}")
    if cf_patch:
        dev_patch["custom_fields"] = cf_patch
    nb.patch_device(dev_id, dev_patch)

    # 4) IP handling (skip if placeholder)
    chosen_ip = _clean_ip(chosen_ip)
    ip_is_placeholder = _is_placeholder_ip(chosen_ip)
    ip_out_for_status = current_nb_ip  # default to current NB IP

    if ip_is_placeholder:
        log(f"IP: {ip_source} reported placeholder '{chosen_ip}', skipping IP changes; keeping NetBox ip(s) as-is")
    else:
        iface_id = nb.ensure_eth0(dev_id)

        ip_rec = nb.get_ip_by_addr(chosen_ip)
        if ip_rec is None:
            log(f"IP: create new {chosen_ip} on iface {iface_id} (source={ip_source})")
            ip_id = nb.create_ip_for_iface(chosen_ip, iface_id)
        else:
            ip_id = ip_rec["id"]
            assigned_type = ip_rec.get("assigned_object_type") or ""
            assigned_id = ip_rec.get("assigned_object_id")
            if not assigned_type:
                if not nb.assign_ip_to_iface(ip_id, iface_id):
                    die(4, f"FAIL assign IP {chosen_ip} to iface {iface_id}")
            elif assigned_type == "dcim.interface":
                if str(assigned_id) != str(iface_id):
                    other_dev = nb.get_device_id_of_interface(assigned_id) if assigned_id else None
                    log(f"IP: moving {chosen_ip} from iface={assigned_id} dev={other_dev} -> iface={iface_id} dev={dev_id} (source={ip_source})")
                    # Clear old device primary if necessary
                    if other_dev:
                        r = S_NB.get(nb._url(f"/api/dcim/devices/{other_dev}/"), headers=nb._h(), timeout=REQ_TIMEOUT)
                        if r.status_code == 200:
                            old_primary_id = (r.json().get("primary_ip4") or {}).get("id")
                            if str(old_primary_id) == str(ip_id):
                                S_NB.patch(nb._url(f"/api/dcim/devices/{other_dev}/"),
                                           headers=nb._h(), data=json.dumps({"primary_ip4": None}),
                                           timeout=REQ_TIMEOUT)
                    if not nb.assign_ip_to_iface(ip_id, iface_id):
                        die(4, f"FAIL move IP {chosen_ip} to iface {iface_id}")
            else:
                die(4, f"FAIL IP {chosen_ip} assigned to {assigned_type}")

        nb.device_set_primary_ip4(dev_id, ip_id)
        ip_out_for_status = chosen_ip

        # PRUNE all other IPs on this device (default behavior)
        for rec in nb.list_device_ips(dev_id):
            rid = rec.get("id")
            if str(rid) == str(ip_id):
                continue
            addr = rec.get("address")
            log(f"IP: pruning stale {addr} (id={rid}) from device {dev_id}")
            nb.delete_ip(rid)

    print(f"OK {hostname} ip={ip_out_for_status or 'NONE'} fw={fw} node={node} sector={sector} small={small}")

# ------------ CLI ------------
if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Sync one NetBox device from Cloud by hostname")
    ap.add_argument("hostname", help="Device name in NetBox")
    ap.add_argument("--chatty", action="store_true", help="Verbose step-by-step logging to stderr")
    args = ap.parse_args()

    CHATTY = bool(args.chatty)
    try:
        run(args.hostname)
    except requests.RequestException as e:
        die(3, f"FAIL network error: {e}")