#!/usr/bin/env python3
"""Reconcile Longhorn disk tags for the Titan longhorn storage classes.

The astreae/asteria storageclasses select Longhorn disks by tag. The current
nodes already have the right disk paths, but the tag fields can drift to empty
after node recovery. This job patches the live Longhorn Node CRs back to the
expected tags so PVC provisioning keeps working.
"""

from __future__ import annotations

import json
import os
import ssl
import urllib.request


LONGHORN_NS = "longhorn-system"
LONGHORN_API = "/apis/longhorn.io/v1beta2/namespaces/{namespace}/nodes"
DESIRED_DISK_TAGS = {
    "/mnt/astreae": ["astreae"],
    "/mnt/asteria": ["asteria", "cassandra-db", "cassandra-artifacts"],
    "/mnt/veles": [
        "veles-oceanus",
        "veles-db",
        "veles-artifacts",
        "cassandra-db",
        "cassandra-artifacts",
    ],
    "/mnt/veles-db": ["veles-oceanus", "veles-db"],
    "/mnt/veles-artifacts": ["veles-oceanus", "veles-artifacts"],
}
DESIRED_NODE_TAGS = {
    "titan-23": ["veles-oceanus"],
}
DESIRED_NODE_DISKS = {
    "titan-23": {
        "veles-oceanus": {
            "path": "/mnt/veles",
            "allowScheduling": True,
            "evictionRequested": False,
            "storageReserved": 0,
            "tags": [
                "veles-oceanus",
                "veles-db",
                "veles-artifacts",
                "cassandra-db",
                "cassandra-artifacts",
            ],
        }
    }
}
DISABLE_DEFAULT_DISK_NODES = {"titan-23"}


def api_base() -> str:
    host = os.environ.get("KUBERNETES_SERVICE_HOST")
    port = os.environ.get("KUBERNETES_SERVICE_PORT", "443")
    if not host:
        raise SystemExit("missing KUBERNETES_SERVICE_HOST")
    return f"https://{host}:{port}"


def token() -> str:
    path = "/var/run/secrets/kubernetes.io/serviceaccount/token"
    with open(path, "r", encoding="utf-8") as fh:
        return fh.read().strip()


def ca_context() -> ssl.SSLContext:
    cafile = "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"
    return ssl.create_default_context(cafile=cafile)


def request_json(method: str, path: str, body: dict | None = None) -> dict:
    req = urllib.request.Request(
        f"{api_base()}{path}",
        method=method,
        headers={
            "Authorization": f"Bearer {token()}",
            "Content-Type": "application/merge-patch+json",
            "Accept": "application/json",
        },
        data=None if body is None else json.dumps(body).encode("utf-8"),
    )
    with urllib.request.urlopen(req, context=ca_context(), timeout=20) as resp:
        payload = resp.read()
    return json.loads(payload) if payload else {}


def list_nodes() -> list[dict]:
    data = request_json("GET", LONGHORN_API.format(namespace=LONGHORN_NS))
    return data.get("items", [])


def merged_tags(current_tags: list[str], desired_tags: list[str]) -> list[str]:
    return sorted(dict.fromkeys([*current_tags, *desired_tags]))


def patch_node_tags(node_name: str, desired_tags: list[str]) -> None:
    body = {"spec": {"tags": desired_tags}}
    request_json(
        "PATCH",
        f"{LONGHORN_API.format(namespace=LONGHORN_NS)}/{node_name}",
        body=body,
    )


def patch_disk_tags(node_name: str, disk_name: str, desired_tags: list[str]) -> None:
    body = {"spec": {"disks": {disk_name: {"tags": desired_tags}}}}
    request_json(
        "PATCH",
        f"{LONGHORN_API.format(namespace=LONGHORN_NS)}/{node_name}",
        body=body,
    )


def patch_disks(node_name: str, disks: dict) -> None:
    body = {"spec": {"disks": disks}}
    request_json(
        "PATCH",
        f"{LONGHORN_API.format(namespace=LONGHORN_NS)}/{node_name}",
        body=body,
    )


def main() -> int:
    changed = 0
    skipped = 0

    for node in list_nodes():
        name = node.get("metadata", {}).get("name", "")
        desired_node_tags = DESIRED_NODE_TAGS.get(name)
        if desired_node_tags:
            current_node_tags = node.get("spec", {}).get("tags") or []
            next_node_tags = merged_tags(current_node_tags, desired_node_tags)
            if current_node_tags != next_node_tags:
                print(f"patching {name} node tags={current_node_tags!r} -> {next_node_tags!r}")
                patch_node_tags(name, next_node_tags)
                changed += 1
            else:
                skipped += 1

        spec_disks = node.get("spec", {}).get("disks", {}) or {}
        desired_disks = DESIRED_NODE_DISKS.get(name, {})
        missing_disks = {
            disk_name: disk_spec
            for disk_name, disk_spec in desired_disks.items()
            if disk_name not in spec_disks
        }
        if missing_disks:
            print(f"adding {name} disks={sorted(missing_disks)}")
            patch_disks(name, missing_disks)
            changed += len(missing_disks)
            spec_disks = {**spec_disks, **missing_disks}

        if name in DISABLE_DEFAULT_DISK_NODES:
            disable_patch = {}
            for disk_name, disk in spec_disks.items():
                disk_path = (disk.get("path") or "").rstrip("/")
                if disk_path == "/var/lib/longhorn" and disk.get("allowScheduling", True):
                    disable_patch[disk_name] = {"allowScheduling": False}
            if disable_patch:
                print(f"disabling default Longhorn scheduling on {name} disks={sorted(disable_patch)}")
                patch_disks(name, disable_patch)
                changed += len(disable_patch)

        for disk_name, disk in spec_disks.items():
            disk_path = disk.get("path")
            desired_disk_tags = DESIRED_DISK_TAGS.get(disk_path)
            if not desired_disk_tags:
                continue
            current_tags = disk.get("tags") or []
            if current_tags == desired_disk_tags:
                skipped += 1
                continue
            print(f"patching {name}:{disk_name} path={disk_path} tags={current_tags!r} -> {desired_disk_tags!r}")
            patch_disk_tags(name, disk_name, desired_disk_tags)
            changed += 1

    print(f"done: changed={changed} skipped={skipped}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
