from __future__ import annotations

from datetime import datetime
import logging
import os

import pandas as pd
from textual.app import ComposeResult
from textual.binding import Binding
from textual.containers import Horizontal, Vertical
from textual.screen import Screen
from textual.widgets import Button, DataTable, Footer, Header, Static

from TUI.allowlistselectionscreen import AllowlistSelectionScreen
from utils.configmanager import load_env

logger = logging.getLogger(__name__)


def _load_working_dir() -> str:
    """
    Load the working directory from environment variables or use the current working directory.
    """
    wd = os.environ.get("WORKING_DIR")
    if wd:
        return wd
    return os.getcwd()


class OTPActivitiesWidget(Static):
    """
    Reusable widget that contains the sessions table (left) and an Activity Preview (right).
    The right side shows an Activity Preview that takes ~75% vertical space, and a lower area
    with Back and Continue buttons. The Continue button pushes ActivityDetailScreen with the
    currently-loaded activities.
    """

    DEFAULT_CSS = """
    OTPActivitiesWidget {
        height: 1fr;
    }
    #main_row {
        width: 100%;
        height: 100%;
        layout: horizontal;
    }
    #left_panel {
        width: 60%;
        min-width: 60;
        border: none;
    }
    #right_panel {
        width: 40%;
        min-width: 40;
        border: none;
        layout: vertical;
    }
    #activity_preview_container {
        height: 1fr;
        border: none;
        padding: 1 1;
    }
    #activity_buttons {
        height: auto;
        min-height: 3;
        padding: 1 1;
        content-align: center middle;
    }
    """

    def compose(self) -> ComposeResult:
        # Layout: horizontal main row with left & right panels
        with Horizontal(id="main_row"):
            # Left: sessions area
            with Vertical(id="left_panel"):
                yield Static("OTP Sessions", classes="panel-title")
                with Vertical(id="sessions_table_container"):
                    self.sessions_table = DataTable(id="sessions_table")
                    self.sessions_table.styles.width = "100%"
                    yield self.sessions_table
            # Right: Activity Preview (top 3/4) + buttons (bottom 1/4)
            with Vertical(id="right_panel"):
                # Activity preview area (takes ~75% of right panel)
                yield Static("Activity Preview", classes="panel-title")
                with Vertical(id="activity_preview_container"):
                    self.activities_table = DataTable(id="activity_preview_table")
                    yield self.activities_table
                # Buttons area at the bottom (Back, Continue)
                with Horizontal(id="activity_buttons"):
                    # Back takes left side, Continue right side
                    self.back_btn = Button("Back", id="activity_back_btn")
                    self.continue_btn = Button("Continue", id="activity_continue_btn")
                    # Stretch buttons nicely
                    self.back_btn.styles.width = "50%"
                    self.continue_btn.styles.width = "50%"
                    yield self.back_btn
                    yield self.continue_btn

    async def on_mount(self) -> None:
        # Configure sessions table and activities preview
        self.sessions_table.clear()
        self.sessions_table.add_columns(
            "otpid", "hostname", "status", "purpose", "granted"
        )
        self.activities_table.clear()
        # activities_table columns are dynamically added when activities are loaded.
        # Selection behavior
        self.sessions_table.cursor_type = "row"
        try:
            self.sessions_table.zebra_stripes = True
        except Exception:
            pass
        self.activities_table.cursor_type = "row"
        try:
            self.activities_table.zebra_stripes = True
        except Exception:
            pass
        # Store state
        self._sessions_df: pd.DataFrame | None = None
        self._activities_df: pd.DataFrame | None = None
        self._selected_session_otpid: str | int | None = None

    async def on_button_pressed(self, event) -> None:  # type: ignore[override]
        """
        Handle Back / Continue buttons for the Activity Preview area.
        """
        # Try to resolve the button object from the event
        btn = (
            getattr(event, "button", None)
            or getattr(event, "sender", None)
            or getattr(event, "control", None)
            or getattr(event, "widget", None)
        )
        btn_id = (
            getattr(btn, "id", None)
            or getattr(event, "button_id", None)
            or getattr(event, "id", None)
        )
        # ---- Back ----
        if btn is self.back_btn or btn_id == getattr(self.back_btn, "id", None):
            while len(self.app.screen_stack) > 2:
                self.app.pop_screen()
            event.stop()
            return
        # ---- Continue ----
        if btn is self.continue_btn or btn_id == getattr(self.continue_btn, "id", None):
            if self._activities_df is None or self._activities_df.empty:
                logger.info("Continue pressed but no activities loaded.")
                await self.post_message(
                    Static("No activities loaded to continue with.")
                )
                return
            # Copy activities DataFrame to pass to new screen
            activities_copy = self._activities_df.copy()
            otpid = self._selected_session_otpid
            # Optionally include hostname if available
            hostname = None
            try:
                if self._sessions_df is not None:
                    df = self._sessions_df.reset_index(drop=True)
                    match = df[df["otpid"] == otpid]
                    if not match.empty:
                        hostname = match.iloc[0].get("hostname")
            except Exception:
                hostname = None
            # Create and push ActivityDetailScreen, handing the data
            try:
                detail_screen = ActivityDetailScreen(
                    activities_copy, otpid=otpid, hostname=hostname
                )
                await self.app.push_screen(detail_screen)
            except Exception as exc:
                logger.exception("Failed to push ActivityDetailScreen: %s", exc)
            return
        # Unknown button on widget
        logger.debug(
            "Unhandled OTPActivitiesWidget button pressed (resolved btn=%r, id=%r)",
            btn,
            btn_id,
        )

    async def on_data_table_row_selected(self, event) -> None:  # type: ignore[override]
        """
        Robust handler for DataTable row-selection across Textual micro-versions.
        Tries many attribute names and shapes:
         - numeric index (row_key, row_index, index)
         - coordinate object or tuple (coordinate.row or (row, col))
         - direct row values (row, values, cells) -> we try to map those back to the sessions DF
         - table.cursor_row fallback
        """
        # 1) Determine the sending table (best-effort)
        sender = None
        for attr in ("sender", "table", "data_table", "control"):
            sender = getattr(event, attr, None)
            if sender is not None:
                break
        if sender is None:
            sender = self.sessions_table  # Assume sessions_table if unknown
        # Only respond to selections in the sessions table
        if sender is not self.sessions_table:
            return

        # Helper to log and return
        def _bad(msg: str, *args):
            logger.warning(msg, *args)
            return None

        # 2) Try to extract a numeric index
        row_key = None
        for attr in ("row_key", "row", "row_index", "index"):
            row_key = getattr(event, attr, None)
            if row_key is not None:
                break
        # If coordinate: try to extract .row or tuple[0]
        if row_key is None:
            coord = getattr(event, "coordinate", None) or getattr(
                event, "cursor_coordinate", None
            )
            if coord is not None:
                if hasattr(coord, "row"):
                    row_key = coord.row
                elif isinstance(coord, (tuple, list)) and len(coord) >= 1:
                    row_key = coord[0]
        # If still nothing, maybe the event provides the row's cell values directly
        row_values = None
        for attr in ("values", "cells", "row", "row_values", "selected_row_values"):
            val = getattr(event, attr, None)
            if val:
                # Prefer actual sequence of cell values
                row_values = val
                break
        # If we have row_values, try to map them back to the sessions DataFrame
        if row_values is not None:
            # Normalize into list of strings for comparison
            try:
                vals = [
                    "" if pd.isna(v) else str(v)
                    for v in (
                        list(row_values)
                        if not isinstance(row_values, str)
                        else [row_values]
                    )
                ]
            except Exception:
                vals = [str(row_values)]
            # Try to match against the expected columns order we render
            if self._sessions_df is None or self._sessions_df.empty:
                logger.warning(
                    "Sessions DataFrame is empty; cannot map selected row values."
                )
                return
            df_ordered = self._sessions_df.reset_index(drop=True)
            expected_cols = ["otpid", "hostname", "status", "purpose", "granted"]

            # Build stringified candidates for each row in df using the same columns we show
            def _row_to_vals(sr):
                out = []
                for c in expected_cols:
                    if c in sr:
                        v = sr[c]
                        out.append("" if pd.isna(v) else str(v))
                    else:
                        out.append("")
                return out

            match_idx = None
            for i, sr in df_ordered.iterrows():
                cand = _row_to_vals(sr)
                # Compare prefix: row values might be a subset (e.g. only first 3 cols), so compare prefix only
                if len(vals) <= len(cand) and all(
                    vals[j] == cand[j] for j in range(len(vals))
                ):
                    match_idx = i
                    break
            if match_idx is None:
                # Try looser match: compare first cell only (otpid)
                first = vals[0] if vals else None
                if first is not None:
                    for i, sr in df_ordered.iterrows():
                        cand0 = "" if pd.isna(sr.get("otpid")) else str(sr.get("otpid"))
                        if cand0 == first:
                            match_idx = i
                            break
            if match_idx is None:
                logger.warning(
                    "Unable to locate DataFrame row matching selected row values: %r",
                    vals,
                )
                return
            idx = int(match_idx)
        else:
            # 3) If we have a row_key, try to normalize to an int index
            if row_key is not None:
                try:
                    idx = int(row_key)
                except Exception:
                    # Try converting via string
                    try:
                        idx = int(str(row_key))
                    except Exception:
                        idx = None
                if idx is None:
                    # Final numeric fallback: use sessions_table.cursor_row if present
                    try:
                        idx = getattr(self.sessions_table, "cursor_row")
                    except Exception:
                        idx = None
                if idx is None:
                    _bad("Failed to normalize row/key from event: %r", row_key)
                    return
            else:
                # 4) Try table cursor_row as last resort
                try:
                    idx = getattr(self.sessions_table, "cursor_row")
                except Exception:
                    logger.warning(
                        "Could not determine selected row from event: %r", event
                    )
                    # Helpful debug hint for you to paste back if still failing:
                    logger.debug("Event repr for debugging: %r", event)
                    return
        # At this point we should have an integer idx
        try:
            idx = int(idx)
        except Exception:
            logger.exception(
                "Final normalization of selected row index failed: %r", idx
            )
            return
        # Validate sessions df
        if self._sessions_df is None or self._sessions_df.empty:
            logger.warning("Sessions DataFrame empty; nothing to select.")
            return
        df_ordered = self._sessions_df.reset_index(drop=True)
        if idx < 0 or idx >= len(df_ordered):
            logger.warning(
                "Selected row index %s out of range (0..%d)", idx, len(df_ordered) - 1
            )
            return
        row_series = df_ordered.iloc[idx]
        otpid = row_series.get("otpid")
        hostname = row_series.get("hostname")
        # Store selected session and fetch activities
        self._selected_session_otpid = otpid
        # Obtain api from app (try multiple places)
        api = (
            getattr(self.app, "api", None)
            or getattr(self, "api", None)
            or getattr(self.app, "airlock_api", None)
        )
        if api is None:
            logger.error("No API available on self.app.api - cannot fetch activities")
            return
        logger.info(
            "Fetching activities for otpid=%s host=%s (selected row=%s)",
            otpid,
            hostname,
            idx,
        )
        await self._fetch_activities_for_otpid(api, otpid, hostname=hostname)

    async def load_sessions_from_api(self, api) -> None:
        """
        Pulls OTP session lists, adds status column, concatenates and populates the sessions table.
        """
        try:
            active = api.otp_find_active()
            awaiting = api.otp_find_awaiting()
            enforced = api.otp_find_enforced()
            revoked = api.otp_find_revoked()
        except Exception as exc:
            logger.exception("Failed to fetch OTP session lists: %s", exc)
            # Present empty
            active = awaiting = enforced = revoked = pd.DataFrame()

        # Ensure DataFrame objects
        def _ensure_df(df):
            return df if isinstance(df, pd.DataFrame) else pd.DataFrame(df)

        active = _ensure_df(active)
        awaiting = _ensure_df(awaiting)
        enforced = _ensure_df(enforced)
        revoked = _ensure_df(revoked)
        for df, status in [
            (active, "active"),
            (awaiting, "awaiting"),
            (enforced, "enforced"),
            (revoked, "revoked"),
        ]:
            if "status" not in df.columns:
                df["status"] = status
        combined = pd.concat([active, awaiting, enforced, revoked], ignore_index=True)
        if "otpid" in combined.columns:
            combined = combined.sort_values(by="otpid", ascending=False)
        self._sessions_df = combined
        # Populate DataTable
        self.sessions_table.clear()
        # Ensure columns exist in DF and when missing add empty column
        expected_cols = ["otpid", "hostname", "status", "purpose", "granted"]
        for col in expected_cols:
            if col not in combined.columns:
                combined[col] = ""
        self.sessions_table.add_columns(*expected_cols)
        # Add rows
        for _, row in combined[expected_cols].iterrows():
            # Convert values to str for safe insertion
            vals = ["" if pd.isna(v) else v for v in row.to_list()]
            self.sessions_table.add_row(*[str(v) for v in vals])
        logger.info("Loaded %d OTP sessions.", len(combined))

    async def _fetch_activities_for_otpid(self, api, otpid, hostname=None) -> None:
        """
        Fetch activities DataFrame for a given otpid and populate activities_table.
        """
        try:
            result = api.otp_get_activities(otpid)
            result_df = (
                result if isinstance(result, pd.DataFrame) else pd.DataFrame(result)
            )
        except Exception as exc:
            logger.exception("Failed to fetch activities for otpid %s: %s", otpid, exc)
            result_df = pd.DataFrame()
        # Attach hostname if provided
        if hostname is not None:
            result_df["hostname"] = hostname
        if result_df.empty:
            logger.info("No activities found for otpid %s (host: %s)", otpid, hostname)
            self._activities_df = pd.DataFrame()
            self.activities_table.clear()
            return
        # Store and render
        self._activities_df = result_df.copy()
        # Rebuild activities_table columns from result_df
        self.activities_table.clear()
        # Ensure stable column order
        for col in result_df.columns:
            self.activities_table.add_column(col)
        # Add rows
        for _, arow in result_df.iterrows():
            values = ["" if pd.isna(v) else v for v in arow.to_list()]
            self.activities_table.add_row(*[str(v) for v in values])
        logger.info(
            "Loaded %d activity rows for otpid %s (host: %s)",
            len(result_df),
            otpid,
            hostname,
        )

    async def export_activities(self) -> None:
        """
        Export currently-loaded activities DataFrame to CSV.
        Can be called directly (programmatically) or from the button handler.
        """
        if self._activities_df is None or self._activities_df.empty:
            logger.info("No activities loaded to export.")
            # On-screen short message
            await self.post_message(Static("No activities to export."))
            return
        working_dir = _load_working_dir()
        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        filename = f"otp_activities_{self._selected_session_otpid}_{timestamp}.csv"
        file_path = os.path.join(working_dir, filename)
        try:
            self._activities_df.to_csv(file_path, index=False)
            logger.info("Exported activities to %s", file_path)
            await self.post_message(Static(f"✅ Exported activities to: {file_path}"))
        except Exception as exc:
            logger.exception("Failed to export activities to %s: %s", file_path, exc)
            await self.post_message(Static("Failed to export activities; check logs."))


