prole/knoe/core/controller.py

from __future__ import annotations
import subprocess
import threading
import os
import shutil
import logging
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Sequence

from knoe.state import InstallerState
from knoe.core.env import _deployment_mode_from_env, _expand_cfg_value, resolve_prole_home

if TYPE_CHECKING:
    from knoe.milestone import Milestone


class ProleController:
    """Business logic for Prole installer, separated from UI."""

    def __init__(self, project_root: Path, verbose: bool = False, cfg_path: Path | None = None):
        self.project_root = project_root
        self.verbose = verbose
        self.cfg_path = cfg_path
        self.logger = logging.getLogger("ProleController")
        self.state = InstallerState(controller=self)

    def check_docker_running(self) -> bool:
        """Check if Docker daemon is responsive."""
        try:
            subprocess.run(["docker", "info"], capture_output=True, check=True)
            return True
        except (subprocess.CalledProcessError, FileNotFoundError):
            return False

    def get_knoe_db_version(self) -> str:
        # Prefer per-user, per-mode version files under ~/.prole (or resolved PROLE_HOME)
        # so that k3d/k3s/k8s do not overwrite each other's DB image versions.
        cluster_env = os.environ.get("CLUSTER_ENV") or ""
        mode_key = (
            _deployment_mode_from_env(cluster_env)
            or (os.environ.get("PROLE_MODE") or "")
            or "default"
        )

        runtime_home = resolve_prole_home(env=os.environ.copy())
        mode_root = runtime_home / "modes" / mode_key

        pg_version_file = mode_root / "conf" / "postgresql" / ".version"
        release_file = mode_root / "knoe-db" / ".version"

        # Fallback to repo-relative files for legacy/dev workflows and tests.
        if not pg_version_file.exists():
            pg_version_file = self.project_root / "conf" / "postgresql" / ".version"
        if not release_file.exists():
            release_file = self.project_root / "knoe-db" / ".version"

        pg_version = pg_version_file.read_text().strip() if pg_version_file.exists() else "17.7"
        release = release_file.read_text().strip() if release_file.exists() else "43"
        if release.isdigit():
            release = release.zfill(3)
        return f"{pg_version}-{release}"

    def run_script(
        self,
        script_name: str,
        args: list[str] | None = None,
        env: dict | None = None,
        stdin_text: str | None = None,
        on_line: Callable[[str], None] | None = None,
        on_stderr_line: Callable[[str], None] | None = None,
        stderr_to_stdout: bool = True,
    ) -> int:
        """Generic runner for etc/ scripts.
        Copies the script to $PROLE_HOME/etc before running it.
        """
        self.logger.debug(
            f"run_script: {script_name} args={args} env_keys={sorted((env or {}).keys())[:20]}"
        )
        if args is None:
            args = []

        base_env = os.environ.copy()
        if env:
            try:
                base_env.update(env)
            except Exception:
                pass

        # Get paths
        source_script = self.project_root / "etc" / script_name

        # Determine target etc directory
        prole_home_val = base_env.get("PROLE_HOME")
        if not prole_home_val:
            prole_home = resolve_prole_home(env={})
        else:
            expanded_home = _expand_cfg_value(
                str(prole_home_val), env=base_env, max_depth=10
            )
            expanded_home = os.path.expanduser(expanded_home)
            prole_home = Path(expanded_home)

        target_etc = prole_home / "etc"
        target_etc.mkdir(parents=True, exist_ok=True)
        target_script = target_etc / script_name

        # Also ensure k8s resources are copied to $PROLE_HOME/k8s for reference
        target_k8s = prole_home / "k8s"
        source_k8s = self.project_root / "k8s"
        if source_k8s.exists():
            if source_k8s.resolve() != target_k8s.resolve():
                if not target_k8s.exists() or (
                    source_k8s.stat().st_mtime > target_k8s.stat().st_mtime
                ):
                    if target_k8s.exists():
                        shutil.rmtree(target_k8s)
                    shutil.copytree(source_k8s, target_k8s)

        # Ensure a minimal conf allowlist is available in PROLE_HOME.
        #
        # IMPORTANT: `$PROLE_HOME/conf` may contain curated, environment-specific
        # configuration files including the `prole.cfg` symlink entrypoint and
        # `$PROLE_CONF/<env>/prole.cfg`. Never delete or wholesale-copy the
        # entire `conf/` directory here.
        target_conf = prole_home / "conf"
        source_conf = self.project_root / "conf"
        if source_conf.exists() and source_conf.resolve() != target_conf.resolve():
            target_conf.mkdir(parents=True, exist_ok=True)

            # conf/postgresql is used for version detection.
            source_pg = source_conf / "postgresql"
            target_pg = target_conf / "postgresql"
            if source_pg.exists():
                try:
                    shutil.copytree(source_pg, target_pg, dirs_exist_ok=True)
                except Exception:
                    pass

            # Copy a small allowlist of non-environment artifacts that etc/ scripts may use.
            for name in ("database_versions.json", "port-mapping.cfg"):
                src = source_conf / name
                if not src.exists():
                    continue
                dst = target_conf / name
                try:
                    if (not dst.exists()) or (src.stat().st_mtime > dst.stat().st_mtime):
                        shutil.copy2(src, dst)
                except Exception:
                    pass

        # Ensure the db image .version is available in PROLE_HOME for image tagging
        source_db_version = self.project_root / "knoe-db" / ".version"
        if source_db_version.exists():
            target_db_dir = prole_home / "knoe-db"
            target_db_dir.mkdir(parents=True, exist_ok=True)
            target_db_version = target_db_dir / ".version"
            try:
                shutil.copy2(source_db_version, target_db_version)
            except Exception:
                pass

        # Ensure shared shell libraries are available for etc/ scripts.
        # Common core scripts source `etc/common_core_lib.sh`, which in turn sources
        # `etc/lib/shell/*` relative to `$PROLE_HOME/etc`.
        try:
            source_etc_lib_shell = self.project_root / "etc" / "lib" / "shell"
            target_etc_lib_shell = target_etc / "lib" / "shell"
            if source_etc_lib_shell.exists():
                target_etc_lib_shell.mkdir(parents=True, exist_ok=True)
                for src in source_etc_lib_shell.rglob("*"):
                    if not src.is_file():
                        continue
                    rel = src.relative_to(source_etc_lib_shell)
                    dst = target_etc_lib_shell / rel
                    dst.parent.mkdir(parents=True, exist_ok=True)
                    try:
                        if (not dst.exists()) or (src.stat().st_mtime > dst.stat().st_mtime):
                            shutil.copy2(src, dst)
                    except Exception:
                        pass
        except Exception:
            pass

        # Copy script and set permissions
        if source_script.exists():
            if source_script.resolve() != target_script.resolve():
                shutil.copy2(source_script, target_script)
            os.chmod(target_script, 0o755)

        # Keep shared etc helpers in sync.
        try:
            for helper in ["prole_cfg.sh", "common_core_lib.sh"]:
                source_helper = self.project_root / "etc" / helper
                target_helper = target_etc / helper
                if not source_helper.exists():
                    continue
                if (not target_helper.exists()) or (
                    source_helper.stat().st_mtime > target_helper.stat().st_mtime
                ):
                    shutil.copy2(source_helper, target_helper)
                    os.chmod(target_helper, 0o755)
        except Exception:
            pass

        # Run from the installed location
        cmd = ["bash", str(target_script)] + args
        self.logger.debug(f"Starting subprocess: {' '.join(cmd)}")

        # Ensure PYTHONUNBUFFERED=1 for any python scripts called within the bash script
        run_env = base_env.copy()
        # Expand PROLE_* paths so subprocesses never see literal `$HOME`/`$PROLE_HOME`.
        try:
            run_env["PROLE_HOME"] = str(prole_home)
            for k in (
                "PROLE_CONF",
                "PROLE_DATA",
                "PROLE_LOGS",
                "PROLE_SERVICE",
                "KUBECONFIG",
            ):
                raw = (run_env.get(k) or "").strip()
                if not raw or "$" not in raw:
                    continue
                expanded = _expand_cfg_value(raw, env=run_env, max_depth=10)
                run_env[k] = os.path.expanduser(expanded)
        except Exception:
            pass
        self.logger.debug(
            f"  cwd={os.getcwd()} KUBECONFIG={run_env.get('KUBECONFIG','<unset>')} NAMESPACE={run_env.get('NAMESPACE','<unset>')} PROLE_MODE={run_env.get('PROLE_MODE','<unset>')}"
        )
        run_env["PYTHONUNBUFFERED"] = "1"
        if self.verbose:
            run_env.setdefault("PROLE_VERBOSE", "1")
            run_env.setdefault("VERBOSE", "1")
            run_env.setdefault("CLICOLOR_FORCE", "1")
            run_env.setdefault("PY_COLORS", "1")
            if not run_env.get("TERM"):
                run_env["TERM"] = "xterm-256color"

        proc = subprocess.Popen(
            cmd,
            stdin=subprocess.PIPE if stdin_text else None,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT if stderr_to_stdout else subprocess.PIPE,
            text=True,
            bufsize=1,
            env=run_env,
        )

        if stdin_text and proc.stdin:
            proc.stdin.write(stdin_text)
            proc.stdin.close()

        def _read_stream(stream, handler):
            if not stream:
                return
            for line in iter(stream.readline, ""):
                if handler:
                    handler(line)

        stderr_thread = None
        if not stderr_to_stdout:
            stderr_thread = threading.Thread(
                target=_read_stream, args=(proc.stderr, on_stderr_line), daemon=True
            )
            stderr_thread.start()

        while True:
            line = proc.stdout.readline() if proc.stdout else None
            if not line and proc.poll() is not None:
                break
            if line and on_line:
                on_line(line)
        if stderr_thread:
            stderr_thread.join(timeout=2)

        self.logger.debug(f"run_script: {script_name} exited with rc={proc.returncode}")
        return proc.returncode

    def run_milestones(
        self,
        milestones: Sequence[Milestone],
        progress_callback: Callable[[str, float], None] | None = None,
    ):
        """Execute a sequence of milestones."""
        import time as _time

        for milestone in milestones:
            self.logger.info(f"==> Executing Milestone: {milestone.title}")
            t0 = _time.monotonic()
            milestone.execute(self.state, progress=progress_callback)
            elapsed = _time.monotonic() - t0
            self.logger.debug(
                f"Milestone '{milestone.title}' completed in {elapsed:.1f}s"
            )
            self.state.mark_completed(milestone.id)