"""CLI commands for managing connection contexts."""
import shlex
import subprocess

import click

from k4s.cli.state import CliState
from k4s.cli.errors import exit_with_error
from k4s.cli.verbosity import apply_command_ui_overrides
from k4s.core.context_store import Context
from k4s.core.executor import Executor


class AliasGroup(click.Group):
    """Click group that supports command aliases."""

    _aliases: dict = {
        "ls": "list",
        "rm": "remove",
    }

    def get_command(self, ctx, cmd_name):
        resolved = self._aliases.get(cmd_name, cmd_name)
        return super().get_command(ctx, resolved)

class ArgsFirstUsageCommand(click.Command):
    """Click command that prints usage as `NAME [OPTIONS]` (args first)."""

    def get_usage(self, ctx) -> str:
        usage = super().get_usage(ctx)
        # Default Click format is: "Usage: ... [OPTIONS] NAME"
        return usage.replace("[OPTIONS] NAME", "NAME [OPTIONS]")

    def format_usage(self, ctx, formatter) -> None:
        # Click's help output doesn't always use get_usage() for the usage line
        # (depending on Click version / command wiring), so enforce the same
        # `NAME [OPTIONS]` ordering here as well.
        pieces = self.collect_usage_pieces(ctx)
        if "[OPTIONS]" in pieces and "NAME" in pieces:
            pieces = [p for p in pieces if p not in ("[OPTIONS]", "NAME")]
            pieces.insert(0, "[OPTIONS]")
            pieces.insert(0, "NAME")
        formatter.write(f"Usage: {ctx.command_path} {' '.join(pieces)}\n")


@click.group(cls=AliasGroup)
@click.pass_context
def context(ctx):
    """Manage connection contexts (VM or Kubernetes)."""
    pass


@context.command("add", short_help="Add or update a connection context.", cls=ArgsFirstUsageCommand)
@click.argument("name", metavar="NAME", required=False, default=None)
@click.option("--type", "ctx_type", type=click.Choice(["vm", "k8s"], case_sensitive=False), default="vm", show_default=True, help="Context type.")
# VM options
@click.option("-h", "--host", default=None, help="Host (IP/hostname). Required for VM contexts.")
@click.option("-u", "--username", default=None, help="SSH username. Required for VM contexts.")
@click.option("-P", "--port", default=22, show_default=True, type=int, help="SSH port.")
@click.option("-i", "--identity-file", help="SSH private key path.")
@click.option("-ip", "--identity-passphrase", default=None, help="SSH key passphrase (stored encrypted).")
@click.option("-aip", "--ask-identity-passphrase", is_flag=True, help="Prompt for SSH key passphrase (stored encrypted).")
@click.option("-p", "--password", help="SSH password (discouraged; prefer keys).")
# K8s options
@click.option("--kubeconfig", default=None, help="Path to kubeconfig file (for k8s contexts).")
@click.option("--kubectl-context", default=None, help="Context name within the kubeconfig file (for k8s contexts).")
@click.option("-n", "--namespace", default=None, help="Default namespace for Helm operations (for k8s contexts).")
# Common
@click.option("--current", is_flag=True, default=False, help="Set as current context.")
@click.pass_context
def add(ctx, name, ctx_type, host, username, port, identity_file, identity_passphrase,
        ask_identity_passphrase, password, kubeconfig, kubectl_context, namespace, current):
    """Add or update a connection context.

    NAME is a label for this connection (e.g. 'prod', 'staging-k8s').
    Use --current to set it as the active context for subsequent commands.

    \b
    Examples (VM):
      k4s context add prod --host 10.0.1.5 --username ubuntu -i ~/.ssh/id_rsa --current
      k4s context add dev --host 192.168.1.10 --username root --password secret

    \b
    Examples (K8s):
      k4s context add staging --type k8s --kubeconfig ~/.kube/k4s/default.yaml --current
      k4s context add customer-x --type k8s --kubectl-context customer-x-aks --namespace starburst
    """
    if not name:
        click.echo("Usage: k4s context add NAME [OPTIONS]\n")
        click.echo("Error: Missing argument 'NAME'.")
        click.echo("Hint:  k4s context add <name> --host <ip> --username <user> -i <key>")
        click.echo("       k4s context add <name> --type k8s --kubeconfig <path>")
        ctx.exit(2)

    state: CliState = ctx.obj["state"]

    if ctx_type == "k8s":
        # K8s context validation
        if not kubeconfig and not kubectl_context:
            state.ui.error("K8s context requires --kubeconfig and/or --kubectl-context.")
            ctx.exit(2)
        state.contexts.upsert(
            Context(
                name=name,
                type="k8s",
                kubeconfig=kubeconfig,
                kubectl_context=kubectl_context,
                namespace=namespace,
            ),
            make_current=current,
        )
    else:
        # VM context validation
        if not host:
            state.ui.error("VM context requires --host.")
            ctx.exit(2)
        if not username:
            state.ui.error("VM context requires --username.")
            ctx.exit(2)
        if not identity_file and not password:
            state.ui.error("Provide either --identity-file or --password.")
            ctx.exit(2)

        if ask_identity_passphrase:
            if not identity_file:
                state.ui.error("--ask-identity-passphrase requires --identity-file.")
                ctx.exit(2)
            identity_passphrase = click.prompt(
                "SSH key passphrase",
                hide_input=True,
                default="",
                show_default=False,
            )
        state.contexts.upsert(
            Context(
                name=name,
                type="vm",
                host=host,
                username=username,
                port=port,
                identity_file=identity_file,
                identity_passphrase=identity_passphrase or None,
                password=password,
            ),
            make_current=current,
        )

    msg = f"Context '{name}' ({ctx_type}) saved."
    if current:
        msg = f"Context '{name}' ({ctx_type}) saved and set as current."
    state.ui.success(msg)


