alnn_rewrite/main.py

import os
import threading
from typing import Any, Dict

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,
    Select,
    TabbedContent,
    TabPane,
)

from util.progress import ProgressTracker
from util.screens import OverwriteConfirmScreen
from util.tasks import PIPELINE_TASKS, SCENARIOS
from util.ui_logger import PipelineLogger


# ==========================================
# TUI Application
# ==========================================
class CNNHarnessApp(App):
    CSS_PATH = "style.tcss"

    def __init__(self):
        super().__init__()
        self.stop_event = threading.Event()
        self.pause_event = threading.Event()

        # Initialize our consolidated, thread-aware logger
        self.ui_logger = PipelineLogger(self, widget_id="#live_log")

    def compose(self) -> ComposeResult:
        yield Header()

        with Vertical(id="main_container"):
            # --- TOP SECTION: Scrollable Configuration Tabs ---
            with TabbedContent(initial="tab-general"):
                with TabPane("General Config", id="tab-general"):
                    with Horizontal(classes="row"):
                        with Vertical(classes="input_group"):
                            yield Label("Load Configuration:")
                            yield Input(placeholder="./config.toml", id="config_path")
                        yield Button("Load Config", variant="primary", id="load_btn")

                    with Horizontal(classes="row"):
                        with Vertical(classes="input_group"):
                            yield Label("Save Directory:", classes="input_label")
                            yield Input(
                                placeholder="./models/experiment_1", id="save_dir"
                            )
                        with Vertical(classes="input_group"):
                            yield Label("Batch Size:", classes="input_label")
                            yield Input(placeholder="32", id="batch_size")

            # --- BOTTOM SECTION: Controls, Progress, and Logs ---
            with Horizontal(id="bottom_area"):
                with Vertical(id="control_panel"):
                    with Horizontal(id="status_row"):
                        scenario_options = [
                            (data["label"], sc_id) for sc_id, data in SCENARIOS.items()
                        ]
                        yield Select(
                            scenario_options,
                            prompt="Select Pipeline Scenario...",
                            id="scenario_select",
                        )
                        yield Label("STATUS: READY", id="pipeline_status_label")

                    yield Rule()

                    yield Button(
                        "Start Pipeline & Save Config",
                        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")

                    yield Rule()

                    with Vertical(id="progress_area"):
                        for i in range(5):
                            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"
                                )

                with Vertical(id="log_panel"):
                    yield Label("Live Logs:")
                    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.")

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

    # ==========================================
    # Configuration Management Methods
    # ==========================================

    def _load_config_file(self, filepath: str) -> Dict[str, Any]:
        with open(filepath, "r") as f:
            return toml.load(f)

    def _save_config_file(self, filepath: str, config: Dict[str, Any]) -> None:
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        with open(filepath, "w") as f:
            toml.dump(config, f)

    def _populate_ui_from_config(self, config: Dict[str, Any]) -> None:
        self.query_one("#save_dir", Input).value = str(config.get("save_dir", ""))
        self.query_one("#batch_size", Input).value = str(config.get("batch_size", ""))

        loaded_scenario = config.get("scenario", "")
        if loaded_scenario in SCENARIOS:
            self.query_one("#scenario_select", Select).value = loaded_scenario

    def _gather_config_from_ui(self) -> Dict[str, Any]:
        scenario_select = self.query_one("#scenario_select", Select)

        scenario_id = (
            scenario_select.value if type(scenario_select.value) is str else None
        )

        config = {
            "save_dir": self.query_one("#save_dir", Input).value or "./outputs/default",
            "batch_size": self.query_one("#batch_size", Input).value or "32",
            "scenario": scenario_id,
        }

        if scenario_id and scenario_id in SCENARIOS:
            config["steps"] = SCENARIOS[scenario_id]["tasks"]
        else:
            config["steps"] = []

        return config

    # ==========================================
    # Event Handlers
    # ==========================================

    def on_button_pressed(self, event: Button.Pressed) -> None:
        if event.button.id == "load_btn":
            config_path = self.query_one("#config_path", Input).value
            if not config_path:
                self.ui_logger.error("Please specify a config path to load.")
                return

            try:
                loaded_config = self._load_config_file(config_path)
                self._populate_ui_from_config(loaded_config)
                self.ui_logger.info(f"Successfully loaded config from {config_path}")
            except Exception as e:
                self.ui_logger.error(f"Failed to load config: {str(e)}")

        elif event.button.id == "btn_start":
            config_dict = self._gather_config_from_ui()

            if not config_dict.get("scenario"):
                self.ui_logger.error(
                    "No scenario selected! Please select a scenario from the dropdown."
                )
                return

            save_dir = config_dict["save_dir"]

            if os.path.exists(save_dir) and os.listdir(save_dir):

                def check_overwrite_callback(proceed: bool) -> None:
                    if proceed:
                        self._execute_pipeline(config_dict)
                    else:
                        self.ui_logger.error(
                            "Pipeline start cancelled by user (folder non-empty)."
                        )

                self.app.push_screen(
                    OverwriteConfirmScreen(save_dir),
                    check_overwrite_callback,  # pyright: ignore
                )
            else:
                self._execute_pipeline(config_dict)

        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, config_dict: Dict[str, Any]) -> None:
        self.query_one(TabbedContent).disabled = True
        self.query_one("#scenario_select").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()
        self.query_one("#btn_pause", Button).label = "Pause"
        self.query_one("#btn_pause", Button).variant = "warning"

        self.run_background_pipeline(config_dict)

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

        try:
            config_out_path = os.path.join(save_dir, "config.toml")
            self._save_config_file(config_out_path, config)

            self.ui_logger.info(f"Initiating Scenario: {scenario_label}")
            self.ui_logger.info(f"Configuration saved to {config_out_path}")

            total_tasks = len(steps)

            pipeline_state: Dict[str, Any] = {
                "events": {"stop": self.stop_event, "pause": self.pause_event}
            }

            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"]

                # Generate a task-specific logger from our PipelineLogger
                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!")

        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:
            self.call_from_thread(self._reset_ui)

    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.query_one(TabbedContent).disabled = False
        self.query_one("#scenario_select").disabled = False


if __name__ == "__main__":
    app = CNNHarnessApp()
    app.run()