from __future__ import annotations

import time
from dataclasses import dataclass
from typing import List

from k4s.core.context_store import Context, ContextStore
from k4s.core.executor import Executor
from k4s.ui.ui import Ui


class ClusterError(Exception):
    """Cluster operation error."""


@dataclass(frozen=True)
class ClusterNode:
    context: Context
    role: str  # "control-plane" | "worker"

def build_nodes_from_contexts(
    store: ContextStore, *, control_planes: list[str], workers: list[str]
) -> List[ClusterNode]:
    if not control_planes:
        raise ClusterError("At least one --control-plane is required.")

    all_names = list(control_planes) + list(workers)
    if len(set(all_names)) != len(all_names):
        raise ClusterError("Duplicate node context names are not allowed.")

    nodes: List[ClusterNode] = []
    for name in control_planes:
        ctx = store.resolve(name)
        if not ctx.host:
            raise ClusterError(f"Context '{name}' has no host.")
        nodes.append(ClusterNode(context=ctx, role="control-plane"))

    for name in workers:
        ctx = store.resolve(name)
        if not ctx.host:
            raise ClusterError(f"Context '{name}' has no host.")
        nodes.append(ClusterNode(context=ctx, role="worker"))

    return nodes


def _use_sudo(ctx: Context) -> bool:
    # Prefer running as root for infra operations. If not root, attempt sudo -n.
    return (ctx.username or "").lower() not in {"root"}


class Rke2ClusterCreator:
    def __init__(self, ui: Ui):
        self.ui = ui

    def create(
        self,
        *,
        name: str,
        nodes: List[ClusterNode],
        cni: str,
        tls_san: list[str] | None = None,
        dry_run: bool = False,
    ) -> str:
        """
        Create an RKE2 cluster.

        Returns:
            Path to written kubeconfig (local path) if fetched, otherwise "".
        """
        primary_cp = next(n for n in nodes if n.role == "control-plane")
        cp_nodes = [n for n in nodes if n.role == "control-plane"]
        worker_nodes = [n for n in nodes if n.role == "worker"]

        self.ui.info(f"Cluster: {name}")
        self.ui.info("Type: rke2")
        self.ui.info(f"Control-plane nodes: {len(cp_nodes)} | Worker nodes: {len(worker_nodes)}")
        self.ui.info(f"CNI: {cni}")

        if dry_run:
            self.ui.log("Dry-run mode. No changes will be applied.")
            self.ui.info("Planned steps:")
            self.ui.info(f"- Verify SSH connectivity to {len(nodes)} node(s)")
            self.ui.info(f"- Install rke2-server on {primary_cp.context.name} ({primary_cp.context.host})")
            if len(cp_nodes) > 1:
                self.ui.info(f"- Join {len(cp_nodes)-1} additional control-plane node(s)")
            if worker_nodes:
                self.ui.info(f"- Join {len(worker_nodes)} worker node(s)")
            self.ui.info("- Fetch kubeconfig from primary control-plane")
            return ""

        token = None

        with self.ui.step("Validating SSH connectivity"):
            self.ui.log("Ensuring all nodes are reachable before making changes.")
            for n in nodes:
                sudo = _use_sudo(n.context)
                with Executor(n.context.to_server_config()) as ex:
                    rc, out, err = ex.execute("whoami && hostname")
                    if rc != 0:
                        raise ClusterError(
                            f"SSH check failed for {n.context.name} ({n.context.host}): {err}"
                        )
                    if self.ui.options.verbosity >= 3 and out:
                        self.ui.debug(f"{n.context.name}: {out}")

        with self.ui.step(f"Installing rke2-server on {primary_cp.context.name}"):
            self.ui.log("Installing control-plane components.")
            sudo = _use_sudo(primary_cp.context)
            with Executor(primary_cp.context.to_server_config()) as ex:
                # Install script (idempotent-ish)
                if sudo:
                    cmd = "curl -sfL https://get.rke2.io | sudo -n sh -"
                else:
                    cmd = "curl -sfL https://get.rke2.io | sh -"
                rc, _, err = ex.execute(cmd, use_sudo=False)
                if rc != 0:
                    raise ClusterError(f"rke2 install failed on {primary_cp.context.host}: {err}")

                # Build config — always include the control-plane's public IP in
                # tls-san so kubeconfig exported to remote VMs works without TLS errors.
                san_entries: list[str] = [primary_cp.context.host] if primary_cp.context.host else []
                for extra in (tls_san or []):
                    if extra not in san_entries:
                        san_entries.append(extra)
                san_yaml = ""
                if san_entries:
                    san_yaml = "tls-san:\n" + "".join(f"  - {s}\n" for s in san_entries)
                config_yaml = f"cni: {cni}\n{san_yaml}"
                ex.execute("mkdir -p /etc/rancher/rke2", use_sudo=sudo)
                ex.write_remote_file_content("/etc/rancher/rke2/config.yaml", config_yaml, use_sudo=sudo)

                rc, _, err = ex.execute("systemctl enable rke2-server --now", use_sudo=sudo)
                if rc != 0:
                    raise ClusterError(f"Failed to start rke2-server on {primary_cp.host}: {err}")

                # Wait for token
                for _ in range(60):
                    rc, out, _ = ex.execute("cat /var/lib/rancher/rke2/server/node-token", use_sudo=sudo)
                    if rc == 0 and out.strip():
                        token = out.strip()
                        break
                    time.sleep(2)
                if not token:
                    raise ClusterError("Timed out waiting for rke2 node token on primary control-plane.")

        server_url = f"https://{primary_cp.context.host}:9345"

        for cp in cp_nodes[1:]:
            with self.ui.step(f"Joining control-plane {cp.context.name}"):
                self.ui.log("Adding an additional control-plane node.")
                sudo = _use_sudo(cp.context)
                with Executor(cp.context.to_server_config()) as ex:
                    if sudo:
                        cmd = "curl -sfL https://get.rke2.io | sudo -n sh -"
                    else:
                        cmd = "curl -sfL https://get.rke2.io | sh -"
                    rc, _, err = ex.execute(cmd, use_sudo=False)
                    if rc != 0:
                        raise ClusterError(f"rke2 install failed on {cp.context.host}: {err}")
                    cfg = f"server: {server_url}\n" f"token: {token}\n" f"cni: {cni}\n"
                    ex.execute("mkdir -p /etc/rancher/rke2", use_sudo=sudo)
                    ex.write_remote_file_content("/etc/rancher/rke2/config.yaml", cfg, use_sudo=sudo)
                    rc, _, err = ex.execute("systemctl enable rke2-server --now", use_sudo=sudo)
                    if rc != 0:
                        raise ClusterError(f"Failed to start rke2-server on {cp.host}: {err}")

        for w in worker_nodes:
            with self.ui.step(f"Joining worker {w.context.name}"):
                self.ui.log("Installing agent and joining it to the cluster.")
                sudo = _use_sudo(w.context)
                with Executor(w.context.to_server_config()) as ex:
                    if sudo:
                        cmd = "curl -sfL https://get.rke2.io | sudo -n env INSTALL_RKE2_TYPE=agent sh -"
                    else:
                        cmd = 'curl -sfL https://get.rke2.io | INSTALL_RKE2_TYPE="agent" sh -'
                    rc, _, err = ex.execute(cmd, use_sudo=False)
                    if rc != 0:
                        raise ClusterError(f"rke2 agent install failed on {w.context.host}: {err}")
                    cfg = f"server: {server_url}\n" f"token: {token}\n"
                    ex.execute("mkdir -p /etc/rancher/rke2", use_sudo=sudo)
                    ex.write_remote_file_content("/etc/rancher/rke2/config.yaml", cfg, use_sudo=sudo)
                    rc, _, err = ex.execute("systemctl enable rke2-agent --now", use_sudo=sudo)
                    if rc != 0:
                        raise ClusterError(f"Failed to start rke2-agent on {w.host}: {err}")

        # Fetch kubeconfig
        from pathlib import Path

        kube_dir = Path.home() / ".kube" / "k4s"
        kube_dir.mkdir(parents=True, exist_ok=True)
        local_kubeconfig = kube_dir / f"{name}.yaml"

        with self.ui.step("Fetching kubeconfig"):
            self.ui.log("Downloading kubeconfig from the primary control-plane.")
            sudo = _use_sudo(primary_cp.context)
            with Executor(primary_cp.context.to_server_config()) as ex:
                content = ex.read_remote_file_content("/etc/rancher/rke2/rke2.yaml", use_sudo=sudo)
            content = content.replace("127.0.0.1", primary_cp.context.host or "127.0.0.1")
            local_kubeconfig.write_text(content, encoding="utf-8")

        self.ui.success(f"Kubeconfig written: {str(local_kubeconfig)}")
        return str(local_kubeconfig)