class ActivityDetailWidget(Static):
    """
    Interactive widget for Activity Detail screen.
    Shows the provided DataFrame in a DataTable and offers Export + Back buttons.
    Now includes Select All/None and Add to Allowlist functionality.
    """

    DEFAULT_CSS = """
    ActivityDetailWidget {
        height: 1fr;
        layout: vertical;
    }
    #detail_table_container {
        height: 1fr;
        padding: 1 1;
    }
    #selection_buttons {
        height: auto;
        min-height: 3;
        padding: 1 1;
        content-align: center middle;
    }
    #detail_buttons {
        height: auto;
        min-height: 3;
        padding: 1 1;
        content-align: center middle;
    }
    """

    def __init__(self, activities_df: pd.DataFrame, otpid=None, hostname=None) -> None:
        super().__init__()
        self.activities_df = (
            activities_df.copy()
            if isinstance(activities_df, pd.DataFrame)
            else pd.DataFrame(activities_df)
        )
        # Add a unique identifier column if not present
        if "_row_id" not in self.activities_df.columns:
            self.activities_df["_row_id"] = range(len(self.activities_df))

        self.otpid = otpid
        self.hostname = hostname
        self.selected_row_ids = set()  # Track selected rows by unique ID
        self.row_key_to_id = {}  # Map DataTable row keys to unique row IDs
        self.table_row_to_id = {}  # Map table row indices to unique row IDs
        self._last_sort = None  # Track last sort column and order

    def compose(self) -> ComposeResult:
        yield Static(
            f"Activity Detail (otpid={self.otpid} host={self.hostname})",
            classes="panel-title",
        )
        # Table container
        with Vertical(id="detail_table_container"):
            self.detail_table = DataTable(id="detail_table")
            yield self.detail_table

        # Original buttons at bottom
        with Horizontal(id="detail_buttons"):
            self.detail_back_btn = Button("Back", id="detail_back_btn")
            self.add_allowlist_btn = Button(
                "📋 Add Selected to Allowlist", id="add_allowlist_btn"
            )
            yield self.add_allowlist_btn
            yield self.detail_back_btn

    async def on_mount(self) -> None:
        await self._build_table(rebuild=True)
        self._update_button_states()

    def _update_button_states(self) -> None:
        """Update button states based on selection."""
        has_selection = len(self.selected_row_ids) > 0
        self.add_allowlist_btn.disabled = not has_selection

        # Update button labels with count
        count = len(self.selected_row_ids)
        total = len(self.activities_df)

        if has_selection:
            self.add_allowlist_btn.label = f"📋 Add {count} Selected to Allowlist"
        else:
            self.add_allowlist_btn.label = "📋 Add Selected to Allowlist"

    async def _build_table(self, rebuild: bool = True) -> None:
        """Rebuild the DataTable. If rebuild=False, only refresh rows."""
        if rebuild:
            # Full rebuild: clear columns and rows
            self.detail_table.clear()
            self.detail_table.columns.clear()
            self.row_key_to_id.clear()
            self.table_row_to_id.clear()

            if self.activities_df is None or self.activities_df.empty:
                logger.info("ActivityDetailWidget mounted with empty dataframe.")
                return

            # Add columns (checkbox + data columns, excluding internal _row_id)
            self.detail_table.add_column("Select", key="select")
            for col in self.activities_df.columns:
                if col != "_row_id":  # Don't display the internal ID column
                    self.detail_table.add_column(col)
        else:
            # Partial rebuild: clear rows only
            self.detail_table.clear()
            self.row_key_to_id.clear()
            self.table_row_to_id.clear()

        # Add rows
        for table_idx, (df_idx, row) in enumerate(self.activities_df.iterrows()):
            # Get the unique row ID
            row_id = row["_row_id"]

            # Build values list (excluding _row_id column)
            vals = []
            for col in self.activities_df.columns:
                if col != "_row_id":
                    v = row[col]
                    vals.append("" if pd.isna(v) else str(v))

            # Check if this row is selected
            checkbox = "☑" if row_id in self.selected_row_ids else "☐"

            # Add row to table
            row_key = self.detail_table.add_row(checkbox, *vals)

            # Map the row key and table index to the unique row ID
            self.row_key_to_id[row_key] = row_id
            self.table_row_to_id[table_idx] = row_id

    async def on_data_table_cell_selected(self, event: DataTable.CellSelected) -> None:
        # Toggle selection when the "Select" column is clicked
        if event.cell_key.column_key.value == "select":
            table_row_index = event.coordinate.row

            # Get the unique row ID for this table row
            row_id = self.table_row_to_id.get(table_row_index)
            if row_id is not None:
                # Get the row key for updating the cell
                row_key = event.cell_key.row_key

                if row_id in self.selected_row_ids:
                    self.selected_row_ids.remove(row_id)
                    self.detail_table.update_cell(row_key, "select", "☐")  # Unchecked
                else:
                    self.selected_row_ids.add(row_id)
                    self.detail_table.update_cell(row_key, "select", "☑")  # Checked

                self._update_button_states()

    async def on_data_table_header_selected(
        self, event: DataTable.HeaderSelected
    ) -> None:
        column_key = event.column_key.value if event.column_key else None
        if not column_key:
            col_index = event.column_index
            if col_index == 0:  # First column is "Select"
                return
            # Adjust for hidden _row_id column
            visible_cols = [
                col for col in self.activities_df.columns if col != "_row_id"
            ]
            if col_index - 1 < len(visible_cols):
                column_key = visible_cols[col_index - 1]
            else:
                return
        if column_key == "select" or column_key == "_row_id":
            return

        ascending = True
        if self._last_sort == (column_key, True):
            ascending = False
        self._last_sort = (column_key, ascending)

        try:
            self.activities_df.sort_values(
                by=column_key, ascending=ascending, inplace=True
            )
        except Exception as exc:
            logger.exception("Failed to sort by column %s: %s", column_key, exc)
            return

        # âœ… Only refresh rows, not columns
        await self._build_table(rebuild=False)

    async def on_button_pressed(self, event) -> None:
        btn = getattr(event, "button", None) or getattr(event, "sender", None)
        btn_id = getattr(btn, "id", None) or getattr(event, "button_id", None)

        if btn is self.detail_back_btn or btn_id == "detail_back_btn":
            while len(self.app.screen_stack) > 2:
                self.app.pop_screen()
            event.stop()
            return

        if btn is self.add_allowlist_btn or btn_id == "add_allowlist_btn":
            await self._open_allowlist_screen()
            return

    async def _select_all(self) -> None:
        """Select all rows in the table."""
        # Add all row IDs to selected set
        self.selected_row_ids = set(self.activities_df["_row_id"].tolist())

        # Update all checkboxes in the table
        for row_key, row_id in self.row_key_to_id.items():
            self.detail_table.update_cell(row_key, "select", "☑")

        self._update_button_states()
        logger.info(f"Selected all {len(self.selected_row_ids)} rows")

    async def _select_none(self) -> None:
        """Deselect all rows in the table."""
        # Clear selected set
        self.selected_row_ids.clear()

        # Update all checkboxes in the table
        for row_key, row_id in self.row_key_to_id.items():
            self.detail_table.update_cell(row_key, "select", "☐")

        self._update_button_states()
        logger.info("Cleared all selections")

    async def _open_allowlist_screen(self) -> None:
        """Open the allowlist selection screen with selected activities."""
        if not self.selected_row_ids:
            self.app.notify("No rows selected", severity="warning")
            return

        # Get selected data
        selected_df = self.get_selected_data()

        # Get API from app
        api = getattr(self.app, "api", None)
        if api is None:
            logger.error("No API available on self.app.api")
            self.app.notify("API not available", severity="error")
            return

        # Create and push AllowlistSelectionScreen
        try:
            allowlist_screen = AllowlistSelectionScreen(
                selected_df, api=api, hostname=self.hostname, otpid=self.otpid
            )
            await self.app.push_screen(allowlist_screen)
            logger.info(
                f"Opened allowlist screen with {len(selected_df)} selected activities"
            )
        except ImportError as e:
            logger.error(f"Failed to import AllowlistSelectionScreen: {e}")
            self.app.notify("Allowlist screen module not found", severity="error")
        except Exception as e:
            logger.exception(f"Failed to open allowlist screen: {e}")
            self.app.notify(
                f"Error opening allowlist screen: {str(e)}", severity="error"
            )

    async def _export_detail_activities(self) -> None:
        if self.activities_df is None or self.activities_df.empty:
            logger.info("No activities to export.")
            await self.mount(
                Static("âŒ No activities to export.", classes="notification")
            )
            return
        if not self.selected_row_ids:
            logger.info("No rows selected for export.")
            await self.mount(
                Static("âŒ No rows selected for export.", classes="notification")
            )
            return
        try:
            working_dir = load_env("WORKING_DIR") or os.getcwd()
            timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
            filename = f"otp_activities_detail_{timestamp}.csv"
            file_path = os.path.join(working_dir, filename)
            selected_df = self.get_selected_data()
            selected_df.to_csv(file_path, index=False)
            logger.info("Exported selected activities to %s", file_path)
            await self.mount(
                Static(
                    f"✅ Exported selected activities to: {filename}",
                    classes="notification",
                )
            )
        except Exception as exc:
            logger.exception("Failed to export detail activities: %s", exc)
            await self.mount(
                Static(
                    "âŒ Failed to export activities; check logs.",
                    classes="notification",
                )
            )

    # âœ… Helper methods
    def get_selected_data(self) -> pd.DataFrame:
        """Return a DataFrame of the selected rows."""
        if not self.selected_row_ids:
            return pd.DataFrame()
        # Filter by selected row IDs and drop the internal _row_id column
        selected_df = self.activities_df[
            self.activities_df["_row_id"].isin(self.selected_row_ids)
        ].copy()
        if "_row_id" in selected_df.columns:
            selected_df = selected_df.drop(columns=["_row_id"])
        return selected_df

    def get_selected_records(self) -> list[dict]:
        """Return selected rows as a list of dicts."""
        if not self.selected_row_ids:
            return []
        # Filter by selected row IDs and drop the internal _row_id column
        selected_df = self.activities_df[
            self.activities_df["_row_id"].isin(self.selected_row_ids)
        ].copy()
        if "_row_id" in selected_df.columns:
            selected_df = selected_df.drop(columns=["_row_id"])
        return selected_df.to_dict(orient="records")


