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"

        # Detect whether k4s itself is running on the primary control-plane node.
        # Executor treats host=localhost/127.0.0.1/None as local (no SSH).
        _cp_is_local = not Executor(primary_cp.context.to_server_config()).is_remote_execution

        raw_content = ""
        # Default private IP to the configured host; overridden by discovery below.
        private_ip: str = primary_cp.context.host or "127.0.0.1"

        with self.ui.step("Fetching kubeconfig"):
            self.ui.log("Downloading kubeconfig from the primary control-plane.")
            sudo = _use_sudo(primary_cp.context)
            if _cp_is_local:
                # k4s is running on the control-plane — read the file directly from disk.
                raw_content = Path("/etc/rancher/rke2/rke2.yaml").read_text(encoding="utf-8")
            else:
                with Executor(primary_cp.context.to_server_config()) as ex:
                    raw_content = ex.read_remote_file_content(
                        "/etc/rancher/rke2/rke2.yaml", use_sudo=sudo
                    )
                    # Discover the control-plane's private (internal) IP.
                    # Workers must use this address to reach the API server.
                    rc, priv_out, _ = ex.execute(
                        "hostname -I 2>/dev/null | awk '{print $1}'"
                    )
                    if rc == 0 and priv_out.strip():
                        private_ip = priv_out.strip()

            if not _cp_is_local:
                # k4s is running on a separate machine — write a kubeconfig that
                # uses the public IP so kubectl works from here.
                local_content = raw_content.replace(
                    "127.0.0.1", primary_cp.context.host or "127.0.0.1"
                )
                local_kubeconfig.write_text(local_content, encoding="utf-8")

        if not _cp_is_local:
            self.ui.success(f"Kubeconfig written: {str(local_kubeconfig)}")

        # Remove the control-plane NoSchedule taint when running worker-less so
        # regular workloads can be scheduled on the control-plane node.
        if not worker_nodes:
            with self.ui.step("Enabling pod scheduling on control-plane"):
                self.ui.log("No workers — removing NoSchedule taint so pods can be scheduled.")
                sudo = _use_sudo(primary_cp.context)
                kube = "/var/lib/rancher/rke2/bin/kubectl --kubeconfig /etc/rancher/rke2/rke2.yaml"
                with Executor(primary_cp.context.to_server_config()) as ex:
                    # Wait up to ~3 min for node to be registered in the API server.
                    for _ in range(36):
                        rc, out, _ = ex.execute(
                            f"{kube} get nodes --no-headers 2>/dev/null",
                            use_sudo=sudo,
                        )
                        if rc == 0 and out.strip():
                            break
                        time.sleep(5)
                    else:
                        self.ui.warning("Node not Ready after timeout; skipping taint removal.")
                        return str(local_kubeconfig)

                    rc, out, err = ex.execute(
                        f"{kube} taint nodes --all "
                        "node-role.kubernetes.io/control-plane:NoSchedule- 2>/dev/null || true",
                        use_sudo=sudo,
                    )
                    # Also remove the legacy master taint if present.
                    ex.execute(
                        f"{kube} taint nodes --all "
                        "node-role.kubernetes.io/master:NoSchedule- 2>/dev/null || true",
                        use_sudo=sudo,
                    )

        # Copy kubeconfig to all cluster nodes so kubectl works out of the box
        # for the context user on each node without any extra setup.
        #
        # Each node gets a kubeconfig pointing to the right API server address:
        #   - control-plane: 127.0.0.1  (API server is local)
        #   - worker:        private IP  (must cross the VPC to reach the CP)
        with self.ui.step("Copying kubeconfig to cluster nodes"):
            self.ui.log("Writing kubeconfig to ~/.kube/config on each node.")
            for n in nodes:
                if n.role == "control-plane":
                    node_content = raw_content  # keeps 127.0.0.1
                else:
                    node_content = raw_content.replace("127.0.0.1", private_ip)

                user = n.context.username or "root"
                home = f"/home/{user}" if user.lower() != "root" else "/root"
                kube_dir_remote = f"{home}/.kube"
                kube_path = f"{kube_dir_remote}/config"
                sudo = _use_sudo(n.context)
                _s = "sudo -n " if sudo else ""
                with Executor(n.context.to_server_config()) as ex:
                    ex.execute(f"{_s}mkdir -p {kube_dir_remote}", use_sudo=False)
                    ex.execute(f"{_s}chmod 700 {kube_dir_remote}", use_sudo=False)
                    if sudo:
                        ex.execute(
                            f"sudo -n chown {user}:{user} {kube_dir_remote}", use_sudo=False
                        )
                    rc, _, err = ex.execute(
                        f"{_s}tee {kube_path} > /dev/null", use_sudo=False, stdin_data=node_content
                    )
                    if rc != 0:
                        self.ui.warning(
                            f"Could not write kubeconfig to {n.context.name}: {err}"
                        )
                        continue
                    ex.execute(f"{_s}chmod 600 {kube_path}", use_sudo=False)
                    if sudo:
                        ex.execute(
                            f"sudo -n chown {user}:{user} {kube_path}", use_sudo=False
                        )

        return str(local_kubeconfig)


def run_preflight(ui: Ui, *, nodes: List[ClusterNode]) -> None:
    """Run basic preflight checks: SSH connectivity and required tools on each node."""
    with ui.step("Preflight checks"):
        ui.log("Checking SSH connectivity and required tools on each node.")
        for n in nodes:
            sudo = _use_sudo(n.context)
            with Executor(n.context.to_server_config()) as ex:
                rc, _, err = ex.execute("whoami && hostname")
                if rc != 0:
                    raise ClusterError(
                        f"SSH check failed for {n.context.name} ({n.context.host}): {err}"
                    )

                rc, _, _ = ex.execute(
                    "command -v curl >/dev/null && command -v systemctl >/dev/null"
                )
                if rc != 0:
                    raise ClusterError(
                        f"Missing required tools on {n.context.name}: "
                        "both curl and systemctl must be installed."
                    )

                if sudo:
                    rc, _, err = ex.execute("true", use_sudo=True)
                    if rc != 0:
                        raise ClusterError(
                            f"Passwordless sudo (sudo -n) not available on {n.context.name}: {err}"
                        )

