"""Managed Talos config persistence for SDK and CLI."""

from __future__ import annotations

import base64
import json
import os
import time
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any

import yaml

if TYPE_CHECKING:
    from omni.client.async_client import AsyncOmniClient

KEY_DEFAULT_CLUSTER = "talos.default_cluster"
KEY_LAST_CLUSTER_BY_INSTANCE = "talos.last_cluster_by_instance"
KEY_TALOSCONFIG_PATH_BY_INSTANCE = "talos.talosconfig_path_by_instance"
KEY_SIDEROV1_KEYS_DIR_BY_INSTANCE = "talos.siderov1_keys_dir_by_instance"
KEY_CONTEXT_LAST_REFRESH = "talos.context_last_refresh"
KEY_LAST_NODES_BY_SCOPE = "talos.last_nodes_by_scope"
KEY_LAST_ENDPOINTS_BY_SCOPE = "talos.last_endpoints_by_scope"
DEFAULT_REFRESH_TTL_SECONDS = 300


@dataclass(slots=True)
class TalosConfigResolution:
    path: Path
    context: str
    cluster: str | None
    refreshed: bool


def _decode_bytes_field(payload: dict[str, Any], key: str) -> bytes:
    value = payload.get(key)
    if isinstance(value, bytes):
        return value
    if not isinstance(value, str):
        return b""
    try:
        return base64.b64decode(value, validate=True)
    except Exception:
        return value.encode("utf-8")