class ActivityDetailScreen(Screen):
    """
    Screen that wraps ActivityDetailWidget. Expects a DataFrame passed on init.
    """

    BINDINGS = [
        Binding("b", "back", "Back"),
        Binding("e", "export", "Export"),
        Binding("a", "select_all", "Select All"),
        Binding("n", "select_none", "Select None"),
    ]

    def __init__(self, activities_df: pd.DataFrame, otpid=None, hostname=None) -> None:
        super().__init__()
        self._activities_df = (
            activities_df.copy()
            if isinstance(activities_df, pd.DataFrame)
            else pd.DataFrame(activities_df)
        )
        self._otpid = otpid
        self._hostname = hostname

    def compose(self) -> ComposeResult:
        self.widget = ActivityDetailWidget(
            self._activities_df, otpid=self._otpid, hostname=self._hostname
        )
        yield Header(show_clock=True)
        yield self.widget
        yield Footer()

    async def action_back(self) -> None:
        try:
            await self.app.pop_screen()
        except Exception:
            logger.debug("ActivityDetailScreen.action_back pop_screen failed.")

    async def action_export(self) -> None:
        # Delegate to widget export helper
        if hasattr(self, "widget") and self.widget is not None:
            await self.widget._export_detail_activities()

    async def action_select_all(self) -> None:
        """Handle 'a' key for select all."""
        if hasattr(self, "widget") and self.widget is not None:
            await self.widget._select_all()

    async def action_select_none(self) -> None:
        """Handle 'n' key for select none."""
        if hasattr(self, "widget") and self.widget is not None:
            await self.widget._select_none()


