prole/installer/ui/screens/docker.py

"""Docker image build, registry management and image pre-pull."""

import os
import platform
import re
import shutil
import subprocess
import sys
import time
from pathlib import Path

from installer.config import get_docker_build_platform_args
from installer.core.env import (
    PROJECT_ROOT,
    _collect_images_from_files,
    _deployment_mode_from_env,
    _find_kubeconfig_file,
    _host_from_url,
    _http_ping_registry,
    _k3d_prole_data_volume_args,
    _local_registry_enabled,
    _normalize_cluster_env,
    _push_docker_image,
    _resolve_supabase_home,
)


class DockerScreenMixin:
    """Docker image build, registry management and image pre-pull."""

    def _k3s_registry_hostport(self) -> str:
        cfg = self.prole_cfg_data or {}

        def _get_cfg_val(key: str) -> str:
            for section in (
                "Global",
                "Service Cluster (k3s)",
                "Initialize Cluster",
                "Docker Build",
            ):
                try:
                    val = ((cfg.get(section, {}) or {}).get(key) or "").strip()
                except Exception:
                    val = ""
                if val:
                    return val
            return ""

        def _sanitize_host(h: str) -> str:
            host = (h or "").strip()
            if not host:
                return ""
            # `0.0.0.0` / `::` are bind addresses and are never a usable client endpoint.
            # `localhost`/`127.0.0.1` must not be used for k3s registry access.
            if host in ("0.0.0.0", "::", "127.0.0.1", "localhost"):
                return ""
            return host

        # Prefer explicit registry host/port (align with `etc/init_k3s_registry.sh`).
        raw_host = (
            _get_cfg_val("K3S_REGISTRY_HOST")
            or os.environ.get("K3S_REGISTRY_HOST")
            or os.environ.get("PROLE_K3S_REGISTRY_HOST")
            or ""
        ).strip()
        raw_port = (
            _get_cfg_val("K3S_REGISTRY_PORT")
            or os.environ.get("K3S_REGISTRY_PORT")
            or os.environ.get("PROLE_K3S_REGISTRY_PORT")
            or ""
        ).strip()
        if raw_host:
            host = _sanitize_host(_host_from_url(raw_host) or raw_host)
            port = raw_port or "5000"
            if host:
                return f"{host}:{port}"

        # Otherwise derive the registry host from the configured k3s API server URL.
        server = (
            (cfg.get("Global", {}) or {}).get("PROLE_K3S_SERVER")
            or (cfg.get("Global", {}) or {}).get("K3S_SERVER_URL")
            or (cfg.get("Initialize Cluster", {}) or {}).get("K3S_SERVER_URL")
            or (cfg.get("Service Cluster (k3s)", {}) or {}).get("K3S_SERVER_URL")
            or os.environ.get("PROLE_K3S_SERVER")
            or os.environ.get("K3S_SERVER_URL")
            or ""
        )
        host = _sanitize_host(_host_from_url(server))
        if host:
            return f"{host}:5000"

        # Last resort: attempt to extract a server endpoint from an existing kubeconfig.
        try:
            kubeconfig = (_find_kubeconfig_file() or "").strip()
            if kubeconfig and Path(kubeconfig).expanduser().is_file():
                text = Path(kubeconfig).expanduser().read_text(encoding="utf-8")
                for line in text.splitlines():
                    m = re.match(r"^\s*server:\s*(\S+)\s*$", line)
                    if not m:
                        continue
                    kc_host = _sanitize_host(_host_from_url(m.group(1)))
                    if kc_host:
                        return f"{kc_host}:5000"
        except Exception:
            return ""

        return ""

    def _registry_namespace(self) -> str:
        cfg = self.prole_cfg_data or {}
        ns = (
            (cfg.get("Global", {}) or {}).get("REGISTRY_NAMESPACE")
            or os.environ.get("REGISTRY_NAMESPACE")
            or ""
        )
        ns = (ns or "").strip()
        if ns:
            return ns

        # Registry is a common-core service; default to service namespace.
        try:
            ns = (self._get_service_namespace() or "").strip()
        except Exception:
            ns = (
                (cfg.get("Global", {}) or {}).get("SERVICE_NAMESPACE")
                or os.environ.get("SERVICE_NAMESPACE")
                or ""
            ).strip()

        return ns or "default"

    def check_docker_running(self):
        """Check if Docker is running"""
        try:
            result = subprocess.run(["docker", "ps"], capture_output=True, timeout=10)
            return result.returncode == 0
        except Exception:
            return False

    def ensure_registry_available(self, env: str) -> str:
        """Detect if k8s.prole.org:5000 is reachable; if so use it. Otherwise ensure a local registry is available.
        Returns the registry URL (host:port) to be used for tagging/pushing.
        For Dev, may create a local k3d-managed registry; for other envs, still prefer the external if reachable.
        """
        mode = _deployment_mode_from_env(
            (self.prole_cfg_data.get("Global", {}) or {}).get("DEPLOYMENT_MODE")
            or env
        )
        if mode == "k3s":
            # k3s must use the k3s registry (never localhost/k3d registries).
            return self._k3s_registry_hostport()

        # Prefer the shared registry if reachable
        if _http_ping_registry("k8s.prole.org", 5000):
            return "k8s.prole.org:5000"

        if not _local_registry_enabled(env):
            return ""

        # Otherwise, ensure a local registry (localhost:5000) exists/started
        if _http_ping_registry("localhost", 5000):
            return "localhost:5000"

        # Use k3d registry helper for Dev; for non-Dev, we still create/use local for pushing
        reg_name = "prole-registry"
        # Check if k3d is installed
        k3d_exists = (
            subprocess.run(["which", "k3d"], capture_output=True).returncode == 0
        )
        if k3d_exists:
            # List registries
            lst = subprocess.run(
                ["k3d", "registry", "list"], capture_output=True, text=True
            )
            if reg_name not in (lst.stdout or ""):
                # Create registry exposed on 0.0.0.0:5000
                subprocess.run(
                    ["k3d", "registry", "create", reg_name, "--port", "0.0.0.0:5000"],
                    check=True,
                )
        return "localhost:5000"

    def ensure_local_registry_available(self):
        """Ensure a local registry is running and return (host_registry, cluster_registry)."""
        mode_hint = None
        try:
            mode_hint = self._deployment_mode()
        except Exception:
            mode_hint = None

        if not _local_registry_enabled(mode_hint):
            try:
                log_path = self._registry_log_path()
                with log_path.open("a", encoding="utf-8") as fp:
                    fp.write(
                        f"\n# Registry disabled (mode={mode_hint or 'unknown'}) @ {time.strftime('%Y-%m-%d %H:%M:%S')}\n"
                    )
            except Exception:
                pass
            return None

        log_path = self._registry_log_path()
        try:
            log_fp = log_path.open("a", encoding="utf-8")
        except Exception:
            log_fp = None

        def _log(msg: str):
            if not log_fp:
                return
            try:
                log_fp.write(msg)
                log_fp.flush()
            except Exception:
                pass

        def _log_cmd(cmd: list[str], res: subprocess.CompletedProcess):
            _log(f"$ {' '.join(cmd)}\n")
            if res.stdout:
                _log(res.stdout)
            if res.stderr:
                _log(res.stderr)
            _log(f"[exit {res.returncode}]\n")

        def _run(cmd: list[str], check: bool = False):
            res = subprocess.run(cmd, capture_output=True, text=True)
            if res.returncode != 0:
                _log_cmd(cmd, res)
                if check:
                    raise RuntimeError(f"Command failed: {' '.join(cmd)}")
            return res

        _log(
            f"\n# Registry check @ {time.strftime('%Y-%m-%d %H:%M:%S')} (mode={mode_hint or 'unknown'})\n"
        )
        reg_name = "prole-registry"
        registry_container = f"k3d-{reg_name}"
        host_registry = "localhost:5000"

        # Prefer the currently selected environment (UI) over the global config.
        # This avoids incorrectly treating Service (k3s) as Dev (k3d) when
        # Global/DEPLOYMENT_MODE is still set to k3d.
        mode_source = ""
        try:
            mode_source = (self._cluster_env_key() or "").strip()
        except Exception:
            mode_source = ""
        if not mode_source:
            try:
                mode_source = (
                    (self.prole_cfg_data.get("Global", {}) or {})
                    .get("DEPLOYMENT_MODE", "")
                    .strip()
                )
            except Exception:
                mode_source = ""
        if not mode_source:
            mode_source = (mode_hint or "").strip()

        mode = _deployment_mode_from_env(mode_source)
        if mode == "k3s":
            k3s_registry = self._k3s_registry_hostport()
            if not k3s_registry:
                _log("k3s mode but K3S_SERVER_URL/PROLE_K3S_SERVER not set; cannot resolve registry.\n")
                if log_fp:
                    log_fp.close()
                return None
            host_registry = k3s_registry
            registry_ns = self._registry_namespace()
            cluster_registry = f"registry.{registry_ns}.svc.cluster.local:5000"
            if not self.check_docker_running():
                _log("Docker not running; cannot push/prepull images to k3s registry.\n")
                if log_fp:
                    log_fp.close()
                return None
            self.local_registry_url = host_registry
            self.local_registry_internal = cluster_registry
            try:
                if "Docker Build" not in self.prole_cfg_data:
                    self.prole_cfg_data["Docker Build"] = {}
                self.prole_cfg_data["Docker Build"]["LOCAL_REGISTRY"] = host_registry
                self.prole_cfg_data["Docker Build"][
                    "LOCAL_REGISTRY_INTERNAL"
                ] = cluster_registry
                self.safe_after(self._save_prole_cfg)
            except Exception:
                pass
            if log_fp:
                log_fp.close()
            return host_registry, cluster_registry

        if mode == "k3d":
            cluster_registry = f"{registry_container}.localhost:5000"
        else:
            cluster_registry = host_registry

        if not self.check_docker_running():
            _log("Docker not running; cannot ensure local registry.\n")
            if log_fp:
                log_fp.close()
            return None

        # 1. If a registry is already responsive on localhost:5000, we use it.
        if _http_ping_registry("localhost", 5000):
            # Check if it's a k3d registry to use the correct cluster internal name
            try:
                lst = _run(["k3d", "registry", "list"])
                if registry_container in (lst.stdout or "") or reg_name in (
                    lst.stdout or ""
                ):
                    cluster_registry = f"{registry_container}.localhost:5000"
                else:
                    cluster_registry = host_registry
            except Exception:
                cluster_registry = host_registry

            self.local_registry_url = host_registry
            self.local_registry_internal = cluster_registry
            try:
                if "Docker Build" not in self.prole_cfg_data:
                    self.prole_cfg_data["Docker Build"] = {}
                self.prole_cfg_data["Docker Build"]["LOCAL_REGISTRY"] = host_registry
                self.prole_cfg_data["Docker Build"][
                    "LOCAL_REGISTRY_INTERNAL"
                ] = cluster_registry
                self.safe_after(self._save_prole_cfg)
            except Exception:
                pass
            if log_fp:
                log_fp.close()
            return host_registry, cluster_registry

        # 2. Try k3d first if available
        k3d_exists = _run(["which", "k3d"]).returncode == 0
        if k3d_exists:
            try:
                # Ensure registry exists
                lst = _run(["k3d", "registry", "list"])
                if registry_container not in (lst.stdout or "") and reg_name not in (
                    lst.stdout or ""
                ):
                    _run(
                        [
                            "k3d",
                            "registry",
                            "create",
                            reg_name,
                            "--port",
                            "0.0.0.0:5000",
                        ],
                        check=True,
                    )

                # Ensure registry container is running
                running = _run(
                    [
                        "docker",
                        "ps",
                        "--filter",
                        f"name={registry_container}",
                        "--format",
                        "{{.Names}}",
                    ],
                )
                if not (running.stdout or "").strip():
                    _run(["docker", "start", registry_container], check=True)

                self.local_registry_url = host_registry
                self.local_registry_internal = f"{registry_container}.localhost:5000"
                if log_fp:
                    log_fp.close()
                return self.local_registry_url, self.local_registry_internal
            except Exception:
                _log("k3d registry setup failed.\n")

        # 3. Fallback to plain Docker registry
        try:
            docker_exists = _run(["which", "docker"]).returncode == 0
            if docker_exists:
                # Check if prole-registry container exists
                exists = _run(
                    [
                        "docker",
                        "ps",
                        "-a",
                        "--filter",
                        "name=^/prole-registry$",
                        "--format",
                        "{{.Names}}",
                    ],
                )
                if (exists.stdout or "").strip():
                    _run(["docker", "start", "prole-registry"], check=True)
                else:
                    _run(
                        [
                            "docker",
                            "run",
                            "-d",
                            "-p",
                            "5000:5000",
                            "--restart",
                            "always",
                            "--name",
                            "prole-registry",
                            "registry:2",
                        ],
                        check=True,
                    )

                # Wait a moment for it to start
                time.sleep(1)
                if _http_ping_registry("localhost", 5000):
                    self.local_registry_url = host_registry
                    self.local_registry_internal = host_registry
                    if log_fp:
                        log_fp.close()
                    return host_registry, host_registry
        except Exception:
            _log("Docker registry fallback failed.\n")

        if log_fp:
            log_fp.close()
        return None

    def build_docker_image(self):
        """Build prole-db Docker image"""
        # If Kerberos is enabled, update pg_hba.conf in conf/postgresql before copying
        if self.kerberos_enabled.get():
            realm = self.kerberos_realm.get().strip() or "EXAMPLE.COM"
            hba_src = PROJECT_ROOT / "conf" / "postgresql" / "pg_hba.conf"
            if hba_src.exists():
                content = hba_src.read_text()
                # Also ensure the Kerberos line exists in the source if it's not there yet
                if "gss" not in content:
                    content += f"\nhost    all             all             all                     gss include_realm=1 krb_realm={realm}\n"
                else:
                    content = content.replace(
                        "krb_realm=EXAMPLE.COM", f"krb_realm={realm}"
                    )
                hba_src.write_text(content)
                print(f"Updated {hba_src} with realm {realm}")

        # Base local image tag (before pushing to registry)
        version = self.get_prole_db_version()
        image_tag = f"prole-db:{version}"

        # Prepare build context: copy conf/postgresql to prole-db/postgresql
        conf_src = PROJECT_ROOT / "conf" / "postgresql"
        conf_dst = PROJECT_ROOT / "prole-db" / "postgresql"

        if conf_dst.exists():
            shutil.rmtree(conf_dst, ignore_errors=True)
        if sys.version_info >= (3, 8):
            shutil.copytree(conf_src, conf_dst, dirs_exist_ok=True)
        else:
            shutil.copytree(conf_src, conf_dst)

        build_cmd = ["docker", "build", "--progress=plain", "-t", image_tag]

        target_env = self._cluster_env_key()
        if getattr(self, "deploy_environment", None):
            try:
                target_env = _normalize_cluster_env(self.deploy_environment.get())
            except Exception:
                pass
        # Add platform flag for target environment (k3s on Pi uses arm64)
        build_cmd.extend(get_docker_build_platform_args(target_env))

        build_cmd.append(".")

        result = subprocess.run(
            build_cmd,
            check=True,
            cwd=PROJECT_ROOT / "prole-db",
            capture_output=True,
            text=True,
        )
        if result.returncode != 0:
            raise Exception(f"Failed to build image: {result.stderr}")
        # Keep a reference for later steps
        self.local_image_tag = image_tag

        # Update K8s manifest
        self.update_k8s_manifest(image_tag)

    def build_mssql_docker_image(
        self, image_tag="prole-mssql-db:latest", target_env: str | None = None
    ):
        """Build mssql Docker image (with platform detection for target env)."""
        build_cmd = ["docker", "build", "-t", image_tag]

        env_key = target_env or self._cluster_env_key()
        # Add platform flag for Apple Silicon / k3s target
        build_cmd.extend(get_docker_build_platform_args(env_key))

        build_cmd.append(".")

        result = subprocess.run(
            build_cmd,
            check=True,
            cwd=PROJECT_ROOT / "mssql",
            capture_output=True,
            text=True,
        )
        if result.returncode != 0:
            raise Exception(f"Failed to build mssql image: {result.stderr}")
        return result

    def tag_docker_image(self):
        """Tag Docker image for registry"""
        if not getattr(self, "registry_url", ""):
            print("Registry disabled; skipping tag.")
            return
        version = self.get_prole_db_version()
        registry = getattr(self, "registry_url", "localhost:5000")
        image_name = f"prole-db:{version}"
        image = getattr(self, "local_image_tag", image_name)
        self.remote_image_tag = f"{registry}/prole-db:{version}"
        subprocess.run(
            ["docker", "tag", image, self.remote_image_tag],
            check=True,
            capture_output=True,
        )

    def push_docker_image(self):
        """Push Docker image to registry"""
        if not getattr(self, "registry_url", ""):
            print("Registry disabled; skipping push.")
            return
        remote_tag = getattr(self, "remote_image_tag", None)
        if not remote_tag:
            raise Exception("Remote image tag not set")
        if not _push_docker_image(remote_tag):
            raise Exception("Failed to push image to registry.")

    def import_k3d_image(self, cluster_name="knoe-dev-cluster"):
        """Import image to k3d cluster (only for Dev)."""
        version = self.get_prole_db_version()
        image_name = f"prole-db:{version}"
        subprocess.run(
            ["k3d", "image", "import", image_name, "-c", cluster_name],
            check=True,
            capture_output=True,
        )

    def update_k8s_manifest(self, image_tag):
        """Update k8s/prole/prole-db.yaml with the new version and Kerberos config"""
        manifest_paths = [
            PROJECT_ROOT / "k8s" / "prole" / "prole-db.yaml",
            PROJECT_ROOT / "k8s" / "prole" / "prole-db-recovery.yaml.tpl",
        ]
        import re

        for manifest_path in manifest_paths:
            if not manifest_path.exists():
                continue
            content = manifest_path.read_text()
            # Update imageName: prole-db:17.7-043
            new_content = re.sub(r"imageName:\s*.*", f"imageName: {image_tag}", content)

            # Update Kerberos realm in manifest if enabled
            if self.kerberos_enabled.get():
                realm = self.kerberos_realm.get().strip() or "EXAMPLE.COM"
                if "gss" not in new_content:
                    # Insert GSS rule after existing scram rules for service accounts
                    new_content = new_content.replace(
                        "      - host all all all scram-sha-256",
                        f"      - host all authenticator all scram-sha-256\n      - host all all all gss include_realm=1 krb_realm={realm}\n      - host all all all scram-sha-256",
                        1,  # replace only first occurrence
                    )
                else:
                    new_content = new_content.replace(
                        "krb_realm=EXAMPLE.COM", f"krb_realm={realm}"
                    )

            if new_content != content:
                manifest_path.write_text(new_content)

    def get_prole_db_version(self):
        return self.controller.get_prole_db_version()

    def create_or_select_cluster(self, env: str):
        """Ensure target cluster depending on environment selection."""
        if env == "Dev":
            self.create_or_recreate_k3d_dev_cluster()
        elif env == "Service":
            kubectl = subprocess.run(["which", "kubectl"], capture_output=True)
            if kubectl.returncode != 0:
                raise Exception(
                    "kubectl not found. Please install kubectl and configure access to the target cluster."
                )
            server, token = self._k3s_connection_info()
            if server and token:
                res = subprocess.run(
                    self._k3s_kubectl_base_cmd() + ["cluster-info"],
                    capture_output=True,
                    text=True,
                )
                if res.returncode != 0:
                    raise Exception(
                        "K3s cluster not reachable with provided server/token."
                    )
            else:
                # Fallback to local kubeconfig if no token provided
                subprocess.run(["kubectl", "version", "--client"], check=True)
        elif env == "Prod":
            kubectl = subprocess.run(["which", "kubectl"], capture_output=True)
            if kubectl.returncode != 0:
                raise Exception(
                    "kubectl not found. Please install kubectl and configure access to the target cluster."
                )
            subprocess.run(["kubectl", "version", "--client"], check=True)
        else:
            raise Exception(f"Unknown environment: {env}")

    def create_or_recreate_k3d_dev_cluster(self, cluster_name="knoe-dev-cluster"):
        """Create or restart local k3d cluster and wire it to the chosen registry."""
        result = subprocess.run(
            ["k3d", "cluster", "list"], capture_output=True, text=True
        )
        if cluster_name in (result.stdout or ""):
            subprocess.run(["k3d", "cluster", "delete", cluster_name], check=True)

        # Determine registry integration args
        reg_args = []
        if getattr(self, "registry_url", None):
            # If using the local k3d registry, we want to create or use it
            if self.registry_url.startswith("localhost:5000"):
                # Creating with --registry-create ensures it's available and integrated
                reg_args = ["--registry-create", f"prole-registry:0.0.0.0:5000"]
            else:
                reg_args = ["--registry-use", self.registry_url]

        prole_data = str(self._resolve_env_dir("PROLE_DATA", "data"))
        volume_args = _k3d_prole_data_volume_args(prole_data)
        cmd = (
            ["k3d", "cluster", "create", cluster_name, "-a", "2", "--wait"]
            + volume_args
            + reg_args
            + ["--timestamps"]
        )
        subprocess.run(cmd, check=True, cwd=PROJECT_ROOT)