"""Recipe for building Dataiku DSS container execution base images.

Runs ``dssadmin build-base-image --type container-exec`` on a target VM
where Dataiku DSS is already installed, and optionally pushes the
resulting image to a Docker registry.
"""
from __future__ import annotations

from dataclasses import dataclass

from k4s.core.executor import Executor, ExecutorError
from k4s.core.products import Step
from k4s.recipes.common.run import check, q, run
from k4s.ui.ui import Ui


@dataclass(frozen=True)
class BaseImagePlan:
    """Plan inputs for building a Dataiku base image."""

    data_dir: str
    os_user: str
    tag: str = "dss-base-image"
    distrib: str = "almalinux8"
    with_cuda: bool = False
    cuda_version: str | None = None
    with_py311: bool = False
    without_r: bool = False
    system_packages: str | None = None
    http_proxy: str | None = None
    no_proxy: str | None = None
    push: bool = False
    registry: str | None = None
    registry_username: str | None = None
    registry_password: str | None = None


def build_dssadmin_args(plan: BaseImagePlan) -> str:
    """Build the argument string for dssadmin build-base-image."""
    args = "--type container-exec"
    args += f" --distrib {q(plan.distrib)}"
    args += f" --tag {q(plan.tag)}"

    if plan.with_cuda:
        args += " --with-cuda"
        if plan.cuda_version:
            args += f" --cuda-version {q(plan.cuda_version)}"

    if plan.with_py311:
        args += " --with-py311"

    if plan.without_r:
        args += " --without-r"

    if plan.system_packages:
        args += f" --system-packages {q(plan.system_packages)}"

    if plan.http_proxy:
        args += f" --http-proxy {q(plan.http_proxy)}"

    if plan.no_proxy:
        args += f" --no-proxy {q(plan.no_proxy)}"

    return args


def build_base_image_steps(
    ui: Ui,
    ex: Executor,
    plan: BaseImagePlan,
) -> list[Step]:
    """Return steps for building (and optionally pushing) a Dataiku base image."""

    def _preflight():
        ui.log(f"Checking dssadmin exists at {plan.data_dir}.")
        rc, _, _ = run(ex, f"test -x {q(plan.data_dir)}/bin/dssadmin")
        if rc != 0:
            raise ExecutorError(
                f"dssadmin not found at {plan.data_dir}/bin/dssadmin.\n"
                "Install DSS first with `k4s install dataiku`, or pass the correct --data-dir."
            )
        ui.log("Checking Docker is available.")
        rc2, _, _ = run(ex, f"sudo -n -u {q(plan.os_user)} docker info", silent=True)
        if rc2 != 0:
            raise ExecutorError(
                f"Docker is not available for user '{plan.os_user}'.\n"
                "Install Docker first with `k4s install docker` and ensure the user is in the docker group."
            )

    def _build():
        args = build_dssadmin_args(plan)
        ui.log(f"Building base image: {plan.tag} (distrib={plan.distrib})")
        cmd = (
            f"sudo -n -u {q(plan.os_user)} -i -- "
            f"{q(plan.data_dir)}/bin/dssadmin build-base-image {args}"
        )
        check(ex, cmd)

    def _docker_login():
        if not plan.registry:
            return
        ui.log(f"Logging in to registry: {plan.registry}")
        login_cmd = f"docker login {q(plan.registry)}"
        if plan.registry_username:
            login_cmd += f" --username {q(plan.registry_username)}"
        if plan.registry_password:
            login_cmd += " --password-stdin"
            check(ex, login_cmd, stdin_data=plan.registry_password + "\n", silent=True)
        else:
            check(ex, login_cmd)

    def _push():
        if not plan.registry:
            raise ExecutorError(
                "Cannot push without a registry. Use --registry to specify the target registry."
            )
        remote_tag = f"{plan.registry.rstrip('/')}/{plan.tag}"
        ui.log(f"Tagging image as {remote_tag}")
        check(ex, f"docker tag {q(plan.tag)} {q(remote_tag)}")
        ui.log(f"Pushing image to {remote_tag}")
        check(ex, f"docker push {q(remote_tag)}")

    steps: list[Step] = [
        Step(title=f"Preflight (base image) on {ex.host}", run=_preflight),
        Step(title=f"Build base image ({plan.tag})", run=_build),
    ]

    if plan.push:
        if plan.registry:
            steps.append(Step(title=f"Docker login to {plan.registry}", run=_docker_login))
        steps.append(Step(title="Push image to registry", run=_push))

    return steps
