"""
Filename: updateevents.py
Description: Methods to update snirf events
Note: Compliant with pylance strict type checking

Author: Tyler de Zeeuw
License: GPL-3.0
"""

# Built-in imports
import os
import json
import concurrent.futures
from enum import Enum, auto
from typing import Any, List, Optional, Sequence, Union, cast

# External library imports
from PySide6.QtWidgets import QWidget, QVBoxLayout, QLabel, QLineEdit, QPushButton, QComboBox, QHBoxLayout, QMessageBox, QFileDialog
from PySide6.QtCore import Qt

from mne import Annotations
from mne.io import read_raw_snirf  #type: ignore
from mne_nirs.io import write_raw_snirf  #type: ignore

from mne.io.base import BaseRaw
from src.shared.shareddata import APP_NAME


def _load_annotations_worker(file_path: str) -> Annotations:
    """Isolated worker process to load SNIRF annotations without leaving HDF5 memory leaks in main process."""
    raw: BaseRaw = read_raw_snirf(file_path, preload=False, verbose=False)
    return cast(Annotations, getattr(raw, "annotations"))


def load_snirf_annotations_isolated(file_path: str) -> Annotations:
    """Spawns a short-lived process to extract annotations and cleanly releases C-memory allocations."""
    with concurrent.futures.ProcessPoolExecutor(max_workers=1) as executor:
        future = executor.submit(_load_annotations_worker, file_path)
        return future.result()


def _write_snirf_worker(
    file_path: str,
    save_path: str,
    onsets: list[float],
    durations: list[float],
    descriptions: list[str]
) -> None:
    """Isolated process worker to read, annotate, write, and immediately free memory/file handles."""
    raw = read_raw_snirf(file_path, preload=True, verbose=False)
    new_annotations = Annotations(onset=onsets, duration=durations, description=descriptions)
    raw.set_annotations(new_annotations)  #type: ignore
    write_raw_snirf(raw, save_path)


def write_snirf_with_annotations_isolated(
    file_path: str,
    save_path: str,
    onsets: list[float],
    durations: list[float],
    descriptions: list[str]
) -> None:
    """Executes SNIRF writing in a dedicated worker process to ensure 100% memory/handle cleanup."""
    with concurrent.futures.ProcessPoolExecutor(max_workers=1) as executor:
        future = executor.submit(
            _write_snirf_worker,
            file_path,
            save_path,
            onsets,
            durations,
            descriptions
        )
        future.result()


def _align_boris_and_write_snirf_worker(
    file_path: str,
    save_path: str,
    boris_events: list[Any],
    time_shift: float
) -> int:
    """
    Isolated process worker to read SNIRF, align BORIS events, write updated file,
    and cleanly exit releasing all HDF5 C-memory allocations and file locks.
    """
    raw = read_raw_snirf(file_path, preload=True, verbose=False)

    # Type-safe extraction of sampling frequency from raw.info
    info_dict = cast(dict[str, Any], raw.info)
    sfreq = float(info_dict.get("sfreq", 10.0))
    min_shift = 1.0 / sfreq if sfreq > 0 else 0.1
    max_attempts = 10

    onsets: list[float] = []
    durations: list[float] = []
    descriptions: list[str] = []

    open_events: dict[str, list[float]] = {}
    label_counts: dict[str, int] = {}
    used_times: set[float] = set()

    for raw_event in boris_events:
        if not isinstance(raw_event, (list, tuple)):
            continue
        
        event = cast(Sequence[Any], raw_event)

        if len(event) < 3:
            continue

        event_time = float(event[0])
        label = str(event[2])

        count = label_counts.get(label, 0) + 1
        label_counts[label] = count

        if label not in open_events:
            open_events[label] = []

        if count % 2 == 1:
            open_events[label].append(event_time)
        else:
            if open_events[label]:
                start_time = open_events[label].pop(0)
                duration = event_time - start_time
                if duration <= 0:
                    continue

                adjusted_time = start_time + time_shift
                attempts = 0
                while round(adjusted_time, 6) in used_times and attempts < max_attempts:
                    adjusted_time += min_shift
                    attempts += 1
                if attempts == max_attempts:
                    continue

                adjusted_time = round(adjusted_time, 6)
                used_times.add(adjusted_time)
                onsets.append(adjusted_time)
                durations.append(round(duration, 6))
                descriptions.append(label)

    # Handle unmatched start markers
    for label, starts in open_events.items():
        for start_time in starts:
            adjusted_time = start_time + time_shift
            attempts = 0
            while round(adjusted_time, 6) in used_times and attempts < max_attempts:
                adjusted_time += min_shift
                attempts += 1
            if attempts == max_attempts:
                continue
            adjusted_time = round(adjusted_time, 6)
            used_times.add(adjusted_time)
            onsets.append(adjusted_time)
            durations.append(0.0)
            descriptions.append(label)

    new_annotations = Annotations(onset=onsets, duration=durations, description=descriptions)
    raw.set_annotations(new_annotations)  #type: ignore
    write_raw_snirf(raw, save_path)

    return len(onsets)


