prole/installer/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 installer.state import InstallerState

if TYPE_CHECKING:
    from installer.milestone import Milestone

class ProleController:
    """Business logic for Prole installer, separated from UI."""
    def __init__(self, project_root: Path, verbose: bool = False):
        self.project_root = project_root
        self.verbose = verbose
        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_prole_db_version(self) -> str:
        pg_version_file = self.project_root / "conf" / "postgresql" / ".version"
        release_file = self.project_root / "prole-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} with args={args}")
        if args is None: args = []
        
        # Get paths
        source_script = self.project_root / "etc" / script_name
        
        # Determine target etc directory
        prole_home_val = (env or os.environ).get("PROLE_HOME")
        if not prole_home_val:
            prole_home = Path.home() / ".prole"
        else:
            prole_home = Path(prole_home_val).expanduser()
            
        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 conf/postgresql is copied for version detection
        target_conf = prole_home / "conf"
        source_conf = self.project_root / "conf"
        if source_conf.exists():
            if source_conf.resolve() != target_conf.resolve():
                if not target_conf.exists() or (source_conf.stat().st_mtime > target_conf.stat().st_mtime):
                    if target_conf.exists():
                        shutil.rmtree(target_conf)
                    shutil.copytree(source_conf, target_conf)

        # Ensure prole-db/.version is available in PROLE_HOME for image tagging
        source_db_version = self.project_root / "prole-db" / ".version"
        if source_db_version.exists():
            target_db_dir = prole_home / "prole-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

        # 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 config helpers in sync (prole_cfg.sh)
        try:
            source_cfg = self.project_root / "etc" / "prole_cfg.sh"
            target_cfg = target_etc / "prole_cfg.sh"
            if source_cfg.exists():
                if not target_cfg.exists() or (source_cfg.stat().st_mtime > target_cfg.stat().st_mtime):
                    shutil.copy2(source_cfg, target_cfg)
                    os.chmod(target_cfg, 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 = (env or os.environ).copy()
        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)
        
        return proc.returncode

    def run_milestones(self, milestones: Sequence[Milestone], progress_callback: Callable[[str, float], None] | None = None):
        """Execute a sequence of milestones."""
        for milestone in milestones:
            self.logger.info(f"==> Executing Milestone: {milestone.title}")
            milestone.execute(self.state, progress=progress_callback)
            self.state.mark_completed(milestone.id)