@context.command("list", short_help="List all contexts.")
@click.pass_context
def list_(ctx):
    """List all saved contexts. The current context is marked with *."""
    state: CliState = ctx.obj["state"]
    current = state.contexts.get_current_name()
    contexts = state.contexts.list()
    if not contexts:
        state.ui.info("No contexts found. Run `k4s context add ...` to create one.")
        return
    for name, c in contexts.items():
        marker = "*" if name == current else " "
        tag = f"[{c.type}]"
        if c.type == "k8s":
            detail = c.kubeconfig or ""
            if c.kubectl_context:
                detail = f"context: {c.kubectl_context}" if not detail else f"{detail} (context: {c.kubectl_context})"
            if c.namespace:
                detail += f" ns: {c.namespace}"
        else:
            host = c.host or "-"
            user = c.username or "-"
            detail = f"{user}@{host}:{c.port}"
        state.ui.info(f"{marker} {name:20} {tag:6} {detail}")


@context.command("use-context", short_help="Set the current context.")
@click.argument("name")
@click.pass_context
def use_context(ctx, name):
    """Set NAME as the current context.

    Commands that accept --context will use the current context when
    --context is omitted.
    """
    state: CliState = ctx.obj["state"]
    try:
        state.contexts.set_current(name)
    except KeyError as e:
        state.ui.error(str(e))
        ctx.exit(2)
    state.ui.success(f"Current context set to '{name}'.")


@context.command("show", short_help="Show context details.")
@click.argument("name", required=False, default=None)
@click.pass_context
def show(ctx, name):
    """Show context details.

    If NAME is given, shows that context. Otherwise shows the current context.
    """
    state: CliState = ctx.obj["state"]
    try:
        c = state.contexts.resolve(name)
    except KeyError as e:
        state.ui.error(str(e))
        ctx.exit(2)
    _print_context(state.ui, c)


def _print_context(ui, c: Context) -> None:
    """Print context details to the UI."""
    ui.info(f"name:            {c.name}")
    ui.info(f"type:            {c.type}")
    if c.type == "k8s":
        ui.info(f"kubeconfig:      {c.kubeconfig or '-'}")
        ui.info(f"kubectl_context: {c.kubectl_context or '-'}")
        ui.info(f"namespace:       {c.namespace or '-'}")
    else:
        ui.info(f"host:            {c.host}")
        ui.info(f"username:        {c.username}")
        ui.info(f"port:            {c.port}")
        ui.info(f"identity_file:   {c.identity_file or '-'}")
        ui.info(f"passphrase:      {'<set>' if c.identity_passphrase else '-'}")
        ui.info(f"password:        {'<set>' if c.password else '-'}")


