###############################################################################
# (c) Copyright 2026 CERN for the benefit of the LHCb Collaboration           #
#                                                                             #
# This software is distributed under the terms of the GNU General Public      #
# Licence version 3 (GPL Version 3), copied verbatim in the file "COPYING".   #
#                                                                             #
# In applying this licence, CERN does not waive the privileges and immunities #
# granted to it by virtue of its status as an Intergovernmental Organization  #
# or submit itself to any jurisdiction.                                       #
###############################################################################
"""
Textual-based TUI for CWL workflow execution monitoring.

Provides a real-time terminal UI for monitoring CWL workflow execution with:
- Workflow status and progress
- Live stdout from CWL executor
- Current step log file viewer
- Workflow graph visualization
- Output file tracking
- Memory usage monitoring
"""

import asyncio
import re
import time
from pathlib import Path
from typing import Optional

import yaml
from textual.app import App, ComposeResult
from textual.containers import Container, VerticalScroll
from textual.reactive import reactive
from textual.widgets import (
    DataTable,
    Footer,
    Header,
    RichLog,
    Static,
    TabbedContent,
    TabPane,
    Tree,
)
from textual_plotext import PlotextPlot


class WorkflowStatus(Static):
    """Header showing workflow execution status"""

    status = reactive("Starting...")
    elapsed_time = reactive("0s")
    current_step = reactive("Initializing...")
    memory_usage = reactive("")
    output_files_info = reactive("")
    max_memory = reactive("")

    def render(self) -> str:
        status_emoji = "🔄" if self.status == "Running" else "⏳"

        content = f"{status_emoji} [bold cyan]{self.status}[/bold cyan]"
        content += f"  │  ⏱️  {self.elapsed_time}"
        content += f"  │  📋 {self.current_step}"

        if self.output_files_info:
            content += f"  │  📂 {self.output_files_info}"

        if self.memory_usage:
            content += f"  │  💾 {self.memory_usage}"

        if self.max_memory:
            content += f"  │  📊 {self.max_memory}"

        return content


class TabbedLogViewer(Container):
    """Tabbed log file viewer for multiple steps"""

    DEFAULT_CSS = """
    TabbedLogViewer {
        border: solid $accent;
        height: 1fr;
    }
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.border_title = "Logs"
        self.log_tabs = {}  # {log_file_path: state_dict}
        self.tabbed_content = None

    def compose(self) -> ComposeResult:
        """Build the tabbed interface"""
        yield TabbedContent()

    def on_mount(self):
        """Get reference to TabbedContent after mounting"""
        self.tabbed_content = self.query_one(TabbedContent)

    async def add_log_file(self, log_file: Path, step_name: str):
        """Add a new log file tab"""
        if not log_file or log_file in self.log_tabs:
            return

        if not self.tabbed_content:
            return

        # Create a short tab label
        tab_label = step_name[-20:] if len(step_name) > 20 else step_name
        tab_id = f"log_tab_{len(self.log_tabs) + 1}"

        # Create a simple TabPane
        pane = TabPane(tab_label, id=tab_id)
        await self.tabbed_content.add_pane(pane)

        # Mount a RichLog widget - designed for scrollable log content
        log_widget = RichLog(max_lines=500, auto_scroll=True, wrap=True)
        await pane.mount(log_widget)

        # Store reference with file tracking state
        self.log_tabs[log_file] = {
            "tab_id": tab_id,
            "widget": log_widget,
            "position": 0,
        }

        # Switch to the new tab
        self.tabbed_content.active = tab_id

    def refresh_all_logs(self):
        """Refresh all log file contents"""
        for log_file, state in self.log_tabs.items():
            if not log_file.exists():
                continue

            try:
                with open(log_file, "r") as f:
                    f.seek(state["position"])
                    new_lines = f.readlines()
                    if new_lines:
                        for line in new_lines:
                            state["widget"].write(line.rstrip())
                        state["position"] = f.tell()
            except Exception:
                pass


class MemoryChart(PlotextPlot):
    """Memory consumption chart using plotext for proper time series visualization"""

    DEFAULT_CSS = """
    MemoryChart {
        border: solid $primary;
        height: 100;
    }
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.border_title = "Memory Usage (RSS+Swap)"
        self._times = []  # Time in seconds from start
        self._rss_values = []  # RSS in GB
        self._swap_values = []  # Swap in GB
        self._total_values = []  # RSS+Swap in GB
        self.max_memory = 0

    def on_mount(self):
        """Initialize the plot"""
        self._replot()

    def add_data_point(self, time_seconds: float, rss_kb: int, swap_kb: int = 0):
        """Add a new memory data point

        Args:
            time_seconds: Time in seconds from workflow start (relative time)
            rss_kb: RSS memory in KB
            swap_kb: Swap memory in KB
        """
        # Convert to GB
        rss_gb = rss_kb / (1024 * 1024)
        swap_gb = swap_kb / (1024 * 1024)
        total_gb = rss_gb + swap_gb

        # Store the data - time_seconds is already relative
        self._times.append(time_seconds)
        self._rss_values.append(rss_gb)
        self._swap_values.append(swap_gb)
        self._total_values.append(total_gb)

        if total_gb > self.max_memory:
            self.max_memory = total_gb

        # Limit data points to prevent memory issues
        max_points = 200
        if len(self._times) > max_points:
            self._times = self._times[-max_points:]
            self._rss_values = self._rss_values[-max_points:]
            self._swap_values = self._swap_values[-max_points:]
            self._total_values = self._total_values[-max_points:]

        self._replot()

    def _replot(self):
        """Redraw the plot with current data"""
        self.plt.clear_figure()

        if not self._times:
            self.plt.title("Waiting for memory data...")
            self.refresh()
            return

        # Plot RSS and Swap as stacked area effect using lines
        self.plt.plot(self._times, self._rss_values, label="RSS", marker="braille")
        if any(v > 0 for v in self._swap_values):
            self.plt.plot(
                self._times, self._total_values, label="RSS+Swap", marker="braille"
            )

        self.plt.xlabel("Time (s)")
        self.plt.ylabel("Memory (GB)")
        # self.plt.title("Memory Usage")

        self.refresh()