def align_boris_and_write_snirf_isolated(
    file_path: str,
    save_path: str,
    boris_events: list[Any],
    time_shift: float
) -> int:
    """Executes BORIS alignment and SNIRF writing in a dedicated worker process."""
    with concurrent.futures.ProcessPoolExecutor(max_workers=1) as executor:
        future = executor.submit(
            _align_boris_and_write_snirf_worker,
            file_path,
            save_path,
            boris_events,
            time_shift
        )
        return future.result()
    

class EventUpdateMode(Enum):
    WRITE_SNIRF = auto()   # destructive
    WRITE_JSON = auto()    # non-destructive


class UpdateEventsWindow(QWidget):
    def __init__(self, parent: Optional[QWidget]=None, mode: EventUpdateMode=EventUpdateMode.WRITE_SNIRF, caller: Optional[Union[str, object]] = None,):
        super().__init__(parent, Qt.WindowType.Window)

        self.mode = mode
        self.caller = caller or self.__class__.__name__
        self.setWindowTitle(f"Update event markers - {APP_NAME.upper()}")
        self.resize(760, 200)

        print("INIT MODE:", mode)

        self.label_file_a = QLabel("SNIRF file:")
        self.line_edit_file_a = QLineEdit()
        self.line_edit_file_a.setReadOnly(True)
        self.btn_browse_a = QPushButton("Browse .snirf")
        self.btn_browse_a.clicked.connect(self.browse_file_a)

        self.label_file_b = QLabel("BORIS file:")
        self.line_edit_file_b = QLineEdit()
        self.line_edit_file_b.setReadOnly(True)
        self.btn_browse_b = QPushButton("Browse .boris")
        self.btn_browse_b.clicked.connect(self.browse_file_b)

        self.label_suffix = QLabel("Filename in BORIS project file:")
        self.combo_suffix = QComboBox()
        self.combo_suffix.setEditable(False)
        self.combo_suffix.currentIndexChanged.connect(self.on_observation_selected)

        self.label_events = QLabel("Events in selected observation:")
        self.combo_events = QComboBox()
        self.combo_events.setEnabled(False)
        
        self.label_snirf_events = QLabel("Events in SNIRF file:")
        self.combo_snirf_events = QComboBox()
        self.combo_snirf_events.setEnabled(False)

        self.btn_clear = QPushButton("Clear")
        self.btn_go = QPushButton("Go")
        self.btn_clear.clicked.connect(self.clear_files)
        self.btn_go.clicked.connect(self.go_action)

        # --- 
        layout = QVBoxLayout()
        self.description = QLabel()
        self.description.setTextFormat(Qt.TextFormat.RichText)
        self.description.setTextInteractionFlags(Qt.TextInteractionFlag.TextBrowserInteraction)
        self.description.setOpenExternalLinks(True)
            
        self.description.setText("The events that are present in a snirf file may not be the events that are to be studied and examined.<br>"
        "Utilizing different software and video recordings, it is easy enough to see when an action actually occured in a file.<br>"
        "The software <a href='https://www.boris.unito.it/'>BORIS</a> is used to add these events to video files, and these events can be applied to the snirf file <br>"
        "selected below by selecting the correct BORIS observation and time syncing it to an event that it shares with the snirf file.")

        layout.addWidget(self.description)

        help_text_a = "Select the SNIRF (.snirf) file to update with new event markers."

        file_a_layout = QHBoxLayout()

        # Help button on the left
        help_btn_a = QPushButton("?")
        help_btn_a.setFixedWidth(25)
        help_btn_a.setToolTip(help_text_a)
        help_btn_a.clicked.connect(lambda: self.show_help_popup(help_text_a))
        file_a_layout.addWidget(help_btn_a)

        # Container for label + line_edit + browse button with tooltip
        file_a_container = QWidget()
        file_a_container_layout = QHBoxLayout()
        file_a_container_layout.setContentsMargins(0, 0, 0, 0)
        file_a_container_layout.addWidget(self.label_file_a)
        file_a_container_layout.addWidget(self.line_edit_file_a)
        file_a_container_layout.addWidget(self.btn_browse_a)
        file_a_container.setLayout(file_a_container_layout)
        file_a_container.setToolTip(help_text_a)

        file_a_layout.addWidget(file_a_container)
        layout.addLayout(file_a_layout)

        help_text_b = "Provide a .boris project file that contains events for this participant."

        file_b_layout = QHBoxLayout()

        help_btn_b = QPushButton("?")
        help_btn_b.setFixedWidth(25)
        help_btn_b.setToolTip(help_text_b)
        help_btn_b.clicked.connect(lambda: self.show_help_popup(help_text_b))
        file_b_layout.addWidget(help_btn_b)

        file_b_container = QWidget()
        file_b_container_layout = QHBoxLayout()
        file_b_container_layout.setContentsMargins(0, 0, 0, 0)
        file_b_container_layout.addWidget(self.label_file_b)
        file_b_container_layout.addWidget(self.line_edit_file_b)
        file_b_container_layout.addWidget(self.btn_browse_b)
        file_b_container.setLayout(file_b_container_layout)
        file_b_container.setToolTip(help_text_b)

        file_b_layout.addWidget(file_b_container)
        layout.addLayout(file_b_layout)

        help_text_suffix = "This participant from the .boris project file matches the .snirf file."

        suffix_layout = QHBoxLayout()

        help_btn_suffix = QPushButton("?")
        help_btn_suffix.setFixedWidth(25)
        help_btn_suffix.setToolTip(help_text_suffix)
        help_btn_suffix.clicked.connect(lambda: self.show_help_popup(help_text_suffix))
        suffix_layout.addWidget(help_btn_suffix)

        suffix_container = QWidget()
        suffix_container_layout = QHBoxLayout()
        suffix_container_layout.setContentsMargins(0, 0, 0, 0)
        suffix_container_layout.addWidget(self.label_suffix)
        suffix_container_layout.addWidget(self.combo_suffix)
        suffix_container.setLayout(suffix_container_layout)
        suffix_container.setToolTip(help_text_suffix)

        suffix_layout.addWidget(suffix_container)
        layout.addLayout(suffix_layout)

        help_text_suffix = "The events extracted from the BORIS project file for the selected observation."

        suffix2_layout = QHBoxLayout()

        help_btn_suffix = QPushButton("?")
        help_btn_suffix.setFixedWidth(25)
        help_btn_suffix.setToolTip(help_text_suffix)
        help_btn_suffix.clicked.connect(lambda: self.show_help_popup(help_text_suffix))
        suffix2_layout.addWidget(help_btn_suffix)
        
        suffix2_container = QWidget()
        suffix2_container_layout = QHBoxLayout()
        suffix2_container_layout.setContentsMargins(0, 0, 0, 0)
        suffix2_container_layout.addWidget(self.label_events)
        suffix2_container_layout.addWidget(self.combo_events)
        suffix2_container.setLayout(suffix2_container_layout)
        suffix2_container.setToolTip(help_text_suffix)

        suffix2_layout.addWidget(suffix2_container)
        layout.addLayout(suffix2_layout)
        
        snirf_events_layout = QHBoxLayout()

        help_text_snirf_events = "The event markers extracted from the SNIRF file."
        help_btn_snirf_events = QPushButton("?")
        help_btn_snirf_events.setFixedWidth(25)
        help_btn_snirf_events.setToolTip(help_text_snirf_events)
        help_btn_snirf_events.clicked.connect(lambda: self.show_help_popup(help_text_snirf_events))
        snirf_events_layout.addWidget(help_btn_snirf_events)

        snirf_events_container = QWidget()
        snirf_events_container_layout = QHBoxLayout()
        snirf_events_container_layout.setContentsMargins(0, 0, 0, 0)
        snirf_events_container_layout.addWidget(self.label_snirf_events)
        snirf_events_container_layout.addWidget(self.combo_snirf_events)
        snirf_events_container.setLayout(snirf_events_container_layout)
        snirf_events_container.setToolTip(help_text_snirf_events)

        snirf_events_layout.addWidget(snirf_events_container)
        layout.addLayout(snirf_events_layout)

        buttons_layout = QHBoxLayout()
        buttons_layout.addStretch()
        buttons_layout.addWidget(self.btn_clear)
        buttons_layout.addWidget(self.btn_go)
        layout.addLayout(buttons_layout)

        self.setLayout(layout)

    def show_help_popup(self, text: str) -> None:
        msg = QMessageBox(self)
        msg.setWindowTitle(f"Parameter Info - {APP_NAME.upper()}")
        msg.setText(text)
        msg.exec()

    def browse_file_a(self) -> None:
        file_path, _ = QFileDialog.getOpenFileName(self, "Select SNIRF File", "", "SNIRF Files (*.snirf)")
        if file_path:
            self.line_edit_file_a.setText(file_path)
            try:
                # Memory leak safe: Extracted in a separate process that terminates immediately
                annotations: Annotations = load_snirf_annotations_isolated(file_path)
                
                print(f"Loaded {len(annotations)} annotations from {file_path}")

                # Build individual event entries
                event_entries: list[str] = []
                for onset, description in zip(annotations.onset, annotations.description):
                    event_str = f"{description} @ {onset:.3f}s"
                    event_entries.append(event_str)

                if not event_entries:
                    QMessageBox.information(self, "No Events", "No events found in SNIRF file.")
                    self.combo_snirf_events.clear()
                    self.combo_snirf_events.setEnabled(False)
                    return

                self.combo_snirf_events.clear()
                self.combo_snirf_events.addItems(event_entries)
                self.combo_snirf_events.setEnabled(True)

            except Exception as e:
                QMessageBox.warning(self, "Error", f"Could not read SNIRF file with MNE:\n{str(e)}")
                self.combo_snirf_events.clear()
                self.combo_snirf_events.setEnabled(False)

    def browse_file_b(self) -> None:
        file_path, _ = QFileDialog.getOpenFileName(self, "Select BORIS File", "", "BORIS project Files (*.boris)")
        if file_path:
            self.line_edit_file_b.setText(file_path)
            
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                self.boris_data = data

                observation_keys = self.extract_boris_observation_keys(data)
                self.combo_suffix.clear()
                self.combo_suffix.addItems(observation_keys)

            except (json.JSONDecodeError, FileNotFoundError, KeyError) as e:
                QMessageBox.warning(self, "Error", f"Failed to parse BORIS file:\n{e}")

                
    def extract_boris_observation_keys(self, data: dict[str, Any]) -> List[str]:
        if "observations" not in data:
            raise KeyError("Missing 'observations' key in BORIS file.")

        observations = cast(dict[str, Any], data["observations"])
        return list(observations.keys())

    
    def on_observation_selected(self):
        selected_obs = self.combo_suffix.currentText()
        if not selected_obs or not hasattr(self, 'boris_data'):
            self.combo_events.clear()
            self.combo_events.setEnabled(False)
            return

        try:
            events = self.boris_data["observations"][selected_obs]["events"]
        except (KeyError, TypeError):
            self.combo_events.clear()
            self.combo_events.setEnabled(False)
            return

        event_entries: list[str] = []
        for raw_event in events:
            if not isinstance(raw_event, (list, tuple)):
                continue

            event = cast(Sequence[Any], raw_event)

            if len(event) >= 3:
                timestamp = float(event[0])
                label = str(event[2])
                display = f"{label} @ {timestamp:.3f}"
                event_entries.append(display)

        self.combo_events.clear()
        self.combo_events.addItems(event_entries)
        self.combo_events.setEnabled(bool(event_entries))

    def clear_files(self) -> None:
        self.line_edit_file_a.clear()
        self.line_edit_file_b.clear()
    
    def go_action(self) -> None:
        
        file_a = self.line_edit_file_a.text()
        suffix = "flare" 
    
        if not hasattr(self, "boris_data") or self.combo_events.count() == 0 or self.combo_snirf_events.count() == 0:
            QMessageBox.warning(self, "Missing data", "Please make sure a BORIS and SNIRF event are selected.")
            return

        # Extract BORIS anchor
        try:
            boris_label, boris_time_str = self.combo_events.currentText().split(" @ ")
            boris_anchor_time = float(boris_time_str.replace("s", "").strip())
        except Exception as e:
            QMessageBox.critical(self, "BORIS Event Error", f"Could not parse BORIS anchor event:\n{e}")
            return

        # Extract SNIRF anchor
        try:
            snirf_label, snirf_time_str = self.combo_snirf_events.currentText().split(" @ ")
            snirf_anchor_time = float(snirf_time_str.replace("s", "").strip())
        except Exception as e:
            QMessageBox.critical(self, "SNIRF Event Error", f"Could not parse SNIRF anchor event:\n{e}")
            return

        time_shift = snirf_anchor_time - boris_anchor_time

        selected_obs = self.combo_suffix.currentText()
        if not selected_obs or selected_obs not in self.boris_data["observations"]:
            QMessageBox.warning(self, "Invalid selection", "Selected observation not found in BORIS file.")
            return

        boris_events = self.boris_data["observations"][selected_obs].get("events", [])
        if not boris_events:
            QMessageBox.warning(self, "No BORIS events", "No events found in selected BORIS observation.")
            return

        snirf_path = self.line_edit_file_a.text()
        if not snirf_path:
            QMessageBox.warning(self, "No SNIRF file", "Please select a SNIRF file.")
            return

        boris_obs = self.boris_data["observations"][selected_obs]

        # --- Extract videos + delays ---
        files = boris_obs.get("file", {})
        offsets = boris_obs.get("media_info", {}).get("offset", {})

        videos = {}
        for key, path in files.items():
            if path:  # only include videos that exist
                delay = offsets.get(key, 0.0)  # default 0 if missing
                videos[key] = {"file": path, "delay": delay}

        base_name = os.path.splitext(os.path.basename(file_a))[0]

        if self.mode == EventUpdateMode.WRITE_SNIRF:
            # Open save dialog for SNIRF
            base_name = os.path.splitext(os.path.basename(file_a))[0]
            suggested_name = f"{base_name}_{suffix}.snirf"
            save_path, _ = QFileDialog.getSaveFileName(
                self,
                "Save SNIRF File As",
                suggested_name,
                "SNIRF Files (*.snirf)"
            )
            if not save_path:
                print("SNIRF save cancelled.")
                return
            if not save_path.lower().endswith(".snirf"):
                save_path += ".snirf"

            try:
                # Memory leak safe: Worker handles alignment calculation, read, write, and process termination
                event_count = align_boris_and_write_snirf_isolated(
                    file_a,
                    save_path,
                    boris_events,
                    time_shift
                )

                QMessageBox.information(self, "Success", f"SNIRF file updated with {event_count} aligned BORIS events.")

            except Exception as e:
                QMessageBox.critical(self, "Error", f"Failed to update SNIRF file:\n{e}")

        elif self.mode == EventUpdateMode.WRITE_JSON:
            # Open save dialog for JSON
            base_name = os.path.splitext(os.path.basename(file_a))[0]
            suggested_name = f"{base_name}_{suffix}_alignment.json"
            save_path, _ = QFileDialog.getSaveFileName(
                self,
                "Save Event Alignment JSON As",
                suggested_name,
                "JSON Files (*.json)"
            )
            if not save_path:
                print("JSON save cancelled.")
                return
            if not save_path.lower().endswith(".json"):
                save_path += ".json"

            # Build JSON dict
            json_data: dict[str, Any] = {
                "observation": selected_obs,
                "snirf_anchor": {"label": snirf_label, "time": snirf_anchor_time},
                "boris_anchor": {"label": boris_label, "time": boris_anchor_time},
                "time_shift": time_shift,
                "videos": videos
            }

            # Write JSON
            try:
                with open(save_path, "w", encoding="utf-8") as f:
                    json.dump(json_data, f, indent=4)
                QMessageBox.information(self, "Success", f"Event alignment saved to:\n{save_path}")
            except Exception as e:
                QMessageBox.critical(self, "Error", f"Failed to write JSON:\n{e}")


