#    Copyright (C) 2025 James Brotosky, Brandon Wickline
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as published
#    by the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.

from __future__ import annotations

import logging
from typing import List, Optional

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

logger = logging.getLogger(__name__)


class OTPRevokeWidget(Static):
    """
    Widget for managing OTP session revocation.
    Displays active OTP sessions and allows selection for revocation.
    """

    class SessionsRevoked(Message):
        """Message sent when sessions are revoked."""

        def __init__(self, revoked_sessions: List[dict]):
            super().__init__()
            self.revoked_sessions = revoked_sessions

    DEFAULT_CSS = """
    OTPRevokeWidget {
        height: 1fr;
    }
    #main_container {
        width: 100%;
        height: 100%;
        layout: vertical;
    }
    #sessions_container {
        height: 1fr;
        border: none;
        padding: 1;
    }
    #button_container {
        height: auto;
        width: 100%;
        padding: 1;
        align: center middle;
    }
    #button_container Button {
        min-width: 16;
        margin: 0 1;
    }
    #result_container {
        height: auto;
        max-height: 10;
        border: solid #444444;
        padding: 1;
        margin: 1;
        overflow-y: auto;
    }
    .panel-title {
        text-style: bold;
        margin: 0 0 1 0;
    }
    """

    def compose(self) -> ComposeResult:
        with Vertical(id="main_container"):
            # Sessions table
            yield Static("OTP Sessions", classes="panel-title")
            with Vertical(id="sessions_container"):
                self.sessions_table = DataTable(id="sessions_table")
                self.sessions_table.styles.width = "100%"
                self.sessions_table.styles.height = "1fr"
                yield self.sessions_table

            # Action buttons
            with Horizontal(id="button_container"):
                yield Button("Refresh", id="refresh_btn")
                yield Button("Select All", id="select_all_btn")
                yield Button("Clear Selection", id="select_none_btn")
                yield Button("Revoke Selected", id="revoke_btn", variant="error")

            # Results display
            with Vertical(id="result_container"):
                yield Static("Revocation Results", classes="panel-title")
                self.results_display = Static("No actions performed yet.")
                yield self.results_display

    async def on_mount(self) -> None:
        """Initialize the widget when mounted."""
        # Configure sessions table
        self.sessions_table.clear()
        self.sessions_table.add_columns(
            "", "OTP ID", "Hostname", "Status", "Purpose", "Granted"
        )

        # Enable row selection with checkbox column
        self.sessions_table.cursor_type = "row"
        try:
            self.sessions_table.zebra_stripes = True
        except Exception:
            pass

        # Initialize state
        self._sessions_df: Optional[pd.DataFrame] = None
        self._filtered_df: Optional[pd.DataFrame] = None
        self._selected_otpids: set = set()

    async def load_sessions_from_api(self, api) -> None:
        """Load active OTP sessions from the API."""
        try:
            # Fetch only active sessions
            active_df = api.otp_find_active()

            # Ensure we have a DataFrame
            if not isinstance(active_df, pd.DataFrame):
                active_df = pd.DataFrame(active_df)

            # Add status column
            active_df["status"] = "active"

            # Sort by otpid if column exists
            if "otpid" in active_df.columns and not active_df.empty:
                active_df = active_df.sort_values(by="otpid", ascending=False)

            # Store the full dataframe
            self._sessions_df = active_df
            self._filtered_df = active_df.copy()

            # Display in table
            await self._refresh_table()

            # Update status
            active_count = len(active_df)

            status_msg = f"Loaded {active_count} active sessions"
            logger.info(status_msg)
            self.results_display.update(status_msg)

        except Exception as e:
            logger.exception(f"Failed to load OTP sessions: {e}")
            self.results_display.update(f"Error loading sessions: {str(e)}")

    async def _refresh_table(self) -> None:
        """Refresh the table display with current filtered data."""
        if self._filtered_df is None or self._filtered_df.empty:
            self.sessions_table.clear()
            return

        # Ensure expected columns exist
        expected_cols = ["otpid", "hostname", "status", "purpose", "granted"]
        for col in expected_cols:
            if col not in self._filtered_df.columns:
                self._filtered_df[col] = ""

        # Clear and repopulate table
        self.sessions_table.clear(columns=False)

        for _, row in self._filtered_df.iterrows():
            otpid = str(row.get("otpid", ""))
            # Check if this row is selected
            checkbox = "☑️" if otpid in self._selected_otpids else "☐"

            self.sessions_table.add_row(
                checkbox,
                str(otpid),
                str(row.get("hostname", "")),
                str(row.get("status", "")),
                str(row.get("purpose", "")),
                str(row.get("granted", "")),
            )

    async def on_button_pressed(self, event) -> None:
        """Handle button presses."""
        btn = event.button

        if btn.id == "refresh_btn":
            # Refresh sessions
            api = getattr(self.app, "api", None)
            if api:
                await self.load_sessions_from_api(api)

        elif btn.id == "select_all_btn":
            # Select all visible rows
            if (
                self._filtered_df is not None
                and not self._filtered_df.empty
                and "otpid" in self._filtered_df.columns
            ):
                self._selected_otpids = set(str(x) for x in self._filtered_df["otpid"])
                await self._refresh_table()

        elif btn.id == "select_none_btn":
            # Clear selection
            self._selected_otpids.clear()
            await self._refresh_table()

        elif btn.id == "revoke_btn":
            # Revoke selected sessions
            await self._revoke_selected()

    async def on_data_table_row_selected(self, event) -> None:
        """Handle row selection in the table."""
        if event.data_table != self.sessions_table:
            return

        try:
            # Get the row index from the cursor row
            row_index = self.sessions_table.cursor_row

            if (
                self._filtered_df is not None
                and not self._filtered_df.empty
                and "otpid" in self._filtered_df.columns
                and row_index < len(self._filtered_df)
            ):
                # Get the OTP ID for this row
                otpid = str(self._filtered_df.iloc[row_index]["otpid"])

                # Toggle selection
                if otpid in self._selected_otpids:
                    self._selected_otpids.remove(otpid)
                else:
                    self._selected_otpids.add(otpid)

                # Refresh table to update checkbox
                await self._refresh_table()

                # Restore cursor position
                self.sessions_table.move_cursor(row=row_index)

        except Exception as e:
            logger.exception(f"Error handling row selection: {e}")

    async def _revoke_selected(self) -> None:
        """Revoke the selected OTP sessions."""
        if not self._selected_otpids:
            self.results_display.update("No sessions selected for revocation")
            return

        api = getattr(self.app, "api", None)
        if not api:
            self.results_display.update("API not available")
            return

        # Collect results
        results = []
        success_count = 0
        failure_count = 0

        for otpid in self._selected_otpids:
            try:
                # Get hostname for this session
                hostname = "Unknown"
                if self._sessions_df is not None:
                    # Convert otpid to same type as in DataFrame for comparison
                    otpid_compare = otpid
                    if len(self._sessions_df) > 0:
                        first_otpid = self._sessions_df["otpid"].iloc[0]
                        if isinstance(first_otpid, int):
                            otpid_compare = int(otpid)

                    match = self._sessions_df[
                        self._sessions_df["otpid"] == otpid_compare
                    ]
                    if not match.empty:
                        hostname = match.iloc[0].get("hostname", "Unknown")

                # Revoke the session
                result = api.otp_revoke(otpid)

                if result and result.get("status") != "error":
                    success_count += 1
                    results.append(f"Revoked OTP {otpid} for {hostname}")
                    logger.info(f"Revoked OTP {otpid} for {hostname}: {result}")
                else:
                    failure_count += 1
                    error_msg = (
                        result.get("message", "Unknown error")
                        if result
                        else "No response"
                    )
                    results.append(
                        f"Failed to revoke OTP {otpid} for {hostname}: {error_msg}"
                    )
                    logger.error(f"Failed to revoke OTP {otpid}: {error_msg}")

            except Exception as e:
                failure_count += 1
                results.append(f"Error revoking OTP {otpid}: {str(e)}")
                logger.exception(f"Exception revoking OTP {otpid}: {e}")

        # Update results display
        summary = (
            f"Revocation complete: {success_count} succeeded, {failure_count} failed\n"
        )
        details = "\n".join(results[-5:])  # Show last 5 results
        if len(results) > 5:
            details = f"... (showing last 5 of {len(results)} results)\n" + details

        self.results_display.update(summary + details)

        # Clear selection and refresh
        self._selected_otpids.clear()
        await self.load_sessions_from_api(api)

        # Post message about revoked sessions
        if success_count > 0:
            self.post_message(self.SessionsRevoked(results))