def run_preflight(ui: Ui, *, nodes: List[ClusterNode]) -> None:
    """Run basic preflight checks on nodes."""
    rows: list[list[str]] = []
    with ui.step("Preflight checks"):
        ui.log("Collecting OS and dependency info on each node.")
        for n in nodes:
            sudo = _use_sudo(n.context)
            with Executor(n.context.to_server_config()) as ex:
                # OS
                rc, os_release, _ = ex.execute("cat /etc/os-release 2>/dev/null || true")
                os_name = "-"
                if rc == 0 and os_release:
                    for line in os_release.splitlines():
                        if line.startswith("PRETTY_NAME="):
                            os_name = line.split("=", 1)[1].strip().strip('"')
                            break

                # systemd
                rc, _, _ = ex.execute("command -v systemctl >/dev/null 2>&1 && echo yes || echo no")
                has_systemd = "yes" if rc == 0 else "no"

                # curl
                rc, _, _ = ex.execute("command -v curl >/dev/null 2>&1 && echo yes || echo no")
                has_curl = "yes" if rc == 0 else "no"

                # sudo -n (only if needed)
                sudo_ok = "n/a"
                if sudo:
                    rc, _, err = ex.execute("true", use_sudo=True)
                    sudo_ok = "yes" if rc == 0 else "no"
                    if rc != 0 and ui.options.verbosity >= 3 and err:
                        ui.debug(f"{n.name}: sudo error: {err}")

                # kernel
                rc, kernel, _ = ex.execute("uname -r || true")
                kernel = (kernel or "-").strip()

                rows.append(
                    [
                        n.context.name,
                        n.role,
                        n.context.host or "-",
                        os_name,
                        kernel,
                        has_systemd,
                        has_curl,
                        sudo_ok,
                    ]
                )

    ui.table(
        title="Preflight summary",
        columns=["node", "role", "host", "os", "kernel", "systemd", "curl", "sudo -n"],
        rows=rows,
    )

