alnn_rewrite/main.py

import argparse
import os
import threading
import gc
from typing import Any, Dict, Optional

# Must be set before first CUDA context initialization.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import toml
from textual import work
from textual.app import App, ComposeResult
from textual.containers import Horizontal, Vertical
from textual.widgets import (
    Button,
    Footer,
    Header,
    Input,
    Label,
    ProgressBar,
    RichLog,
    Rule,
)

from tasks import PIPELINE_TASKS, SCENARIOS, describe_protected_conflicts
from util.config_manager import handle_directory_config
from util.progress import MAX_LEVELS, ProgressTracker
from util.screens import OverwriteConfirmScreen
from util.ui_logger import PipelineLogger


class CNNHarnessApp(App[None]):
    CSS_PATH = "style.tcss"

    def __init__(
        self,
        work_dir: Optional[str] = None,
        autorun: bool = False,
        autoquit: bool = False,
        force: bool = False,
    ):
        super().__init__()
        self.stop_event = threading.Event()
        self.pause_event = threading.Event()
        self.current_config: Dict[str, Any] = {}
        self.ui_logger = PipelineLogger(self, widget_id="#live_log")

        # CLI overrides. These force the corresponding behavior ON; when not
        # passed, the effective value falls back to the loaded config's [run]
        # section (see _effective_* helpers).
        self._cli_work_dir = work_dir
        self._cli_autorun = autorun
        self._cli_autoquit = autoquit
        self._cli_force = force
        self._last_pipeline_ok = False

    # ==========================================
    # UI Layout
    # ==========================================
    def compose(self) -> ComposeResult:
        yield Header()

        with Vertical(id="main_container"):
            # --- TOP SECTION: Two-Column Configuration ---
            with Horizontal(id="config_section"):
                # Left Column: Path entry and buttons
                with Vertical(id="config_left_col"):
                    yield Label("Working Directory:")
                    yield Input(placeholder="./outputs/experiment_1", id="work_dir")
                    yield Button("Load / Init Dir", variant="primary", id="load_btn")
                    yield Button("Reload Config", variant="default", id="reload_btn")

                # Right Column: Read-only parameters display
                with Vertical(id="config_right_col"):
                    yield Label("Current Parameters:")
                    yield RichLog(id="config_display", highlight=True, markup=True)

            yield Rule()

            # --- BOTTOM SECTION: Controls, Progress, and Logs ---
            with Horizontal(id="bottom_area"):
                # --- Pipeline Controls ---
                with Vertical(id="control_panel"):
                    yield Button("Start Pipeline", variant="success", id="btn_start")

                    with Horizontal(id="active_controls"):
                        yield Button("Pause", variant="warning", id="btn_pause")
                        yield Button("Stop", variant="error", id="btn_stop")

                    # --- Progress Trackers ---
                    with Vertical(id="progress_area"):
                        for i in range(MAX_LEVELS):
                            with Horizontal(
                                id=f"progress_row_{i}", classes="progress_row"
                            ):
                                yield Label(
                                    id=f"progress_title_{i}", classes="progress_title"
                                )
                                yield ProgressBar(
                                    id=f"progress_bar_{i}",
                                    classes="progress_bar",
                                    show_eta=True,
                                )
                                yield Label(
                                    id=f"progress_stats_{i}", classes="progress_stats"
                                )

                # --- Live Log Output ---
                with Vertical(id="log_panel"):
                    with Horizontal(id="log_status_row"):
                        yield Label("STATUS: READY", id="pipeline_status_label")
                    with Horizontal(id="log_header_row"):
                        yield Label("Live Logs:", classes="section_label")
                    yield RichLog(id="live_log", highlight=True, markup=True)

        yield Footer()

    def on_mount(self) -> None:
        self.query_one("#active_controls").display = False
        self.ui_logger.info("Application initialized and ready.")

        # Script mode: a CLI work_dir auto-loads its config; autorun then starts
        # the pipeline automatically once the config is available.
        if self._cli_work_dir:
            self.query_one("#work_dir", Input).value = self._cli_work_dir
            if self._load_directory(self._cli_work_dir):
                if self._effective_autorun():
                    self.call_after_refresh(self._request_start, False)
            elif self._effective_autoquit():
                # Load failed and we're running headless: exit with error code.
                self.call_after_refresh(self._quit_after_run, False)

    # ==========================================
    # Helper Methods
    # ==========================================
    def _set_pipeline_status(self, text: str) -> None:
        try:
            self.query_one("#pipeline_status_label", Label).update(text)
        except Exception:
            pass

    def _update_config_display(self) -> None:
        config_log = self.query_one("#config_display", RichLog)
        config_log.clear()

        if not self.current_config:
            config_log.write("[italic]No configuration loaded.[/italic]")
            return

        toml_string = toml.dumps(self.current_config)
        config_log.write(toml_string)

    def _toggle_config_inputs(self, disabled: bool) -> None:
        self.query_one("#work_dir", Input).disabled = disabled
        self.query_one("#load_btn", Button).disabled = disabled
        self.query_one("#reload_btn", Button).disabled = disabled

    # ------------------------------------------
    # Script-mode helpers (autorun / autoquit / force)
    # ------------------------------------------
    def _run_cfg(self) -> Dict[str, Any]:
        run = self.current_config.get("run", {}) if self.current_config else {}
        return run if isinstance(run, dict) else {}

    def _effective_autorun(self) -> bool:
        return self._cli_autorun or bool(self._run_cfg().get("autorun", False))

    def _effective_autoquit(self) -> bool:
        return self._cli_autoquit or bool(self._run_cfg().get("autoquit", False))

    def _effective_force(self) -> bool:
        return self._cli_force or bool(self._run_cfg().get("force_overwrite", False))

    def _load_directory(self, work_dir: str) -> bool:
        """Load-or-init a work_dir's config into current_config. Returns success."""
        success, message, config_data = handle_directory_config(work_dir)
        if success:
            self.current_config = config_data
            self.current_config["work_dir"] = work_dir
            self._update_config_display()
            self.ui_logger.info(message)
            return True
        self.ui_logger.error(message)
        return False

    def _quit_after_run(self, ok: bool) -> None:
        """Exit the app with a script-friendly return code (0 ok, 1 failure)."""
        self.exit(return_code=0 if ok else 1)

    def _request_start(self, interactive: bool) -> None:
        """Validate and start the pipeline. In non-interactive (autorun) mode a
        protected-overwrite conflict aborts unless force is set, since a modal
        prompt cannot be answered by a script."""
        if not self.current_config:
            self.ui_logger.error(
                "No configuration loaded! Please load a directory first."
            )
            if not interactive and self._effective_autoquit():
                self._quit_after_run(False)
            return

        scenario_id = self.current_config.get("scenario")
        if not scenario_id or scenario_id not in SCENARIOS:
            self.ui_logger.error(
                f"Invalid or missing scenario '{scenario_id}' in config.toml!"
            )
            if not interactive and self._effective_autoquit():
                self._quit_after_run(False)
            return

        work_dir = self.current_config["work_dir"]
        # Only PROTECTED (hard-to-regenerate) outputs trigger a warning. An
        # analysis-only rerun writes nothing protected, so it proceeds silently.
        conflicts = describe_protected_conflicts(work_dir, scenario_id)

        if conflicts and not self._effective_force():
            if interactive:
                def check_overwrite_callback(proceed: bool) -> None:
                    if proceed:
                        self._execute_pipeline()
                    else:
                        self.ui_logger.error("Pipeline start cancelled by user.")

                self.app.push_screen(
                    OverwriteConfirmScreen(work_dir, conflicts),
                    check_overwrite_callback,  # pyright: ignore
                )
                return
            # Non-interactive with a protected conflict and no force: refuse.
            self.ui_logger.error(
                "Autorun aborted: scenario would overwrite protected outputs "
                f"({'; '.join(conflicts)}). Set run.force_overwrite = true or pass "
                "--force to allow."
            )
            if self._effective_autoquit():
                self._quit_after_run(False)
            return

        self._execute_pipeline()

    # ==========================================
    # Event Handlers
    # ==========================================
    def on_button_pressed(self, event: Button.Pressed) -> None:

        if event.button.id in ("load_btn", "reload_btn"):
            work_dir = self.query_one("#work_dir", Input).value
            self._load_directory(work_dir)

        elif event.button.id == "btn_start":
            self._request_start(interactive=True)

        elif event.button.id == "btn_pause":
            if self.pause_event.is_set():
                self.pause_event.clear()
                event.button.label = "Pause"
                event.button.variant = "warning"
                self._set_pipeline_status("[bold green]PIPELINE RUNNING[/bold green]")
                self.ui_logger.info("Pipeline Resumed.")
            else:
                self.pause_event.set()
                event.button.label = "Resume"
                event.button.variant = "success"
                self._set_pipeline_status("[bold yellow]PIPELINE PAUSED[/bold yellow]")
                self.ui_logger.info(
                    "Pipeline Paused. Waiting for current operation to yield..."
                )

        elif event.button.id == "btn_stop":
            self._set_pipeline_status("[bold red]STOPPING PIPELINE...[/bold red]")
            self.ui_logger.error("Stop requested! Terminating gracefully...")
            event.button.disabled = True
            self.stop_event.set()

    # ==========================================
    # Pipeline Execution
    # ==========================================
    def _execute_pipeline(self) -> None:
        self._toggle_config_inputs(disabled=True)
        self.query_one("#btn_start").display = False
        self.query_one("#active_controls").display = True

        self._set_pipeline_status("[bold green]PIPELINE RUNNING[/bold green]")

        self.stop_event.clear()
        self.pause_event.clear()

        btn_pause = self.query_one("#btn_pause", Button)
        btn_pause.label = "Pause"
        btn_pause.variant = "warning"

        self.run_background_pipeline(self.current_config.copy())

    @work(exclusive=True, thread=True)
    def run_background_pipeline(self, config: Dict[str, Any]) -> None:
        scenario_id = config["scenario"]
        steps = SCENARIOS[scenario_id]["tasks"]
        scenario_label = SCENARIOS[scenario_id]["label"]

        root_tracker = ProgressTracker(self, level=0)
        pipeline_state: Dict[str, Any] = {
            "events": {"stop": self.stop_event, "pause": self.pause_event}
        }
        self._last_pipeline_ok = False

        try:
            self.ui_logger.info(f"Initiating Scenario: {scenario_label}")
            total_tasks = len(steps)

            root_tracker.reset()
            root_tracker.set_title("Pipeline Status")
            root_tracker.update(total=total_tasks, advance=0)

            for i, task_id in enumerate(steps):
                if self.stop_event.is_set():
                    break

                task_name = PIPELINE_TASKS[task_id]["task_name"]
                task_func = PIPELINE_TASKS[task_id]["task_func"]

                task_logger = self.ui_logger.get_task_logger(task_name)
                self.ui_logger.info(f"Starting Task {i + 1}/{total_tasks}: {task_name}")

                task_tracker = root_tracker.get_sub_tracker(task_name)

                task_func(task_tracker, task_logger, config, pipeline_state)

                if not self.stop_event.is_set():
                    task_logger.info("Task completed.")
                    root_tracker.update(advance=1)

            if self.stop_event.is_set():
                self.call_from_thread(
                    self._set_pipeline_status, "[bold red]PIPELINE STOPPED[/bold red]"
                )
                self.ui_logger.error("Pipeline stopped by user.")
                root_tracker.reset()
            else:
                self.call_from_thread(
                    self._set_pipeline_status,
                    "[bold blue]PIPELINE COMPLETE[/bold blue]",
                )
                self.ui_logger.info("All tasks finished successfully!")
                self._last_pipeline_ok = True

        except InterruptedError as e:
            self.ui_logger.error(f"STOPPED: {str(e)}")
            self.call_from_thread(
                self._set_pipeline_status, "[bold red]PIPELINE STOPPED[/bold red]"
            )
            root_tracker.reset()

        except Exception as e:
            self.ui_logger.file_logger.error(
                f"Pipeline failed: {str(e)}", exc_info=True
            )
            self.ui_logger.error(f"ERROR: {str(e)}")
            self.call_from_thread(
                self._set_pipeline_status, "[bold red]PIPELINE FAILED[/bold red]"
            )

        finally:
            pipeline_state.clear()
            gc.collect()
            self.call_from_thread(self._reset_ui)
            # Script mode: quit once the run finishes, with a return code that
            # reflects success/failure so it can be used in shell pipelines.
            if self._effective_autoquit():
                self.call_from_thread(self._quit_after_run, self._last_pipeline_ok)

    def _reset_ui(self) -> None:
        self.query_one("#btn_stop", Button).disabled = False
        self.query_one("#btn_start").display = True
        self.query_one("#active_controls").display = False
        self._toggle_config_inputs(disabled=False)


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "ALNN evaluation harness (TUI). Pass a working directory to auto-load "
            "it; combine with autorun/autoquit (CLI flags or the config [run] "
            "section) to use it as a simple script."
        )
    )
    parser.add_argument(
        "work_dir",
        nargs="?",
        default=None,
        help="Working/output directory to load on startup (its config.toml is "
        "loaded, or the default config is copied in if absent).",
    )
    parser.add_argument(
        "--autorun",
        action="store_true",
        help="Start the pipeline automatically once the config loads "
        "(overrides run.autorun).",
    )
    parser.add_argument(
        "--autoquit",
        action="store_true",
        help="Exit the app when the pipeline finishes, with a 0/1 return code "
        "(overrides run.autoquit).",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="In autorun, overwrite protected outputs (models/evaluations) "
        "without prompting (overrides run.force_overwrite).",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = _parse_args()
    app = CNNHarnessApp(
        work_dir=args.work_dir,
        autorun=args.autorun,
        autoquit=args.autoquit,
        force=args.force,
    )
    app.run()
    raise SystemExit(app.return_code or 0)