"""Starburst Enterprise upgrade recipe (helm upgrade)."""

from __future__ import annotations

import shlex

from k4s.core.executor import Executor, ExecutorError
from k4s.core.products import Step
from k4s.recipes.common.helm import (
    force_delete_pods,
    postcheck_pods,
    release_exists,
    snapshot_release_pods,
    wait_for_rollout,
)
from k4s.recipes.common.run import check
from k4s.recipes.starburst.install import build_common_steps, build_helm_value_args
from k4s.recipes.starburst.model import StarburstInstallPlan
from k4s.ui.ui import Ui


def build_upgrade_steps(
    ui: Ui, ex: Executor, plan: StarburstInstallPlan, *, force: bool = False,
) -> list[Step]:
    """Build Starburst upgrade steps using ``helm upgrade``.

    If the release does not exist, raises an error and hints ``k4s install``.

    ``helm upgrade`` always runs without ``--wait``.  Readiness is verified
    via ``kubectl rollout status`` which gives better diagnostics on failure.

    When *force* is ``True`` the pre-upgrade pods are additionally
    force-deleted (``--grace-period=0 --force``) before the rollout wait.
    """
    steps = build_common_steps(ui, ex, plan, is_upgrade=True)

    # Mutable list shared between the _upgrade and _force_delete closures.
    _pre_upgrade_pods: list[str] = []

    def _upgrade():
        if not release_exists(ex, plan):
            raise ExecutorError(
                f"Helm release '{plan.release_name}' not found in namespace '{plan.namespace}'.\n"
                "Install first with `k4s install starburst`."
            )

        # Capture current pods before helm changes anything.
        if force:
            _pre_upgrade_pods.extend(snapshot_release_pods(ui, ex, plan))

        ui.log("Running helm upgrade with OCI chart reference")
        cmd_parts: list[str] = [
            "helm", "upgrade",
            plan.release_name, plan.oci_chart_ref,
            "--version", plan.chart_version,
            "--namespace", plan.namespace,
        ] + plan.helm_flags() + build_helm_value_args(plan)

        cmd = " ".join(shlex.quote(p) for p in cmd_parts)
        check(ex, cmd)

    steps.append(Step(title=f"Upgrade Starburst release '{plan.release_name}'", run=_upgrade))
    if force:
        steps.append(Step(
            title="Delete pods now (best-effort)",
            run=lambda: force_delete_pods(ui, ex, plan, _pre_upgrade_pods),
        ))
    steps.append(Step(title="Wait for rollout to complete", run=lambda: wait_for_rollout(ui, ex, plan)))
    steps.append(Step(title="Post-check (pods)", run=lambda: postcheck_pods(ui, ex, plan)))
    return steps