class WorkflowGraph(VerticalScroll):
    """Workflow structure visualization"""

    DEFAULT_CSS = """
    WorkflowGraph {
        border: solid $success;
        padding: 1;
    }
    """

    def __init__(self, workflow_steps: list, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.border_title = "Workflow Steps"
        self.workflow_steps = workflow_steps
        self._completed_steps = set()
        self._current_step_id = None
        self._tree = None
        self._step_outputs = {}  # Track files per step: {step_id: [(file, size), ...]}
        self._last_outputs_hash = None  # Track hash to detect changes

    def compose(self) -> ComposeResult:
        """Build the workflow graph tree"""
        self._tree = Tree("Workflow", id="workflow-tree")
        self._tree.root.expand()
        self._rebuild_tree()
        yield self._tree

    def _rebuild_tree(self):
        """Rebuild the tree with current status"""
        if not self._tree:
            return

        # Clear existing tree
        self._tree.clear()
        self._tree.root.expand()

        def add_steps_to_tree(parent_node, steps):
            for step in steps:
                step_id = step["id"]
                step_label = step.get("label", step_id)
                full_id = step.get("full_id", step_id)

                # Add status indicator
                if step_id in self._completed_steps:
                    icon = "✅"
                elif step_id == self._current_step_id:
                    icon = "▶️"
                else:
                    icon = "⏸️"

                node = parent_node.add(f"{icon} {step_label}", data=step)
                node.expand()

                # Add output files for this step
                step_files = []
                # Try multiple matching strategies (same as old Rich UI)
                for key in [step_id, full_id]:
                    if key in self._step_outputs:
                        step_files.extend(self._step_outputs[key])

                # Also try matching keys that contain step_id
                for key in self._step_outputs.keys():
                    if step_id in key or key in step_id:
                        if key not in [step_id, full_id]:
                            step_files.extend(self._step_outputs[key])

                # Add file nodes (deduplicate)
                seen_files = set()
                for file_path, size in step_files:
                    if file_path not in seen_files:
                        seen_files.add(file_path)
                        # Format size
                        if size > 1024 * 1024:
                            size_str = f"{size / (1024 * 1024):.1f} MB"
                        elif size > 1024:
                            size_str = f"{size / 1024:.1f} KB"
                        else:
                            size_str = f"{size} B"
                        node.add_leaf(f"📄 {file_path} ({size_str})")

                # Recursively add substeps
                if step.get("substeps"):
                    add_steps_to_tree(node, step["substeps"])

        add_steps_to_tree(self._tree.root, self.workflow_steps)

    def mark_step_running(self, step_id: str):
        """Mark a step as currently running"""
        self._current_step_id = step_id
        self._rebuild_tree()

    def mark_step_complete(self, step_id: str):
        """Mark a step as completed"""
        self._completed_steps.add(step_id)
        self._rebuild_tree()

    def update_step_outputs(self, step_outputs: dict):
        """Update the output files for all steps"""
        # Only rebuild if outputs actually changed
        # Use a fast hash-based comparison instead of JSON serialization
        try:
            # Create a simple hash from the dict structure
            outputs_repr = str(
                sorted((k, tuple(sorted(v))) for k, v in step_outputs.items())
            )
            outputs_hash = hash(outputs_repr)

            if outputs_hash != self._last_outputs_hash:
                self._last_outputs_hash = outputs_hash
                self._step_outputs = step_outputs.copy()
                self._rebuild_tree()
        except (TypeError, ValueError):
            # If hashing fails, always rebuild (safe fallback)
            self._step_outputs = step_outputs.copy()
            self._rebuild_tree()


class OutputFiles(DataTable):
    """Table showing output files and their sizes"""

    DEFAULT_CSS = """
    OutputFiles {
        border: solid $warning;
        height: 12;
    }
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.border_title = "Output Files"
        self.files = {}

    def on_mount(self):
        """Set up the data table columns"""
        self.add_column("File", key="file")
        self.add_column("Size", key="size", width=15)
        self.add_column("Step", key="step", width=20)

    def update_files(self, files_dict: dict, step_outputs: dict):
        """Update the files table"""
        # Clear existing rows
        self.clear()

        # Add rows for each file
        for file_path, size in sorted(files_dict.items()):
            size_str = self._format_size(size)

            # Find which step produced this file
            step_name = "Unknown"
            for step, file_list in step_outputs.items():
                if any(fp == file_path for fp, _ in file_list):
                    step_name = step
                    break

            self.add_row(file_path, size_str, step_name)

    @staticmethod
    def _format_size(size_bytes: int) -> str:
        """Format file size in human-readable format"""
        for unit in ["B", "KB", "MB", "GB"]:
            if size_bytes < 1024:
                return f"{size_bytes:.1f} {unit}"
            size_bytes /= 1024
        return f"{size_bytes:.1f} TB"


class CWLMonitorApp(App):
    """Textual app for monitoring CWL workflow execution"""

    CSS = """
    Screen {
        layout: grid;
        grid-size: 2 4;
        grid-rows: auto 10 1fr 1fr;
        grid-columns: 30 1fr;
        height: 100vh;
    }

    WorkflowStatus {
        column-span: 2;
        height: 3;
        padding: 1;
        background: $boost;
    }

    MemoryChart {
        column-span: 2;
    }

    WorkflowGraph {
        row-span: 2;
    }

    TabbedLogViewer {
        row-span: 2;
    }

    OutputFiles {
        column-span: 2;
        height: auto;
    }
    """

    BINDINGS = [
        ("q", "quit", "Quit"),
        ("s", "screenshot", "Screenshot"),
    ]

    def __init__(self, workflow_steps: list, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.workflow_steps = workflow_steps
        self.status_widget = None
        self.log_viewer = None
        self.graph_widget = None
        self.memory_chart = None
        self.output_files_widget = None
        self.start_time = None
        self.update_timer = None
        self.max_memory_kb = 0
        self.memory_warning_shown = False
        self.memory_critical_shown = False
        self.workflow_start_time = None
        self.workflow_completed = False

    def compose(self) -> ComposeResult:
        """Build the UI layout"""
        yield Header()

        self.status_widget = WorkflowStatus()
        yield self.status_widget

        self.memory_chart = MemoryChart()
        yield self.memory_chart

        self.graph_widget = WorkflowGraph(self.workflow_steps)
        yield self.graph_widget

        # Create log viewer (tabs will be added dynamically)
        self.log_viewer = TabbedLogViewer()
        yield self.log_viewer

        self.output_files_widget = OutputFiles()
        yield self.output_files_widget

        yield Footer()

    def on_mount(self):
        """Start the update timer when app mounts"""
        self.start_time = time.time()
        self.workflow_start_time = time.time()
        self.update_timer = self.set_interval(0.1, self.update_elapsed_time)
        self.log_refresh_timer = self.set_interval(0.2, self.refresh_log_viewer)
        # Force initial refresh to ensure UI is interactive immediately
        self.refresh()

    def update_elapsed_time(self):
        """Update the elapsed time display"""
        try:
            if self.start_time and self.status_widget and not self.workflow_completed:
                elapsed = time.time() - self.start_time
                mins = int(elapsed // 60)
                secs = int(elapsed % 60)
                self.status_widget.elapsed_time = f"{mins}m {secs}s"
        except Exception:
            # Silently ignore errors in timer updates
            pass

    def refresh_log_viewer(self):
        """Refresh all log viewer tabs"""
        try:
            if self.log_viewer:
                self.log_viewer.refresh_all_logs()
        except Exception:
            # Silently ignore errors in log refresh
            pass

    def set_current_step(self, step_name: str, step_id: Optional[str] = None):
        """Update the current step"""
        if self.status_widget:
            self.status_widget.current_step = step_name

        if step_id and self.graph_widget:
            self.graph_widget.mark_step_running(step_id)

    def mark_step_complete(self, step_id: str):
        """Mark a step as completed"""
        if self.graph_widget:
            self.graph_widget.mark_step_complete(step_id)

    def set_log_file(self, log_file: Optional[Path], step_name: str = "Step"):
        """Add a new log file tab"""
        if self.log_viewer and log_file:
            self.call_later(self.log_viewer.add_log_file, log_file, step_name)

    def update_files(self, files_dict: dict, step_outputs: dict):
        """Update the output files table and status info"""
        if self.output_files_widget:
            self.output_files_widget.update_files(files_dict, step_outputs)

        # Update workflow graph with step outputs
        if self.graph_widget:
            self.graph_widget.update_step_outputs(step_outputs)

        # Update output files info in status
        if self.status_widget and files_dict:
            file_count = len(files_dict)
            total_size = sum(files_dict.values())
            if total_size > 1024 * 1024:
                size_str = f"{total_size / (1024 * 1024):.1f} MB"
            elif total_size > 1024:
                size_str = f"{total_size / 1024:.1f} KB"
            else:
                size_str = f"{total_size} bytes"
            self.status_widget.output_files_info = f"{file_count} file(s) ({size_str})"

    def set_memory_usage(self, memory_str: str, memory_kb: int):
        """Update memory usage display and check thresholds"""
        if self.status_widget:
            self.status_widget.memory_usage = memory_str

            # Track max memory and update display
            if memory_kb > self.max_memory_kb:
                self.max_memory_kb = memory_kb
                max_memory_gb = memory_kb / (1024 * 1024)

                # Determine style and warnings
                if max_memory_gb >= 5.0:
                    warning_icon = " ⚠️  CRITICAL"
                    if not self.memory_critical_shown:
                        self.memory_critical_shown = True
                elif max_memory_gb >= 3.5:
                    warning_icon = " ⚠️"
                    if not self.memory_warning_shown:
                        self.memory_warning_shown = True
                else:
                    warning_icon = ""

                self.status_widget.max_memory = (
                    f"Max: {max_memory_gb:.1f} GB{warning_icon}"
                )

    def set_status(self, status: str):
        """Update workflow status"""
        if self.status_widget:
            self.status_widget.status = status

    def mark_workflow_complete(self, success: bool = True):
        """Mark workflow as completed to stop timers"""
        self.workflow_completed = True
        # Update final elapsed time
        if self.start_time and self.status_widget:
            elapsed = time.time() - self.start_time
            mins = int(elapsed // 60)
            secs = int(elapsed % 60)
            self.status_widget.elapsed_time = f"{mins}m {secs}s"

        # Show completion notification
        if success:
            self.notify(
                "✅ CWL workflow executed successfully!",
                title="Workflow Complete",
                severity="information",
                timeout=10,
            )
        else:
            self.notify(
                "❌ CWL workflow execution failed!",
                title="Workflow Failed",
                severity="error",
                timeout=10,
            )


async def run_cwl_with_textual_ui(cmd, env, output_dir, cwl_file):
    """
    Run CWL workflow with Textual TUI for monitoring.

    Args:
        cmd: Command list to execute
        env: Environment dictionary to pass to subprocess
        output_dir: CWL output directory where workflow output files are written
        cwl_file: Path to the CWL workflow file for graph visualization

    Returns:
        tuple: (exit_code, list of temp directories used, max_memory_kb)
    """

    # Parse CWL workflow to extract steps
    workflow_steps = []

    def parse_workflow_steps(workflow_data, parent_id=""):
        """Recursively parse workflow steps"""
        steps = []
        if "steps" not in workflow_data:
            return steps

        for step in workflow_data["steps"]:
            step_id = step.get("id", "unknown")
            full_id = f"{parent_id}/{step_id}" if parent_id else step_id

            run_section = step.get("run", {})

            if isinstance(run_section, dict):
                step_label = run_section.get("label", step_id)
                step_class = run_section.get("class", "unknown")

                steps.append(
                    {
                        "id": step_id,
                        "full_id": full_id,
                        "label": step_label,
                        "class": step_class,
                        "level": parent_id.count("/"),
                        "substeps": [],
                    }
                )

                if step_class == "Workflow":
                    substeps = parse_workflow_steps(run_section, full_id)
                    steps[-1]["substeps"] = substeps
            else:
                steps.append(
                    {
                        "id": step_id,
                        "full_id": full_id,
                        "label": step_id,
                        "class": "external",
                        "level": parent_id.count("/"),
                        "substeps": [],
                    }
                )

        return steps

    try:
        with open(cwl_file, "r") as f:
            cwl_data = yaml.safe_load(f)
            workflow_steps = parse_workflow_steps(cwl_data)
    except Exception:
        workflow_steps = [
            {
                "id": "workflow",
                "full_id": "workflow",
                "label": "Workflow",
                "class": "unknown",
                "level": 0,
                "substeps": [],
            }
        ]

    # Create the Textual app
    app = CWLMonitorApp(workflow_steps)

    # State tracking
    current_step_name = "Initializing..."
    current_log_file = None
    current_tmpdir = None
    current_output_prefix = None
    expecting_output_prefix = False
    temp_directories = []
    output_files = {}
    step_outputs = {}
    max_memory_usage = 0
    last_file_scan = 0
    file_scan_interval = 0.5
    process = None
    exit_code = 0
    prmon_data_points_read = {}  # Track lines read per prmon file to avoid duplicates

    # Patterns for log parsing
    step_pattern = re.compile(r"\[(job|step|workflow) ([^\]]+)\]")
    job_pattern = re.compile(r"\[job\s+([^\]]+)\]\s+(/tmp/[^\s$]+)\$")
    workflow_step_pattern = re.compile(
        r"\[(workflow|step)\s+([^\]]+)\]\s+(starting step|start)"
    )
    # tmpdir_pattern = re.compile(r"tmp[a-z0-9_-]{6,}", re.IGNORECASE)
    # output_prefix_pattern = re.compile(r"--output-prefix")
    # prmon_pattern = re.compile(r"prmon.*\.txt")

    def find_current_log_file(line):
        """Extract tmpdir and output_prefix from CWL executor output to locate log file"""
        nonlocal current_log_file, current_step_name, current_tmpdir
        nonlocal current_output_prefix, expecting_output_prefix, temp_directories

        # Pattern 1a: Look for workflow/step starting messages
        workflow_step_match = workflow_step_pattern.search(line)
        if workflow_step_match:
            step = workflow_step_match.group(2)
            current_step_name = step

        # Pattern 1b: Look for cwltool job execution messages
        job_match = job_pattern.search(line)
        if job_match:
            step = job_match.group(1)
            tmpdir = Path(job_match.group(2))

            # Update tracking variables
            current_step_name = step
            current_tmpdir = tmpdir
            # Track this temp directory for later log file discovery
            if tmpdir not in temp_directories:
                temp_directories.append(tmpdir)

        # Pattern 2: Look for --output-prefix argument
        if "--output-prefix" in line:
            # Next non-empty, non-backslash line will have the prefix value
            expecting_output_prefix = True
        elif expecting_output_prefix:
            # This line should contain the actual output prefix value
            stripped = line.strip()
            # Skip lines that are just backslashes or empty
            if stripped and stripped != "\\":
                current_output_prefix = stripped
                expecting_output_prefix = False

        # Pattern 3: If we have both tmpdir and output_prefix, try to find the log file
        if current_tmpdir and current_output_prefix and current_tmpdir.exists():
            # The log file pattern is: {tmpdir}/*_{output_prefix}.log
            log_pattern = f"*_{current_output_prefix}.log"
            log_files = list(current_tmpdir.glob(log_pattern))

            if log_files:
                # Use the first (should be only one) matching log file
                if log_files[0] != current_log_file:
                    current_log_file = log_files[0]
                    app.set_log_file(current_log_file, current_step_name)
                    return True

        return False

    async def monitor_process():
        """Monitor the CWL process and update UI"""
        nonlocal process, exit_code, current_step_name, current_log_file, current_tmpdir
        nonlocal temp_directories, output_files, step_outputs, max_memory_usage, last_file_scan
        nonlocal current_output_prefix, expecting_output_prefix, prmon_data_points_read

        # Wait for app to be ready (widgets mounted)
        while not app.log_viewer:
            await asyncio.sleep(0.1)

        # Create executor log file in output directory
        executor_log_path = output_dir / "executor.log"
        executor_log_file = open(executor_log_path, "w", buffering=1)  # Line buffered

        # Add executor log to the UI (now that app is ready)
        app.set_log_file(executor_log_path, "CWL Executor")

        # Start the subprocess
        process = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
            env=env,
        )

        app.set_status("Running")

        # Read stdout line by line
        async for line_bytes in process.stdout:
            line = line_bytes.decode("utf-8", errors="replace").rstrip()

            # Write to executor log file
            executor_log_file.write(line + "\n")
            executor_log_file.flush()

            # Try to find log file references in the output
            find_current_log_file(line)

            # Parse for step information to update the graph
            step_match = step_pattern.search(line)
            if step_match:
                step_type = step_match.group(1)
                step_id = step_match.group(2)
                app.set_current_step(f"[{step_type} {step_id}]", step_id)

            # Periodic file scanning
            current_time = time.time()
            if current_time - last_file_scan >= file_scan_interval:
                last_file_scan = current_time

                # Scan for output files
                new_files = {}
                current_step_files = []

                # Scan final output directory
                for ext in ["*.root", "*.dst", "*.xdigi", "*.xml", "*.json", "*.txt"]:
                    for f in output_dir.glob(f"**/{ext}"):
                        if f.is_file():
                            size = f.stat().st_size
                            rel_path = str(f.relative_to(output_dir))
                            new_files[rel_path] = size

                            # Track new files for current step
                            if (
                                rel_path not in output_files
                                and current_step_name != "Initializing..."
                            ):
                                current_step_files.append((rel_path, size))

                # Also scan tmpdir
                if current_tmpdir and current_tmpdir.exists():
                    for ext in [
                        "*.root",
                        "*.dst",
                        "*.xdigi",
                        "*.xml",
                        "*.json",
                        "*.txt",
                    ]:
                        for f in current_tmpdir.glob(ext):
                            if f.is_file():
                                size = f.stat().st_size
                                file_key = f.name  # Use just the filename as key
                                new_files[file_key] = size

                                # Track new files for the current step
                                if (
                                    file_key not in output_files
                                    and current_step_name != "Initializing..."
                                ):
                                    current_step_files.append((file_key, size))

                # Add new files to the current step's output list
                if current_step_files:
                    clean_step_name = (
                        current_step_name.replace("[job ", "")
                        .replace("[step ", "")
                        .replace("[workflow ", "")
                        .replace("]", "")
                        .strip()
                    )
                    if clean_step_name not in step_outputs:
                        step_outputs[clean_step_name] = []
                    step_outputs[clean_step_name].extend(current_step_files)

                # Update sizes of existing files in step_outputs
                for _, files_list in step_outputs.items():
                    for i, (file_path, _) in enumerate(files_list):
                        if file_path in new_files:
                            # Update the size tuple in place
                            files_list[i] = (file_path, new_files[file_path])

                output_files = new_files
                app.update_files(output_files, step_outputs)

                # Check memory usage from prmon files
                for tmpdir in temp_directories:
                    if tmpdir.exists():
                        prmon_files = list(tmpdir.glob("prmon*.txt"))
                        for prmon_file in prmon_files:
                            try:
                                content = prmon_file.read_text()
                                # Parse all lines for time-series data (Time, RSS, PSS, Swap, ...)
                                lines = [
                                    line
                                    for line in content.splitlines()
                                    if line and not line.startswith("#")
                                ]

                                if not lines:
                                    continue

                                # Track which lines we've already processed for this file
                                prmon_key = str(prmon_file)
                                lines_already_read = prmon_data_points_read.get(
                                    prmon_key, 0
                                )

                                # Process only new lines to avoid duplicates
                                current_max_mem = 0
                                for i, line in enumerate(lines):
                                    # Skip lines we've already processed
                                    if i < lines_already_read:
                                        continue

                                    parts = line.split()
                                    if (
                                        len(parts) >= 4
                                    ):  # Time,1 WTime, 2 RSS, 3 PSS, 4 Swap, ...
                                        try:
                                            timestamp_seconds = float(
                                                parts[1]
                                            )  # Time in seconds from start
                                            rss = float(parts[3])  # KB
                                            swap = (
                                                float(parts[4]) if len(parts) > 3 else 0
                                            )  # KB
                                            total_mem = rss + swap

                                            # Add data point to chart - timestamp_seconds is already relative
                                            if app.memory_chart:
                                                app.memory_chart.add_data_point(
                                                    timestamp_seconds,
                                                    int(rss),
                                                    int(swap),
                                                )

                                            # Track max for this scan
                                            if total_mem > current_max_mem:
                                                current_max_mem = total_mem
                                        except (ValueError, IndexError):
                                            # Skip malformed lines
                                            continue

                                # Update how many lines we've read for this file
                                prmon_data_points_read[prmon_key] = len(lines)

                                # Update status widget with latest max
                                if current_max_mem > max_memory_usage:
                                    max_memory_usage = current_max_mem
                                    mem_gb = current_max_mem / (1024 * 1024)
                                    mem_str = f"RSS: {mem_gb:.2f} GB"
                                    app.set_memory_usage(mem_str, int(current_max_mem))
                            except Exception:
                                # Silently ignore prmon read errors
                                pass

        # Wait for process to complete
        exit_code = await process.wait()

        # Close executor log file
        executor_log_file.close()

        if exit_code == 0:
            app.set_status("✅ Completed")
            # Mark workflow as complete with success
            app.mark_workflow_complete(success=True)
        else:
            app.set_status(f"❌ Failed (exit {exit_code})")
            # Mark workflow as complete with failure
            app.mark_workflow_complete(success=False)

    # Run the app and process monitoring concurrently
    async def run_app():
        await app.run_async()

    # Start both tasks
    monitor_task = asyncio.create_task(monitor_process())
    app_task = asyncio.create_task(run_app())

    # Wait for both to complete, catching any exceptions
    try:
        await asyncio.gather(monitor_task, app_task, return_exceptions=False)
    except Exception as e:
        # If an exception occurs, ensure we still return properly
        import sys

        print(f"Error in workflow execution: {e}", file=sys.stderr)

    return exit_code, temp_directories, int(max_memory_usage)