class OTPActivitiesScreen(Screen):
    """
    A Screen intended to be pushed into an existing Textual App.
    Usage:
        app.push_screen(OTPActivitiesScreen())
    or create this screen and call `await screen.load()` inside your app lifecycle.
    The screen expects `self.app.api` to exist and be an AirlockAPIWrapper instance.
    """

    BINDINGS = [
        Binding("r", "refresh_sessions", "Refresh Sessions"),
        Binding("e", "export_activities", "Export activities"),
        Binding("q", "quit", "Quit"),
    ]

    def compose(self) -> ComposeResult:
        yield Header()
        self.widget = OTPActivitiesWidget()
        yield self.widget
        yield Footer()

    async def on_show(self) -> None:
        """Restore focus to the left sessions table when the screen becomes visible."""
        if hasattr(self, "widget") and hasattr(self.widget, "sessions_table"):
            self.widget.sessions_table.focus()

    async def on_mount(self) -> None:
        # Try to load sessions immediately
        api = getattr(self.app, "api", None)
        if api is None:
            logger.warning("OTPActivitiesScreen mounted but no self.app.api found.")
        else:
            await self.widget.load_sessions_from_api(api)

    # Simple actions bound to keys
    async def action_refresh_sessions(self) -> None:
        api = getattr(self.app, "api", None)
        if api is None:
            logger.error("No API on app; cannot refresh sessions.")
            return
        logger.info("Refreshing OTP sessions via API.")
        await self.widget.load_sessions_from_api(api)

    async def action_quit(self) -> None:
        # Pop the screen or exit app
        await self.app.pop_screen()

    # If you want an explicit method to fetch activities for a particular otpid from outside:
    async def fetch_activities_for_otpid(self, otpid, hostname=None) -> None:
        api = getattr(self.app, "api", None)
        if api is None:
            logger.error("No API on app; cannot fetch activities.")
            return
        await self.widget._fetch_activities_for_otpid(api, otpid, hostname=hostname)
