prole/installer/ui/screens/database_options.py

"""Database Options screen for selecting distribution, version and extensions."""

import json
import os
import re
import threading
from pathlib import Path
import tkinter as tk
from tkinter import ttk, messagebox
from installer import screen as ui
from installer.core.env import PROJECT_ROOT, _safe_str

class DatabaseOptionsScreenMixin:
    """Mixin for the Database Options screen."""

    def _init_database_options_state(self):
        # Load saved inputs (available when prole.cfg was previously saved)
        _saved = getattr(self, "prole_cfg_data", {}).get("Inputs", {})

        self.db_at_rest_encryption = tk.BooleanVar(value=True)
        _saved_dist = _saved.get("database_options.distribution", "percona")
        self.db_distribution = tk.StringVar(value=_saved_dist or "percona")
        _saved_vtype = _saved.get("database_options.version_type", "v18")
        self.db_version_type = tk.StringVar(value=_saved_vtype or "v18")
        # Derive selected version from version_type; default to pg18
        _vtype_to_ver = {"stable": "15", "current": "16", "latest": "17", "v18": "18"}
        self.db_selected_version = tk.StringVar(
            value=_vtype_to_ver.get(self.db_version_type.get(), "18")
        )

        # Extensions state: dict of id -> BooleanVar
        self.db_extensions = {}

        self.extensions_list = [
            {"id": "postgis", "name": "PostGIS", "description": "Spatial and geographic objects for PostgreSQL"},
            {"id": "pgvector", "name": "pgvector", "description": "Vector similarity search"},
            {"id": "pgcrypto", "name": "pgcrypto", "description": "Cryptographic functions", "in_contrib": True},
            {"id": "pgaudit", "name": "pgAudit", "description": "Audit logging"},
            {"id": "pg_repack", "name": "pg_repack", "description": "Reorganize tables with minimal locks"},
            {"id": "pg_stat_statements", "name": "pg_stat_statements", "description": "Track execution statistics", "in_contrib": True},
            {"id": "pg_buffercache", "name": "pg_buffercache", "description": "Examine shared buffer cache", "in_contrib": True},
            {"id": "pg_freespacemap", "name": "pg_freespacemap", "description": "Examine free space map", "in_contrib": True},
            {"id": "pgrowlocks", "name": "pgrowlocks", "description": "Show row-level locking information", "in_contrib": True},
            {"id": "postgres_fdw", "name": "postgres_fdw", "description": "Foreign-data wrapper for PostgreSQL", "in_contrib": True},
            {"id": "dblink", "name": "dblink", "description": "Connect to other databases", "in_contrib": True},
            {"id": "pg_stat_monitor", "name": "pg_stat_monitor", "description": "Statistics collector (Percona only)", "percona_only": True},
            {"id": "pgbadger", "name": "pgBadger", "description": "Log analyzer (Package only, no extension)", "no_extension": True},
        ]

        # Default enabled extensions (all on for percona v18)
        default_enabled = ["postgis", "pgvector", "pgcrypto", "pgaudit", "pg_repack", "pg_stat_statements", "pg_buffercache", "pg_freespacemap", "pgrowlocks", "postgres_fdw", "dblink", "pg_stat_monitor", "pgbadger"]

        for ext in self.extensions_list:
            saved_key = f"database_options.ext.{ext['id']}"
            if saved_key in _saved:
                is_enabled = _saved[saved_key].lower() == "true"
            else:
                is_enabled = ext["id"] in default_enabled
            self.db_extensions[ext["id"]] = tk.BooleanVar(value=is_enabled)

        self.db_versions_data = {}
        self._load_database_versions()

    def _load_database_versions(self):
        # Trigger refresh script
        def refresh():
            try:
                prole_home = self.prole_cfg_data.get("System Environment", {}).get("PROLE_HOME")
                if not prole_home:
                    prole_home = str(Path.home() / ".prole")
                
                self.controller.run_script("init_database_options.sh", env={"PROLE_HOME": prole_home})
                # Reload after script finishes
                version_file = Path(prole_home) / "conf" / "database_versions.json"
                if version_file.exists():
                    with open(version_file, "r") as f:
                        self.db_versions_data = json.load(f)
                    # We might need to schedule a UI refresh if the screen is already visible
                    if hasattr(self, "bg_canvas"):
                        self.root.after(0, self._refresh_database_options_ui)
            except Exception as e:
                print(f"Error refreshing database versions: {e}")

        threading.Thread(target=refresh, daemon=True).start()

        prole_home = self.prole_cfg_data.get("System Environment", {}).get("PROLE_HOME")
        if not prole_home:
            prole_home = str(Path.home() / ".prole")
        version_file = Path(prole_home) / "conf" / "database_versions.json"
        if version_file.exists():
            try:
                with open(version_file, "r") as f:
                    self.db_versions_data = json.load(f)
            except Exception:
                pass
        
        if not self.db_versions_data:
            # Fallback defaults
            self.db_versions_data = {
                "postgresql": {"stable": "15", "current": "16", "latest": "17"},
                "percona": {"stable": "15", "current": "16", "latest": "17", "v18": "18"}
            }

    def _render_database_options_page(self):
        self._render_title("Database Options", y=150)
        self._render_paragraph(
            "Configure your database distribution, version, and extensions. At-rest encryption requires Percona Distribution.",
            y=200,
        )

        x_label = 48
        x_field = 300
        y = 260

        # At-Rest Encryption
        self._canvas_items.append(
            ui.canvas_text(
                self,
                x_label,
                y,
                "At-Rest Encryption:",
                fill="black",
                font=("SF Pro Text", 12, "bold"),
            )
        )

        def on_encryption_toggle():
            if self.db_at_rest_encryption.get():
                self.db_distribution.set("percona")
            else:
                self.db_distribution.set("postgresql")
            self._refresh_database_options_ui()
            self._save_prole_cfg()

        enc_cb = tk.Checkbutton(
            self.bg_canvas,
            text="Enabled (Uses Percona Distribution)",
            variable=self.db_at_rest_encryption,
            command=on_encryption_toggle,
            bg="white",
            fg="black",
            activebackground="white",
            selectcolor="white",
            highlightbackground="#F5F5DC",
            highlightthickness=0,
            relief="flat",
            font=("SF Pro Text", 11),
        )
        enc_window = self.bg_canvas.create_window(x_field, y - 10, window=enc_cb, anchor="nw")
        self._canvas_items.append(enc_window)
        self._overlay_widgets.append(enc_cb)

        y += 50
        # Distribution
        self._canvas_items.append(
            ui.canvas_text(
                self,
                x_label,
                y,
                "Distribution:",
                fill="black",
                font=("SF Pro Text", 12, "bold"),
            )
        )
        dist_label = tk.Label(
            self.bg_canvas,
            textvariable=self.db_distribution,
            bg="white",
            fg="black",
            font=("SF Pro Text", 11),
        )
        dist_window = self.bg_canvas.create_window(x_field, y - 10, window=dist_label, anchor="nw")
        self._canvas_items.append(dist_window)
        self._overlay_widgets.append(dist_label)

        y += 50
        # Version Selection (Combobox)
        self._canvas_items.append(
            ui.canvas_text(
                self,
                x_label,
                y,
                "Database Version:",
                fill="black",
                font=("SF Pro Text", 12, "bold"),
            )
        )

        self.db_version_combo = ttk.Combobox(
            self.bg_canvas,
            textvariable=self.db_selected_version,
            state="readonly",
            width=30,
            font=("SF Pro Text", 11),
        )
        v_win = self.bg_canvas.create_window(x_field, y - 10, window=self.db_version_combo, anchor="nw")
        self._canvas_items.append(v_win)
        self._overlay_widgets.append(self.db_version_combo)
        self.db_version_combo.bind("<<ComboboxSelected>>", lambda _: self._save_prole_cfg())

        y += 50
        # Extensions
        self._canvas_items.append(
            ui.canvas_text(
                self,
                x_label,
                y,
                "Extensions Browser:",
                fill="black",
                font=("SF Pro Text", 12, "bold"),
            )
        )
        
        ext_frame = tk.Frame(
            self.bg_canvas,
            bg="white",
            highlightbackground="#E0E0E0",
            highlightthickness=1,
        )
        ext_window = self.bg_canvas.create_window(
            x_label,
            y + 30,
            window=ext_frame,
            anchor="nw",
            width=900,
            height=300,
        )
        self._canvas_items.append(ext_window)
        self._overlay_widgets.append(ext_frame)

        # Extension table (browser style)
        columns = ("enabled", "name", "description")
        self.ext_tree = ttk.Treeview(ext_frame, columns=columns, show="headings", height=10)
        self.ext_tree.heading("enabled", text="Enabled")
        self.ext_tree.heading("name", text="Extension Name")
        self.ext_tree.heading("description", text="Description")
        self.ext_tree.column("enabled", width=80, anchor="center")
        self.ext_tree.column("name", width=200, anchor="w")
        self.ext_tree.column("description", width=600, anchor="w")
        self.ext_tree.pack(side="left", fill="both", expand=True)

        scroll = ttk.Scrollbar(ext_frame, orient="vertical", command=self.ext_tree.yview)
        self.ext_tree.configure(yscrollcommand=scroll.set)
        scroll.pack(side="right", fill="y")
        self._overlay_widgets.append(self.ext_tree)
        self._overlay_widgets.append(scroll)

        def on_ext_click(event):
            item = self.ext_tree.identify_row(event.y)
            if item:
                vals = self.ext_tree.item(item, "values")
                ext_id = vals[1].lower()
                if ext_id in self.db_extensions:
                    current = self.db_extensions[ext_id].get()
                    self.db_extensions[ext_id].set(not current)
                    self._refresh_extensions_table()

        self.ext_tree.bind("<ButtonRelease-1>", on_ext_click)

        self._refresh_extensions_table()
        self._refresh_database_options_ui()

    def _refresh_extensions_table(self):
        for item in self.ext_tree.get_children():
            self.ext_tree.delete(item)
        
        for ext in self.extensions_list:
            enabled = "[ ✓ ]" if self.db_extensions.get(ext["id"], tk.BooleanVar()).get() else "[   ]"
            self.ext_tree.insert("", "end", values=(enabled, ext["name"], ext["description"]))

    def _refresh_database_options_ui(self):
        dist = self.db_distribution.get()
        v_data = self.db_versions_data.get(dist, {})
        
        combo_values = []
        if "latest" in v_data:
            combo_values.append(f"{v_data['latest']} (Latest)")
        if "current" in v_data:
            combo_values.append(f"{v_data['current']} (Current)")
        if "stable" in v_data:
            combo_values.append(f"{v_data['stable']} (Stable)")
        if dist == "percona" and "v18" in v_data:
            combo_values.append(f"{v_data['v18']} (Percona 18)")
        
        if hasattr(self, "db_version_combo"):
            self.db_version_combo.configure(values=combo_values)
            
            # Match current selection to one of the values if possible
            current_ver = self.db_selected_version.get()
            found = False
            for val in combo_values:
                if val.startswith(current_ver):
                    self.db_selected_version.set(val)
                    found = True
                    break
            if not found and combo_values:
                self.db_selected_version.set(combo_values[0])

    def _generate_prole_db_dockerfile(self):
        dist = _safe_str(self.db_distribution.get())
        version_full = _safe_str(self.db_selected_version.get())
        # Extract major version from "17 (Latest)" etc.
        version = version_full.split()[0] if version_full else "17"
        
        template_file = PROJECT_ROOT / "prole-db" / f"Dockerfile.{dist}.template"
        if not template_file.exists():
            messagebox.showerror("Error", f"Template not found: {template_file}")
            return False
        
        template = template_file.read_text()
        
        # Prepare extension steps
        install_steps_list = []
        create_steps_list = []
        
        for ext in self.extensions_list:
            if self.db_extensions.get(ext["id"], tk.BooleanVar()).get():
                ext_id = ext["id"]
                
                # Skip percona-only extensions if not percona
                if ext.get("percona_only") and dist != "percona":
                    continue
                
                # 1. Handle Installation Steps
                if not ext.get("in_contrib"):
                    pkg = ""
                    if dist == "percona":
                        if ext_id == "postgis":
                            pkg = f"percona-postgresql-{version}-postgis-3"
                        elif ext_id == "pg_repack":
                            # Percona uses 'repack' instead of 'pg_repack' in package name
                            pkg = f"percona-postgresql-{version}-repack"
                        elif ext_id == "pgbadger":
                            pkg = "percona-pgbadger"
                        elif ext_id == "pg_stat_monitor":
                            pkg = f"percona-pg-stat-monitor{version}"
                        else:
                            pkg = f"percona-postgresql-{version}-{ext_id}"
                    else:
                        # Standard PostgreSQL
                        if ext_id == "postgis":
                            pkg = f"postgresql-{version}-postgis-3"
                        elif ext_id == "pg_repack":
                            pkg = f"postgresql-{version}-repack"
                        elif ext_id == "pgbadger":
                            pkg = "pgbadger"
                        else:
                            pkg = f"postgresql-{version}-{ext_id}"
                    
                    if pkg:
                        install_steps_list.append(f"  apt-get install -y --no-install-recommends {pkg};")
                
                # 2. Handle Extension Creation (SQL)
                if not ext.get("no_extension"):
                    sql_name = ext_id
                    if ext_id == "pgvector":
                        sql_name = "vector"
                    
                    if sql_name == 'postgis':
                        create_steps_list.append('  echo "    CREATE EXTENSION IF NOT EXISTS postgis;" >> /docker-entrypoint-initdb.d/20_create_extensions.sh;')
                        create_steps_list.append('  echo "    CREATE EXTENSION IF NOT EXISTS postgis_topology;" >> /docker-entrypoint-initdb.d/20_create_extensions.sh;')
                    else:
                        create_steps_list.append(f'  echo "    CREATE EXTENSION IF NOT EXISTS {sql_name};" >> /docker-entrypoint-initdb.d/20_create_extensions.sh;')

        # Join with backslash continuation and ensure no trailing newline to avoid blank lines in Dockerfile.
        # If empty, we use a no-op ':' to maintain valid shell syntax across continuations.
        install_steps = " \\\n".join(install_steps_list) + " \\" if install_steps_list else "  :; \\"
        create_steps = " \\\n".join(create_steps_list) + " \\" if create_steps_list else "  :; \\"

        content = template.replace("{{MAJOR_VERSION}}", version)
        content = content.replace("{{EXTENSION_INSTALL_STEPS}}", install_steps)
        content = content.replace("{{EXTENSION_CREATE_STEPS}}", create_steps)
        
        dockerfile_path = PROJECT_ROOT / "prole-db" / "Dockerfile"
        dockerfile_path.write_text(content)
        
        # Increment version
        self._increment_prole_db_version(version)
        return True

    def _increment_prole_db_version(self, major_minor):
        version_file = PROJECT_ROOT / "prole-db" / ".version"
        current_counter = 0
        if version_file.exists():
            try:
                val = version_file.read_text().strip()
                if val.isdigit():
                    current_counter = int(val)
            except Exception:
                pass
        
        new_counter = current_counter + 1
        version_file.write_text(str(new_counter))
        
        pg_version_file = PROJECT_ROOT / "conf" / "postgresql" / ".version"
        pg_version_file.parent.mkdir(parents=True, exist_ok=True)
        pg_version_file.write_text(major_minor)
        
        return f"{major_minor}-{str(new_counter).zfill(3)}"