class UpdateEventsBlazesWindow(QWidget):

    def __init__(self, parent: Optional[QWidget]=None, mode: EventUpdateMode=EventUpdateMode.WRITE_SNIRF, caller: Optional[Union[str, object]] = None,):
        super().__init__(parent, Qt.WindowType.Window)

        self.mode = mode
        self.caller = caller or self.__class__.__name__
        self.setWindowTitle("Update event markers (BLAZES)")
        self.resize(760, 200)

        self.label_file_a = QLabel("SNIRF file:")
        self.line_edit_file_a = QLineEdit()
        self.line_edit_file_a.setReadOnly(True)
        self.btn_browse_a = QPushButton("Browse .snirf")
        self.btn_browse_a.clicked.connect(self.browse_file_a)

        self.label_file_b = QLabel("BLAZES file:")
        self.line_edit_file_b = QLineEdit()
        self.line_edit_file_b.setReadOnly(True)
        self.btn_browse_b = QPushButton("Browse .blaze")
        self.btn_browse_b.clicked.connect(self.browse_file_b)

        self.label_events = QLabel("Events in selected blazes file:")
        self.combo_events = QComboBox()
        self.combo_events.setEnabled(False)
        
        self.label_snirf_events = QLabel("Events in SNIRF file:")
        self.combo_snirf_events = QComboBox()
        self.combo_snirf_events.setEnabled(False)

        self.btn_clear = QPushButton("Clear")
        self.btn_go = QPushButton("Go")
        self.btn_clear.clicked.connect(self.clear_files)
        self.btn_go.clicked.connect(self.go_action)

        # --- 
        layout = QVBoxLayout()
        self.description = QLabel()
        self.description.setTextFormat(Qt.TextFormat.RichText)
        self.description.setTextInteractionFlags(Qt.TextInteractionFlag.TextBrowserInteraction)
        self.description.setOpenExternalLinks(True)
            
        self.description.setText("The events that are present in a snirf file may not be the events that are to be studied and examined.<br>"
        "Utilizing different software and video recordings, it is easy enough to see when an action actually occured in a file.<br>"
        "The software <a href='https://git.research.dezeeuw.ca/tyler/blazes/'>BLAZES</a> is used to create these events in video files, and these events can be applied to the snirf file <br>"
        "selected below by time syncing it to an event that it shares with the snirf file.")

        layout.addWidget(self.description)

        help_text_a = "Select the SNIRF (.snirf) file to update with new event markers."

        file_a_layout = QHBoxLayout()

        # Help button on the left
        help_btn_a = QPushButton("?")
        help_btn_a.setFixedWidth(25)
        help_btn_a.setToolTip(help_text_a)
        help_btn_a.clicked.connect(lambda: self.show_help_popup(help_text_a))
        file_a_layout.addWidget(help_btn_a)

        # Container for label + line_edit + browse button with tooltip
        file_a_container = QWidget()
        file_a_container_layout = QHBoxLayout()
        file_a_container_layout.setContentsMargins(0, 0, 0, 0)
        file_a_container_layout.addWidget(self.label_file_a)
        file_a_container_layout.addWidget(self.line_edit_file_a)
        file_a_container_layout.addWidget(self.btn_browse_a)
        file_a_container.setLayout(file_a_container_layout)
        file_a_container.setToolTip(help_text_a)

        file_a_layout.addWidget(file_a_container)
        layout.addLayout(file_a_layout)

        help_text_b = "Provide a .blaze output file that contains events for this participant."

        file_b_layout = QHBoxLayout()

        help_btn_b = QPushButton("?")
        help_btn_b.setFixedWidth(25)
        help_btn_b.setToolTip(help_text_b)
        help_btn_b.clicked.connect(lambda: self.show_help_popup(help_text_b))
        file_b_layout.addWidget(help_btn_b)

        file_b_container = QWidget()
        file_b_container_layout = QHBoxLayout()
        file_b_container_layout.setContentsMargins(0, 0, 0, 0)
        file_b_container_layout.addWidget(self.label_file_b)
        file_b_container_layout.addWidget(self.line_edit_file_b)
        file_b_container_layout.addWidget(self.btn_browse_b)
        file_b_container.setLayout(file_b_container_layout)
        file_b_container.setToolTip(help_text_b)

        file_b_layout.addWidget(file_b_container)
        layout.addLayout(file_b_layout)

        help_text_suffix = "The events extracted from the blaze file."

        suffix2_layout = QHBoxLayout()

        help_btn_suffix = QPushButton("?")
        help_btn_suffix.setFixedWidth(25)
        help_btn_suffix.setToolTip(help_text_suffix)
        help_btn_suffix.clicked.connect(lambda: self.show_help_popup(help_text_suffix))
        suffix2_layout.addWidget(help_btn_suffix)
        
        suffix2_container = QWidget()
        suffix2_container_layout = QHBoxLayout()
        suffix2_container_layout.setContentsMargins(0, 0, 0, 0)
        suffix2_container_layout.addWidget(self.label_events)
        suffix2_container_layout.addWidget(self.combo_events)
        suffix2_container.setLayout(suffix2_container_layout)
        suffix2_container.setToolTip(help_text_suffix)

        suffix2_layout.addWidget(suffix2_container)
        layout.addLayout(suffix2_layout)
        
        snirf_events_layout = QHBoxLayout()

        help_text_snirf_events = "The event markers extracted from the SNIRF file."
        help_btn_snirf_events = QPushButton("?")
        help_btn_snirf_events.setFixedWidth(25)
        help_btn_snirf_events.setToolTip(help_text_snirf_events)
        help_btn_snirf_events.clicked.connect(lambda: self.show_help_popup(help_text_snirf_events))
        snirf_events_layout.addWidget(help_btn_snirf_events)

        snirf_events_container = QWidget()
        snirf_events_container_layout = QHBoxLayout()
        snirf_events_container_layout.setContentsMargins(0, 0, 0, 0)
        snirf_events_container_layout.addWidget(self.label_snirf_events)
        snirf_events_container_layout.addWidget(self.combo_snirf_events)
        snirf_events_container.setLayout(snirf_events_container_layout)
        snirf_events_container.setToolTip(help_text_snirf_events)

        snirf_events_layout.addWidget(snirf_events_container)
        layout.addLayout(snirf_events_layout)

        buttons_layout = QHBoxLayout()
        buttons_layout.addStretch()
        buttons_layout.addWidget(self.btn_clear)
        buttons_layout.addWidget(self.btn_go)
        layout.addLayout(buttons_layout)

        self.setLayout(layout)


    def show_help_popup(self, text: str) -> None:
        msg = QMessageBox(self)
        msg.setWindowTitle(f"Parameter Info - {APP_NAME.upper()}")
        msg.setText(text)
        msg.exec()

    def browse_file_a(self) -> None:
        file_path, _ = QFileDialog.getOpenFileName(self, "Select SNIRF File", "", "SNIRF Files (*.snirf)")
        if file_path:
            self.line_edit_file_a.setText(file_path)
            try:
                # Memory leak safe: Extracted in a separate process that terminates immediately
                annotations: Annotations = load_snirf_annotations_isolated(file_path)
                
                print(f"Loaded {len(annotations)} annotations from {file_path}")

                # Build individual event entries
                event_entries: list[str] = []
                for onset, description in zip(annotations.onset, annotations.description):
                    event_str = f"{description} @ {onset:.3f}s"
                    event_entries.append(event_str)

                if not event_entries:
                    QMessageBox.information(self, "No Events", "No events found in SNIRF file.")
                    self.combo_snirf_events.clear()
                    self.combo_snirf_events.setEnabled(False)
                    return

                self.combo_snirf_events.clear()
                self.combo_snirf_events.addItems(event_entries)
                self.combo_snirf_events.setEnabled(True)

            except Exception as e:
                QMessageBox.warning(self, "Error", f"Could not read SNIRF file with MNE:\n{str(e)}")
                self.combo_snirf_events.clear()
                self.combo_snirf_events.setEnabled(False)

    def browse_file_b(self) -> None:
        file_path, _ = QFileDialog.getOpenFileName(self, "Select JSON Timeline File", "", "JSON Files (*.json)")
        if file_path:
            self.line_edit_file_b.setText(file_path)
            
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                self.json_data = data

                obs_keys = self.extract_json_observation_strings(data)
                self.combo_events.clear()
                if obs_keys:
                    self.combo_events.addItems(obs_keys)
                    self.combo_events.setEnabled(True)
                else:
                    QMessageBox.information(self, "No Events", "No events found in JSON file.")
                    self.combo_events.setEnabled(False)

            except (json.JSONDecodeError, FileNotFoundError, KeyError, TypeError) as e:
                QMessageBox.warning(self, "Error", f"Failed to parse JSON file:\n{e}")
                self.combo_events.clear()
                self.combo_events.setEnabled(False)
                

    def extract_json_observation_strings(self, data: dict[str, Any]) -> List[str]:
        if "events" not in data:
            raise KeyError("Missing 'events' key in JSON file.")

        event_strings: List[str] = []

        # Flat list chronologically ordered
        for event in data["events"]:
            track_name = event.get("track_name", "Unknown")
            onset = event.get("start_sec", 0.0)
            
            display_str = f"{track_name} @ {onset:.3f}s"
            event_strings.append(display_str)

        return event_strings
    

    def clear_files(self) -> None:
        self.line_edit_file_a.clear()
        self.line_edit_file_b.clear()
    
    
    def go_action(self) -> None:
        file_a = self.line_edit_file_a.text()
        _ = self.line_edit_file_b.text()
        suffix = APP_NAME
    
        if not hasattr(self, "json_data") or self.combo_events.count() == 0 or self.combo_snirf_events.count() == 0:
            QMessageBox.warning(self, "Missing data", "Please make sure a JSON and SNIRF event are selected.")
            return

        try:
            json_text = self.combo_events.currentText()
            _, json_time_str = json_text.split(" @ ")
            json_anchor_time = float(json_time_str.replace("s", "").strip())
        except Exception as e:
            QMessageBox.critical(self, "JSON Event Error", f"Could not parse JSON anchor:\n{e}")
            return

        try:
            snirf_text = self.combo_snirf_events.currentText()
            _, snirf_time_str = snirf_text.split(" @ ")
            snirf_anchor_time = float(snirf_time_str.replace("s", "").strip())
        except Exception as e:
            QMessageBox.critical(self, "SNIRF Event Error", f"Could not parse SNIRF anchor:\n{e}")
            return

        time_shift = snirf_anchor_time - json_anchor_time

        onsets: list[float] = []
        durations: list[float] = []
        descriptions: list[str] = []
        skipped_count = 0
        
        try:
            events_list = self.json_data.get("events", [])
            
            for event in events_list:
                track_name = event.get("track_name", "Unknown")
                clean_name = track_name.replace("AI: ", "").strip()
                
                original_start = event.get("start_sec", 0.0)
                original_end = event.get("end_sec", original_start)
                duration = original_end - original_start

                # Minimum 0.1s duration
                if duration < 0.1:
                    skipped_count += 1
                    continue

                # Apply shift
                adjusted_onset = original_start + time_shift

                onsets.append(round(adjusted_onset, 6))
                durations.append(round(duration, 6))
                descriptions.append(clean_name) 

        except Exception as e:
            QMessageBox.critical(self, "Track Error", f"Failed to process tracks: {e}")
            return

        if not onsets:
            QMessageBox.warning(self, "No Data", f"No events met the 0.1s threshold. (Skipped {skipped_count})")
            return

        if self.mode == EventUpdateMode.WRITE_SNIRF:
            suggested_name = f"{os.path.splitext(os.path.basename(file_a))[0]}_{suffix}.snirf"
            save_path, _ = QFileDialog.getSaveFileName(self, "Save SNIRF", suggested_name, "SNIRF Files (*.snirf)")
            
            if not save_path: return
            if not save_path.lower().endswith(".snirf"): save_path += ".snirf"

            try:
                # Memory leak safe: Worker handles read, write, and complete memory reclamation upon exit
                write_snirf_with_annotations_isolated(file_a, save_path, onsets, durations, descriptions)

                QMessageBox.information(self, "Success", 
                    f"Aligned {len(onsets)} events.\n(Filtered out {skipped_count} short events)")
            except Exception as e:
                QMessageBox.critical(self, "Error", f"Failed to update SNIRF file:\n{e}")