"""Persistent storage for connection contexts (VM and Kubernetes)."""
from __future__ import annotations

import base64
import os
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

import yaml

try:
    from cryptography.fernet import Fernet
    _HAS_CRYPTO = True
except ImportError:
    _HAS_CRYPTO = False


@dataclass(frozen=True)
class Context:
    """Connection context for k4s operations.

    A context is either a VM (SSH) connection or a Kubernetes
    (kubeconfig) reference, determined by the ``type`` field.
    """

    name: str
    type: str = "vm"  # "vm" | "k8s"

    # --- VM fields (SSH) ---
    host: Optional[str] = None
    username: Optional[str] = None
    port: int = 22
    identity_file: Optional[str] = None
    identity_passphrase: Optional[str] = None
    password: Optional[str] = None

    # --- K8s fields ---
    kubeconfig: Optional[str] = None
    kubectl_context: Optional[str] = None
    namespace: Optional[str] = None

    def to_server_config(self, *, host_override: Optional[str] = None) -> Dict[str, Any]:
        """Convert to a server_config dict understood by Executor (VM contexts only)."""
        host = host_override or self.host
        cfg: Dict[str, Any] = {"host": host or "localhost"}
        if self.username:
            cfg["user"] = self.username
        if self.port:
            cfg["port"] = self.port
        if self.identity_file:
            cfg["ssh_key_path"] = self.identity_file
        if self.identity_passphrase:
            cfg["ssh_key_passphrase"] = self.identity_passphrase
        if self.password:
            cfg["password"] = self.password
        return cfg

    def to_kubectl_args(self) -> List[str]:
        """Return kubectl CLI flags for this k8s context."""
        args: List[str] = []
        if self.kubeconfig:
            args += ["--kubeconfig", self.kubeconfig]
        # kubectl uses --context (NOT --kube-context).
        if self.kubectl_context:
            args += ["--context", self.kubectl_context]
        return args

    def to_helm_args(self) -> List[str]:
        """Return helm CLI flags for this k8s context."""
        args: List[str] = []
        if self.kubeconfig:
            args += ["--kubeconfig", self.kubeconfig]
        # helm uses --kube-context.
        if self.kubectl_context:
            args += ["--kube-context", self.kubectl_context]
        return args


