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.core.env import (
    PROJECT_ROOT,
    _collect_images_from_files,
    _deployment_mode_from_env,
    _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 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.
        """
        # 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"

        mode = _deployment_mode_from_env(
            self.prole_cfg_data.get("Global", {}).get("DEPLOYMENT_MODE")
        )
        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 _collect_dependent_images(
        self, include_supabase: bool, include_kerberos_proxy: bool
    ) -> list[str]:
        images = set()
        for rel_dir in ("k8s/prole", "k8s/openbao"):
            base_dir = PROJECT_ROOT / rel_dir
            if not base_dir.exists():
                continue
            images.update(_collect_images_from_files(list(base_dir.glob("*.yaml"))))

        if include_supabase:
            supa_home = _resolve_supabase_home(PROJECT_ROOT)
            if supa_home:
                docker_dir = supa_home / "docker"
                compose_files = [docker_dir / "docker-compose.yml"]
                if os.environ.get("SUPABASE_USE_DEV_COMPOSE") == "1":
                    compose_files.append(docker_dir / "dev" / "docker-compose.dev.yml")
                images.update(_collect_images_from_files(compose_files))

        if not include_supabase:
            images = {img for img in images if "supabase" not in img}

        if include_kerberos_proxy:
            krb_img = os.environ.get("KRB5_AD_PROXY_IMAGE", "alpine/socat")
            if krb_img:
                images.add(krb_img)

        return sorted(images)

    def _prepull_images_to_registry(
        self, include_supabase: bool, include_kerberos_proxy: bool, log=None
    ) -> bool:
        def _log(msg: str):
            if log:
                try:
                    log(msg)
                except Exception:
                    pass

        info = self.ensure_local_registry_available()
        if not info:
            _log("Local registry unavailable; skipping image pre-pull.\n")
            return False

        registry, _cluster_registry = info
        images = self._collect_dependent_images(
            include_supabase, include_kerberos_proxy
        )
        if not images:
            _log("No dependent images found to pre-pull.\n")
            return True

        overall_ok = True
        import_dir = self.docker_import_dir.get().strip()

        for image in images:
            local_tag = image
            if not image.startswith(f"{registry}/"):
                local_tag = f"{registry}/{image}"

            # 1. Check if already in local registry
            _log(f"Checking if {image} exists in local registry...\n")
            check_reg = subprocess.run(
                ["docker", "pull", local_tag], capture_output=True, text=True
            )
            if check_reg.returncode == 0:
                _log(f"[OK] {image} already exists in local registry as {local_tag}\n")
                continue

            # 2. Check if we already have it in local docker daemon
            check_local = subprocess.run(
                ["docker", "image", "inspect", image], capture_output=True, text=True
            )
            found = check_local.returncode == 0

            if not found and import_dir and os.path.isdir(import_dir):
                # 3. Check import directory
                safe_name = image.replace("/", "_").replace(":", "_")
                tar_path = Path(import_dir) / f"{safe_name}.tar"
                if tar_path.exists():
                    _log(f"Found {tar_path} in import directory, loading...\n")
                    load = subprocess.run(
                        ["docker", "load", "-i", str(tar_path)],
                        capture_output=True,
                        text=True,
                    )
                    if load.returncode == 0:
                        found = True
                    else:
                        _log(f"[WARN] Failed to load {tar_path}: {load.stderr}\n")

            if not found:
                # 4. Pull from Docker Hub
                _log(f"Pulling {image} from Docker Hub...\n")
                pull = subprocess.run(
                    ["docker", "pull", image], capture_output=True, text=True
                )
                if pull.returncode != 0:
                    overall_ok = False
                    _log(pull.stdout or "")
                    _log(pull.stderr or "")
                    _log(f"[ERROR] docker pull failed for {image}\n")
                    continue
                found = True

            # If we have the image locally, tag and push to local registry
            if found:
                if local_tag != image:
                    tag = subprocess.run(
                        ["docker", "tag", image, local_tag],
                        capture_output=True,
                        text=True,
                    )
                    if tag.returncode != 0:
                        overall_ok = False
                        _log(tag.stdout or "")
                        _log(tag.stderr or "")
                        _log(f"[ERROR] docker tag failed for {image} to {local_tag}\n")
                        continue

                _log(f"Pushing {local_tag} to local registry...\n")
                if not _push_docker_image(local_tag, log_fn=log):
                    overall_ok = False
                    continue
                _log(f"[OK] Stored {image} in local registry as {local_tag}\n")

        return overall_ok

    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)