"""Headscale REST API client for Hive Mode coordination.

Headscale (https://headscale.net) is a self-hosted Tailscale control server.
LLMHosts uses it as the coordination plane for Hive member nodes.

ACL policy: Headscale v0.22 does not expose a REST API for ACL/policy. Policy
is file-based (config policy.path → acls.hujson). Use generate_acl_policy()
to build the Tailscale-format ACL dict and HeadscaleClient.write_acl_policy_file()
to write it to a path that Headscale reads. Set HEADSCALE_ACL_PATH when running
llmhost hive create|join|leave so ACLs are synced after each change.

Usage:
    client = HeadscaleClient.from_env()
    if client:
        user = await client.create_user("my-hive")
        key = await client.create_preauth_key("my-hive")
        policy = generate_acl_policy([{"id": "h1", "headscale_user": "my-hive"}])
        client.write_acl_policy_file("/path/to/acls.hujson", policy)
"""

from __future__ import annotations

import json
import logging
import os
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

_DEFAULT_URL = "http://127.0.0.1:8080"
_PREAUTH_EXPIRY_DAYS = 180


def generate_acl_policy(
    hives: list[dict[str, Any]],
    *,
    admin_users: list[str] | None = None,
) -> dict[str, Any]:
    """Build Tailscale-format ACL policy: one group per Hive, same-Hive-only traffic.

    Headscale/Tailscale ACL format: groups, tagOwners, acls. Each Hive gets
    group "group:hive-<id>" with the single Headscale user name (one user per
    Hive). ACLs allow src group:hive-X → dst group:hive-X:* only; no cross-Hive
    accept. Deny by default (no catch-all).

    Parameters
    ----------
    hives:
        List of dicts with "id" (str) and "headscale_user" (str). Hives without
        headscale_user are skipped.
    admin_users:
        Optional list of Headscale user names that can reach all (group:admin).
        If provided, adds group:admin and accept admin -> *:*.

    Returns
    -------
    dict
        Tailscale ACL: {"groups": {...}, "tagOwners": {}, "acls": [...]}
    """
    groups: dict[str, list[str]] = {}
    acls: list[dict[str, Any]] = []

    for h in hives:
        hid = h.get("id")
        user = h.get("headscale_user")
        if not hid or not user:
            continue
        group_name = f"group:hive-{hid}"
        groups[group_name] = [user]
        acls.append({"action": "accept", "src": [group_name], "dst": [f"{group_name}:*"]})

    if admin_users:
        groups["group:admin"] = list(admin_users)
        acls.append({"action": "accept", "src": ["group:admin"], "dst": ["*:*"]})

    return {"groups": groups, "tagOwners": {}, "acls": acls}


class HeadscaleError(Exception):
    """Raised when the Headscale API returns an error."""


class HeadscaleClient:
    """Async HTTP client for the Headscale REST API.

    Authentication uses a bearer API key generated via:
        headscale apikeys create -e 365d
    """

    def __init__(self, url: str, api_key: str) -> None:
        self._url = url.rstrip("/")
        self._api_key = api_key

    @classmethod
    def from_env(cls) -> HeadscaleClient | None:
        """Create client from HEADSCALE_URL and HEADSCALE_API_KEY env vars.

        Returns None if env vars not set (Headscale integration disabled).
        """
        url = os.environ.get("HEADSCALE_URL", _DEFAULT_URL)
        api_key = os.environ.get("HEADSCALE_API_KEY", "")
        if not api_key:
            return None
        return cls(url, api_key)

    async def create_user(self, name: str) -> dict[str, Any]:
        """Create a Headscale user (one per Hive).

        Parameters
        ----------
        name:
            Headscale user name (e.g. "my-hive"). Must be unique.

        Returns
        -------
        dict
            Headscale user object: {"id": "1", "name": "...", "createdAt": "..."}
        """
        import httpx

        async with httpx.AsyncClient(headers=self._auth_headers()) as client:
            resp = await client.post(
                f"{self._url}/api/v1/user",
                json={"name": name},
            )
            if resp.status_code == 409:
                # User already exists — fetch and return it
                return await self.get_user(name)
            self._raise_for_status(resp, f"create user '{name}'")
            return dict(resp.json()["user"])

    async def get_user(self, name: str) -> dict[str, Any]:
        """Fetch a Headscale user by name."""
        import httpx

        async with httpx.AsyncClient(headers=self._auth_headers()) as client:
            resp = await client.get(f"{self._url}/api/v1/user/{name}")
            self._raise_for_status(resp, f"get user '{name}'")
            return dict(resp.json()["user"])

    async def create_preauth_key(
        self,
        user: str,
        *,
        reusable: bool = True,
        expiration_days: int = _PREAUTH_EXPIRY_DAYS,
    ) -> str:
        """Create a reusable pre-auth key for nodes to join the tailnet.

        Parameters
        ----------
        user:
            Headscale user name (matches the Hive user).
        reusable:
            Allow multiple nodes to use the same key (default: True).
        expiration_days:
            Days until key expires (default: 180).

        Returns
        -------
        str
            The pre-auth key string (long random hex).
        """
        import httpx

        expiry = (datetime.now(tz=timezone.utc) + timedelta(days=expiration_days)).strftime("%Y-%m-%dT%H:%M:%SZ")

        async with httpx.AsyncClient(headers=self._auth_headers()) as client:
            resp = await client.post(
                f"{self._url}/api/v1/preauthkey",
                json={
                    "user": user,
                    "reusable": reusable,
                    "ephemeral": False,
                    "expiration": expiry,
                    "aclTags": [],
                },
            )
            self._raise_for_status(resp, f"create pre-auth key for user '{user}'")
            return str(resp.json()["preAuthKey"]["key"])

    async def list_nodes(self, user: str) -> list[dict[str, Any]]:
        """List all nodes registered under a Headscale user."""
        import httpx

        async with httpx.AsyncClient(headers=self._auth_headers()) as client:
            resp = await client.get(f"{self._url}/api/v1/node", params={"user": user})
            self._raise_for_status(resp, f"list nodes for user '{user}'")
            return list(resp.json().get("nodes", []))

    async def health(self) -> bool:
        """Return True if Headscale is reachable and healthy."""
        import httpx

        try:
            async with httpx.AsyncClient() as client:
                resp = await client.get(f"{self._url}/health", timeout=3.0)
                data: dict[str, Any] = resp.json()
                return data.get("status") == "pass"
        except Exception:
            return False

    def _auth_headers(self) -> dict[str, str]:
        return {"Authorization": f"Bearer {self._api_key}"}

    def write_acl_policy_file(self, path: str | os.PathLike[str], policy: dict[str, Any]) -> None:
        """Write ACL policy dict to a file (huJSON or JSON) for Headscale to load.

        Headscale does not provide a policy REST API in v0.22; policy is file-based.
        After writing, reload Headscale (e.g. systemctl reload headscale) or use
        a file watcher so Headscale picks up changes.

        Parameters
        ----------
        path:
            File path (e.g. config/acls.hujson). Parent directory must exist.
        policy:
            Tailscale-format ACL dict from generate_acl_policy().
        """
        raw = json.dumps(policy, indent=2, default=str)
        Path(path).write_text(raw, encoding="utf-8")
        logger.debug("Wrote ACL policy to %s", path)

    @staticmethod
    def _raise_for_status(resp: Any, action: str) -> None:
        if resp.status_code >= 400:
            raise HeadscaleError(f"Headscale API error during {action}: {resp.status_code} {resp.text}")