def _coerce_yaml_mapping(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {}
    try:
        loaded = yaml.safe_load(path.read_text(encoding="utf-8"))
    except Exception:
        return {}
    return loaded if isinstance(loaded, dict) else {}


class TalosConfigManager:
    """Refresh and persist talosconfig contexts for an Omni client instance."""

    def __init__(self, client: AsyncOmniClient) -> None:
        self.client = client
        self.state = client.state
        self.instance = client.instance_name

    def default_path(self) -> Path:
        if self.state.path is not None:
            base_dir = self.state.path.parent
        else:
            base_dir = Path.home() / ".config" / "omnisdk"
        return (base_dir / "talosconfigs" / f"{self.instance}.yaml").expanduser().resolve()

    def default_siderov1_keys_dir(self) -> Path:
        if self.state.path is not None:
            base_dir = self.state.path.parent
        else:
            base_dir = Path.home() / ".config" / "omnisdk"
        return (base_dir / "siderov1-keys").expanduser().resolve()

    def resolve_path(self, *, explicit: str | Path | None = None) -> Path:
        if explicit is not None:
            path = Path(explicit).expanduser().resolve()
            self._set_path_for_instance(path)
            return path

        paths = self.state.get_json_preference(KEY_TALOSCONFIG_PATH_BY_INSTANCE)
        existing = paths.get(self.instance)
        if existing:
            return Path(existing).expanduser().resolve()

        default = self.default_path()
        self._set_path_for_instance(default)
        return default

    def resolve_siderov1_keys_dir(self, *, explicit: str | Path | None = None) -> Path:
        if explicit is not None:
            path = Path(explicit).expanduser().resolve()
            self._set_siderov1_keys_dir_for_instance(path)
            return path

        paths = self.state.get_json_preference(KEY_SIDEROV1_KEYS_DIR_BY_INSTANCE)
        existing = paths.get(self.instance)
        if existing:
            return Path(existing).expanduser().resolve()

        default = self.default_siderov1_keys_dir()
        self._set_siderov1_keys_dir_for_instance(default)
        return default

    def resolve_cluster(self, explicit: str | None = None) -> str | None:
        if explicit:
            return explicit

        last_map = self.state.get_json_preference(KEY_LAST_CLUSTER_BY_INSTANCE)
        if self.instance in last_map:
            return last_map[self.instance]

        talos_default = self.state.get_preference(KEY_DEFAULT_CLUSTER)
        if talos_default:
            return talos_default

        global_default = self.state.get_preference("default_cluster")
        if global_default:
            return global_default

        return self.client.cluster

    def context_name(self, cluster: str | None) -> str:
        cluster_name = cluster if cluster else "generic"
        return f"omni/{self.instance}/{cluster_name}"

    def scope_key(self, cluster: str | None) -> str:
        return f"{self.instance}:{cluster or 'generic'}"

    def _set_path_for_instance(self, path: Path) -> None:
        paths = self.state.get_json_preference(KEY_TALOSCONFIG_PATH_BY_INSTANCE)
        paths[self.instance] = str(path)
        self.state.set_json_preference(KEY_TALOSCONFIG_PATH_BY_INSTANCE, paths)

    def _set_siderov1_keys_dir_for_instance(self, path: Path) -> None:
        paths = self.state.get_json_preference(KEY_SIDEROV1_KEYS_DIR_BY_INSTANCE)
        paths[self.instance] = str(path)
        self.state.set_json_preference(KEY_SIDEROV1_KEYS_DIR_BY_INSTANCE, paths)

    def set_default_cluster(self, cluster: str) -> None:
        self.state.set_preference(KEY_DEFAULT_CLUSTER, cluster)

    def remember_cluster(self, cluster: str) -> None:
        self._set_last_cluster_for_instance(cluster)

    def _set_last_cluster_for_instance(self, cluster: str) -> None:
        clusters = self.state.get_json_preference(KEY_LAST_CLUSTER_BY_INSTANCE)
        clusters[self.instance] = cluster
        self.state.set_json_preference(KEY_LAST_CLUSTER_BY_INSTANCE, clusters)

    def set_talosconfig_path_for_instance(self, path: str | Path) -> Path:
        resolved = Path(path).expanduser().resolve()
        self._set_path_for_instance(resolved)
        return resolved

    def set_last_targets(
        self,
        *,
        cluster: str | None,
        nodes: str | None = None,
        endpoints: str | None = None,
    ) -> None:
        key = self.scope_key(cluster)

        if nodes:
            nodes_map = self.state.get_json_preference(KEY_LAST_NODES_BY_SCOPE)
            nodes_map[key] = nodes
            self.state.set_json_preference(KEY_LAST_NODES_BY_SCOPE, nodes_map)

        if endpoints:
            endpoints_map = self.state.get_json_preference(KEY_LAST_ENDPOINTS_BY_SCOPE)
            endpoints_map[key] = endpoints
            self.state.set_json_preference(KEY_LAST_ENDPOINTS_BY_SCOPE, endpoints_map)

    def get_last_targets(self, *, cluster: str | None) -> tuple[str | None, str | None]:
        key = self.scope_key(cluster)
        nodes_map = self.state.get_json_preference(KEY_LAST_NODES_BY_SCOPE)
        endpoints_map = self.state.get_json_preference(KEY_LAST_ENDPOINTS_BY_SCOPE)
        return nodes_map.get(key), endpoints_map.get(key)

    async def resolve_nodes(self, nodes: str, *, cluster: str | None = None) -> str:
        selectors = [item.strip() for item in nodes.split(",") if item.strip()]
        if not selectors:
            return nodes

        alias_map = await self._node_alias_map(cluster=cluster)
        resolved = [alias_map.get(selector.lower(), selector) for selector in selectors]
        return ",".join(resolved)

    async def _node_alias_map(self, *, cluster: str | None = None) -> dict[str, str]:
        try:
            payload = await self.client.resources.list(
                {"namespace": "default", "type": "MachineStatuses.omni.sidero.dev"}
            )
        except Exception:
            return {}

        items = payload.get("items")
        if not isinstance(items, list):
            return {}

        aliases: dict[str, str] = {}
        collisions: set[str] = set()

        for raw in items:
            item = self._coerce_item(raw)
            if item is None:
                continue

            metadata = item.get("metadata")
            if not isinstance(metadata, dict):
                continue

            node_id = metadata.get("id")
            if not isinstance(node_id, str) or not node_id.strip():
                continue
            node_id = node_id.strip()

            spec = item.get("spec")
            if not isinstance(spec, dict):
                spec = {}

            item_cluster = spec.get("cluster")
            if not isinstance(item_cluster, str) or not item_cluster:
                labels = metadata.get("labels")
                if isinstance(labels, dict):
                    raw_label = labels.get("omni.sidero.dev/cluster")
                    if isinstance(raw_label, str):
                        item_cluster = raw_label

            if cluster and isinstance(item_cluster, str) and item_cluster and item_cluster != cluster:
                continue

            candidates = {node_id}
            network = spec.get("network")
            if isinstance(network, dict):
                hostname = network.get("hostname")
                if isinstance(hostname, str) and hostname.strip():
                    candidates.add(hostname.strip())

            for key in ("hostname", "node_name", "nodename", "machine_name", "name"):
                value = spec.get(key)
                if isinstance(value, str) and value.strip():
                    candidates.add(value.strip())

            for alias in candidates:
                alias_key = alias.lower()
                previous = aliases.get(alias_key)
                if previous is not None and previous != node_id:
                    collisions.add(alias_key)
                else:
                    aliases[alias_key] = node_id

        for alias_key in collisions:
            aliases.pop(alias_key, None)

        return aliases

    @staticmethod
    def _coerce_item(raw: Any) -> dict[str, Any] | None:
        if isinstance(raw, dict):
            return raw
        if isinstance(raw, str):
            try:
                loaded = json.loads(raw)
            except json.JSONDecodeError:
                return None
            if isinstance(loaded, dict):
                return loaded
        return None

    def summary(self) -> dict[str, Any]:
        paths = self.state.get_json_preference(KEY_TALOSCONFIG_PATH_BY_INSTANCE)
        clusters = self.state.get_json_preference(KEY_LAST_CLUSTER_BY_INSTANCE)
        current_cluster = clusters.get(self.instance) or self.state.get_preference(KEY_DEFAULT_CLUSTER)
        return {
            "instance": self.instance,
            "default_ttl_seconds": DEFAULT_REFRESH_TTL_SECONDS,
            "default_cluster": self.state.get_preference(KEY_DEFAULT_CLUSTER),
            "last_cluster_for_instance": clusters.get(self.instance),
            "effective_cluster": current_cluster,
            "talosconfig_path_for_instance": paths.get(self.instance),
            "service_account_key_available": self._service_account_key_b64() is not None,
        }

    def service_account_key_b64(self) -> str | None:
        return self._service_account_key_b64()

    def _service_account_key_b64(self) -> str | None:
        cfg_key = self.client.instance_config.auth.service_account_key
        if isinstance(cfg_key, str) and cfg_key.strip():
            return cfg_key.strip()

        for alias in self.client.instance_config.auth.service_account_env_aliases:
            value = os.getenv(alias)
            if value and value.strip():
                return value.strip()

        return None

    def _service_account_armored_key(self) -> str | None:
        value_b64 = self._service_account_key_b64()
        if value_b64 is None:
            return None

        try:
            raw = base64.b64decode(value_b64.encode("utf-8"), validate=True)
            payload = json.loads(raw.decode("utf-8"))
        except Exception:
            return None

        if not isinstance(payload, dict):
            return None

        armored = payload.get("pgp_key")
        if not isinstance(armored, str) or not armored.strip():
            return None

        return armored

    def _context_identity(self, *, talosconfig_path: Path, context: str) -> str | None:
        config = _coerce_yaml_mapping(talosconfig_path)
        contexts = config.get("contexts")
        if not isinstance(contexts, dict):
            return None

        selected = contexts.get(context)
        if not isinstance(selected, dict):
            return None

        auth = selected.get("auth")
        if not isinstance(auth, dict):
            return None

        siderov1 = auth.get("siderov1")
        if not isinstance(siderov1, dict):
            return None

        identity = siderov1.get("identity")
        if not isinstance(identity, str):
            return None
        identity = identity.strip()
        return identity if identity else None

    def ensure_siderov1_keys_dir(
        self,
        *,
        talosconfig_path: Path,
        context: str,
        path: str | Path | None = None,
    ) -> Path | None:
        identity = self._context_identity(talosconfig_path=talosconfig_path, context=context)
        if identity is None:
            return None

        armored = self._service_account_armored_key()
        if armored is None:
            return None

        keys_dir = self.resolve_siderov1_keys_dir(explicit=path)
        key_path = (keys_dir / f"{context}-{identity}.pgp").expanduser().resolve()
        key_path.parent.mkdir(parents=True, exist_ok=True)

        content = armored if armored.endswith("\n") else f"{armored}\n"
        key_path.write_text(content, encoding="utf-8")
        try:
            os.chmod(key_path, 0o600)
        except OSError:
            pass

        return keys_dir

    def _get_last_refresh(self, context: str) -> float:
        key = f"{self.instance}:{context}"
        mapping = self.state.get_json_preference(KEY_CONTEXT_LAST_REFRESH)
        raw = mapping.get(key)
        if raw is None:
            return 0.0
        try:
            return float(raw)
        except ValueError:
            return 0.0

    def _set_last_refresh(self, context: str, value: float) -> None:
        key = f"{self.instance}:{context}"
        mapping = self.state.get_json_preference(KEY_CONTEXT_LAST_REFRESH)
        mapping[key] = str(value)
        self.state.set_json_preference(KEY_CONTEXT_LAST_REFRESH, mapping)

    async def ensure(
        self,
        *,
        cluster: str | None = None,
        break_glass: bool = False,
        force_refresh: bool = False,
        ttl_seconds: int = DEFAULT_REFRESH_TTL_SECONDS,
        path: str | Path | None = None,
    ) -> TalosConfigResolution:
        resolved_path = self.resolve_path(explicit=path)
        resolved_cluster = self.resolve_cluster(explicit=cluster)
        context = self.context_name(resolved_cluster)
        now = time.time()

        current = _coerce_yaml_mapping(resolved_path)
        current_contexts = current.get("contexts")
        has_context = isinstance(current_contexts, dict) and context in current_contexts
        stale = (now - self._get_last_refresh(context)) > max(1, ttl_seconds)
        refresh = force_refresh or (not resolved_path.exists()) or (not has_context) or stale

        if refresh:
            request = {"raw": False, "break_glass": break_glass}
            metadata: dict[str, str] | None = {"cluster": resolved_cluster} if resolved_cluster else None
            if force_refresh:
                if metadata is None:
                    metadata = {}
                metadata["x-omnictl-refresh"] = str(now)
            payload = await self.client.management.talosconfig(request, metadata=metadata)
            talosconfig = _decode_bytes_field(payload, "talosconfig")
            fetched_text = talosconfig.decode("utf-8", errors="replace")
            fetched = yaml.safe_load(fetched_text)
            if not isinstance(fetched, dict):
                raise ValueError("management talosconfig response is not a YAML mapping")
            fetched_contexts = fetched.get("contexts")
            if not isinstance(fetched_contexts, dict) or not fetched_contexts:
                raise ValueError("management talosconfig does not include any contexts")

            source_context_name = str(fetched.get("context", ""))
            if source_context_name not in fetched_contexts:
                source_context_name = next(iter(fetched_contexts))

            source_context = fetched_contexts[source_context_name]
            if not isinstance(source_context, dict):
                raise ValueError("management talosconfig context payload is invalid")

            merged = current if isinstance(current, dict) else {}
            merged_contexts = merged.get("contexts")
            if not isinstance(merged_contexts, dict):
                merged_contexts = {}
            merged_contexts[context] = source_context
            merged["contexts"] = merged_contexts
            merged["context"] = context

            resolved_path.parent.mkdir(parents=True, exist_ok=True)
            resolved_path.write_text(
                yaml.safe_dump(merged, sort_keys=False, default_flow_style=False),
                encoding="utf-8",
            )

            self._set_last_refresh(context, now)
            if resolved_cluster:
                self._set_last_cluster_for_instance(resolved_cluster)

            return TalosConfigResolution(
                path=resolved_path,
                context=context,
                cluster=resolved_cluster,
                refreshed=True,
            )

        return TalosConfigResolution(
            path=resolved_path,
            context=context,
            cluster=resolved_cluster,
            refreshed=False,
        )