class ContextStore:
    """Store contexts under ``~/.k4s/contexts.yml``.

    Passwords are encrypted at rest using Fernet symmetric encryption
    (requires the ``cryptography`` package).  A machine-local key is
    stored at ``~/.k4s/.key`` with 0600 permissions.
    """

    def __init__(self, path: Optional[str] = None):
        self.path = path or os.path.expanduser("~/.k4s/contexts.yml")
        self._key_path = os.path.join(os.path.dirname(self.path), ".key")

    # ---- encryption helpers ----

    def _get_fernet(self) -> Optional["Fernet"]:
        if not _HAS_CRYPTO:
            return None
        if os.path.exists(self._key_path):
            with open(self._key_path, "rb") as f:
                key = f.read().strip()
        else:
            key = Fernet.generate_key()
            parent = os.path.dirname(self._key_path)
            os.makedirs(parent, exist_ok=True)
            with open(self._key_path, "wb") as f:
                f.write(key)
            os.chmod(self._key_path, 0o600)
        return Fernet(key)

    def _encrypt_password(self, plaintext: str) -> str:
        f = self._get_fernet()
        if f is None:
            return plaintext
        token = f.encrypt(plaintext.encode("utf-8"))
        return "enc:" + base64.urlsafe_b64encode(token).decode("ascii")

    def _decrypt_password(self, stored: str) -> str:
        if not stored or not stored.startswith("enc:"):
            return stored  # plain text (legacy or no crypto)
        f = self._get_fernet()
        if f is None:
            return stored  # can't decrypt without cryptography
        token = base64.urlsafe_b64decode(stored[4:])
        return f.decrypt(token).decode("utf-8")

    # ---- file helpers ----

    def _ensure_parent_dir(self) -> None:
        parent = os.path.dirname(self.path)
        os.makedirs(parent, exist_ok=True)

    def _load_raw(self) -> Dict[str, Any]:
        if not os.path.exists(self.path):
            return {"current": None, "contexts": {}}
        with open(self.path, "r", encoding="utf-8") as f:
            data = yaml.safe_load(f) or {}
        return {
            "current": data.get("current"),
            "contexts": data.get("contexts", {}) or {},
        }

    def _save_raw(self, data: Dict[str, Any]) -> None:
        self._ensure_parent_dir()
        with open(self.path, "w", encoding="utf-8") as f:
            yaml.safe_dump(data, f, sort_keys=False)
        os.chmod(self.path, 0o600)

    def _parse_context(self, name: str, cfg: Dict[str, Any]) -> Context:
        """Build a Context from a raw YAML dict (handles legacy 'user' key)."""
        raw_pw = cfg.get("password")
        password = self._decrypt_password(raw_pw) if raw_pw else None
        raw_pp = cfg.get("identity_passphrase")
        identity_passphrase = self._decrypt_password(raw_pp) if raw_pp else None
        ctx_type = cfg.get("type", "vm")
        return Context(
            name=name,
            type=ctx_type,
            # VM fields
            host=cfg.get("host"),
            username=cfg.get("username") or cfg.get("user"),
            port=int(cfg.get("port", 22)),
            identity_file=cfg.get("identity_file") or cfg.get("ssh_key_path"),
            identity_passphrase=identity_passphrase,
            password=password,
            # K8s fields
            kubeconfig=cfg.get("kubeconfig"),
            kubectl_context=cfg.get("kubectl_context"),
            namespace=cfg.get("namespace"),
        )

    # ---- public API ----

    def list(self) -> Dict[str, Context]:
        raw = self._load_raw()
        out: Dict[str, Context] = {}
        for name, cfg in (raw.get("contexts") or {}).items():
            out[name] = self._parse_context(name, cfg)
        return out

    def get_current_name(self) -> Optional[str]:
        return self._load_raw().get("current")

    def set_current(self, name: str) -> None:
        raw = self._load_raw()
        if name not in (raw.get("contexts") or {}):
            raise KeyError(f"Context '{name}' not found.")
        raw["current"] = name
        self._save_raw(raw)

    def upsert(self, ctx: Context, *, make_current: bool = False) -> None:
        raw = self._load_raw()
        contexts = raw.setdefault("contexts", {})
        entry: Dict[str, Any] = {"type": ctx.type}

        if ctx.type == "k8s":
            # K8s-specific fields
            if ctx.kubeconfig:
                entry["kubeconfig"] = ctx.kubeconfig
            if ctx.kubectl_context:
                entry["kubectl_context"] = ctx.kubectl_context
            if ctx.namespace:
                entry["namespace"] = ctx.namespace
        else:
            # VM-specific fields
            entry["host"] = ctx.host
            entry["username"] = ctx.username
            entry["port"] = ctx.port
            entry["identity_file"] = ctx.identity_file
            if ctx.identity_passphrase:
                entry["identity_passphrase"] = self._encrypt_password(ctx.identity_passphrase)
            if ctx.password:
                entry["password"] = self._encrypt_password(ctx.password)

        contexts[ctx.name] = entry
        if make_current:
            raw["current"] = ctx.name
        self._save_raw(raw)

    def remove(self, name: str) -> None:
        raw = self._load_raw()
        contexts = raw.get("contexts") or {}
        if name not in contexts:
            raise KeyError(f"Context '{name}' not found.")
        del contexts[name]
        if raw.get("current") == name:
            raw["current"] = None
        raw["contexts"] = contexts
        self._save_raw(raw)

    def resolve(self, name: Optional[str]) -> Context:
        raw = self._load_raw()
        contexts = raw.get("contexts") or {}

        ctx_name = name or raw.get("current")
        if not ctx_name:
            raise KeyError(
                "No context selected.\n"
                "  Add one:\n"
                "    VM:   k4s context add <name> --host <ip> --username <user> -i <key> --current\n"
                "    K8s:  k4s context add <name> --type k8s --kubeconfig <path> --current\n"
                "  Or set:   k4s context use-context <name>\n"
                "  Or pass:  --context <name>"
            )
        cfg = contexts.get(ctx_name)
        if not cfg:
            raise KeyError(f"Context '{ctx_name}' not found.")

        return self._parse_context(ctx_name, cfg)