class OTPRevokeScreen(Screen):
    """
    Main screen for OTP session revocation workflow.
    This replaces the otp_revoke function from otp.py.
    """

    BINDINGS = [
        Binding("escape", "go_back", "Back"),
        Binding("q", "main_menu", "Main Menu"),
        Binding("r", "refresh", "Refresh"),
        Binding("a", "select_all", "Select All"),
        Binding("n", "select_none", "Clear Selection"),
        Binding("d", "revoke", "Revoke Selected"),
    ]

    def compose(self) -> ComposeResult:
        yield Header(show_clock=True)
        self.widget = OTPRevokeWidget()
        yield self.widget
        yield Footer()

    async def on_mount(self) -> None:
        """Load sessions when screen mounts."""
        api = getattr(self.app, "api", None)
        if api:
            await self.widget.load_sessions_from_api(api)
        else:
            logger.warning("OTPRevokeScreen mounted but no self.app.api found.")

    async def action_refresh(self) -> None:
        """Refresh the sessions list."""
        api = getattr(self.app, "api", None)
        if api:
            await self.widget.load_sessions_from_api(api)

    async def action_select_all(self) -> None:
        """Select all visible sessions."""
        if (
            self.widget._filtered_df is not None
            and not self.widget._filtered_df.empty
            and "otpid" in self.widget._filtered_df.columns
        ):
            self.widget._selected_otpids = set(
                str(x) for x in self.widget._filtered_df["otpid"]
            )
            await self.widget._refresh_table()

    async def action_select_none(self) -> None:
        """Clear all selections."""
        self.widget._selected_otpids.clear()
        await self.widget._refresh_table()

    async def action_revoke(self) -> None:
        """Revoke selected sessions."""
        await self.widget._revoke_selected()

    async def action_go_back(self) -> None:
        """Go back to previous screen."""
        await self.app.pop_screen()

    async def action_main_menu(self) -> None:
        """Go back to main menu."""
        while len(self.app.screen_stack) > 2:
            await self.app.pop_screen()