@context.command("edit", short_help="Edit a context in $EDITOR.")
@click.argument("name")
@click.pass_context
def edit(ctx, name):
    """Edit an existing context using an editor.

    This edits non-sensitive fields only. Secrets (password/passphrase) are
    preserved from the existing context.
    """
    import os
    import tempfile
    import yaml

    state: CliState = ctx.obj["state"]
    ui = state.ui

    try:
        c = state.contexts.resolve(name)
    except KeyError as e:
        ui.error(str(e))
        ctx.exit(2)

    # Only expose non-sensitive fields for editing.
    data = {
        "name": c.name,
        "type": c.type,
    }
    if c.type == "k8s":
        data.update(
            {
                "kubeconfig": c.kubeconfig,
                "kubectl_context": c.kubectl_context,
                "namespace": c.namespace,
            }
        )
    else:
        data.update(
            {
                "host": c.host,
                "username": c.username,
                "port": c.port,
                "identity_file": c.identity_file,
            }
        )

    editor = (
        os.environ.get("K4S_EDITOR")
        or os.environ.get("VISUAL")
        or os.environ.get("EDITOR")
        or "vi"
    )

    with tempfile.NamedTemporaryFile("w", delete=False, suffix=".yml") as f:
        path = f.name
        f.write("# Edit context fields and save. Close the editor to apply.\n")
        f.write("# Secrets are not editable here and will remain unchanged.\n\n")
        yaml.safe_dump(data, f, sort_keys=False)

    try:
        rc = subprocess.call([editor, path])
        if rc != 0:
            raise RuntimeError(f"Editor exited with code {rc}")
        with open(path, "r", encoding="utf-8") as rf:
            new_data = yaml.safe_load(rf) or {}
    finally:
        try:
            os.unlink(path)
        except OSError:
            pass

    if not isinstance(new_data, dict):
        ui.error("Invalid YAML. Expected a mapping at the root.")
        ctx.exit(2)

    # Name/type are not editable.
    new_type = (new_data.get("type") or c.type)
    if new_type != c.type:
        ui.error("Context type cannot be changed via edit. Create a new context instead.")
        ctx.exit(2)

    if c.type == "k8s":
        updated = Context(
            name=c.name,
            type="k8s",
            kubeconfig=new_data.get("kubeconfig") or None,
            kubectl_context=new_data.get("kubectl_context") or None,
            namespace=new_data.get("namespace") or None,
        )
    else:
        updated = Context(
            name=c.name,
            type="vm",
            host=new_data.get("host") or None,
            username=new_data.get("username") or None,
            port=int(new_data.get("port") or 22),
            identity_file=new_data.get("identity_file") or None,
            identity_passphrase=c.identity_passphrase,
            password=c.password,
        )

    state.contexts.upsert(updated, make_current=(state.contexts.get_current_name() == c.name))
    ui.success(f"Context '{c.name}' updated.")

@context.command("remove", short_help="Remove a context.")
@click.argument("name")
@click.pass_context
def remove(ctx, name):
    """Remove a saved context. Alias: rm"""
    state: CliState = ctx.obj["state"]
    try:
        state.contexts.remove(name)
    except KeyError as e:
        state.ui.error(str(e))
        ctx.exit(2)
    state.ui.success(f"Context '{name}' removed.")


@context.command("ping", short_help="Check connectivity.")
@click.argument("name", required=False)
@click.option("-q", "--quiet", is_flag=True, help="Only print final results and errors.")
@click.option("-v", "--verbose", count=True, help="Increase verbosity (-v logs, -vv debug).")
@click.pass_context
def ping(ctx, name, quiet, verbose):
    """Check connectivity for a context (defaults to current).

    For VM contexts, tests SSH connectivity.
    For K8s contexts, runs 'kubectl get nodes' to verify cluster access.
    """
    state: CliState = ctx.obj["state"]
    apply_command_ui_overrides(ctx, state, quiet=quiet, verbose=verbose)
    ui = state.ui

    try:
        c = state.contexts.resolve(name)

        if c.type == "k8s":
            with ui.step("Checking cluster connectivity"):
                kube_args = c.to_kubectl_args()
                cmd = ["kubectl"] + kube_args + ["get", "nodes"]
                ui.log(f"Running: {' '.join(shlex.quote(a) for a in cmd)}")
                result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
                if result.returncode != 0:
                    err_msg = result.stderr.strip() or result.stdout.strip()
                    raise RuntimeError(f"kubectl get nodes failed: {err_msg}")
            ui.success("Cluster connectivity OK.")
        else:
            with ui.step(f"Connecting to {c.username}@{c.host}:{c.port}"):
                with Executor(c.to_server_config()) as ex:
                    rc, out, err = ex.execute("whoami && hostname")
                    if rc != 0:
                        raise RuntimeError(f"SSH check failed: {err}")
            ui.success("SSH connectivity OK.")
    except Exception as e:
        exit_with_error(ctx, ui, e, code=1)
