#    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 datetime import datetime, timedelta
import logging
import os
from typing import List

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,
    Label,
    Select,
    Static,
)

from models.agent import Agent
from models.execution import ExecutionHistoryRecord
from utils.configmanager import load_env

logger = logging.getLogger(__name__)


class ExecutionHistoryScreen(Screen):
    """
    A screen for viewing and exporting execution history for selected agents.

    This screen allows users to:
    1. Select a start date and end date using dropdown selects
    2. Fetch execution history for all selected agents
    3. View the results in a DataTable
    4. Export the results to CSV using a keybinding

    Attributes:
        agents (List[Agent]): List of agents to fetch execution history for
        execution_data (pd.DataFrame): Combined execution history data
        working_dir (str): Directory for CSV exports
    """

    DEFAULT_CSS = """
    ExecutionHistoryScreen {
        align: center top;
    }
    
    #main_container {
        width: 95%;
        height: 1fr;
        border: solid $primary;
        padding: 1;
    }
    
    #title {
        text-style: bold;
        color: $text;
        text-align: center;
        margin-bottom: 1;
    }
    
    #date_container {
        height: auto;
        margin-bottom: 1;
    }
    
    #start_date_row, #end_date_row {
        height: auto;
        align-horizontal: left;
        margin-bottom: 1;
    }
    
    .date_label {
        width: 8;
        margin-right: 1;
    }
    
    .date_selector {
        width: 18;
        margin: 0 1;
    }
    
    #quick_buttons_row {
        height: auto;
        align-horizontal: center;
        margin-bottom: 1;
    }
    
    .quick_select_btn {
        margin: 0 1;
    }
    
    #button_row {
        height: auto;
        align-horizontal: center;
        margin-top: 1;
        margin-bottom: 1;
    }
    
    Button {
        margin: 0 1;
    }
    
    #status_label {
        text-align: center;
        color: $accent;
        margin-bottom: 1;
    }
    
    #results_container {
        height: 1fr;
        display: none;
    }
    
    #results_button_row {
        height: auto;
        align-horizontal: center;
        margin-bottom: 1;
    }
    
    #history_table {
        height: 1fr;
        border: solid $primary;
    }
    
    DataTable > .datatable--header {
        text-style: bold;
        background: $primary 20%;
    }
    """

    BINDINGS = [
        Binding("escape", "close_screen", "Close"),
        Binding("e", "export_csv", "Export CSV"),
        Binding("q", "close_screen", "Quit"),
    ]

    def __init__(self, agents: List[Agent]):
        """
        Initialize the ExecutionHistoryScreen.

        Args:
            agents (List[Agent]): List of agents to fetch execution history for
        """
        super().__init__()
        self.agents = agents
        self.execution_data = pd.DataFrame()
        self.working_dir = load_env("WORKING_DIR") or os.getcwd()

        # Generate dropdown options
        today = datetime.now().date()

        # Month options - format is (display_text, value)
        self.month_options = [
            ("January", "01"),
            ("February", "02"),
            ("March", "03"),
            ("April", "04"),
            ("May", "05"),
            ("June", "06"),
            ("July", "07"),
            ("August", "08"),
            ("September", "09"),
            ("October", "10"),
            ("November", "11"),
            ("December", "12"),
        ]

        # Day options (1-31) - format is (display_text, value)
        self.day_options = [(f"{i}", f"{i:02d}") for i in range(1, 32)]

        # Year options (current year back 5 years) - format is (display_text, value)
        current_year = today.year
        self.year_options = [
            (str(year), str(year)) for year in range(current_year, current_year - 6, -1)
        ]

        # Default dates: last 30 days
        start_date = today - timedelta(days=30)
        self.start_month = f"{start_date.month:02d}"
        self.start_day = f"{start_date.day:02d}"
        self.start_year = str(start_date.year)

        self.end_month = f"{today.month:02d}"
        self.end_day = f"{today.day:02d}"
        self.end_year = str(today.year)

    def compose(self) -> ComposeResult:
        """Build the UI layout."""
        yield Header(show_clock=True, icon="📊")

        with Vertical(id="main_container"):
            title_text = f"Execution History - {len(self.agents)} Agent(s)"
            yield Static(title_text, id="title")

            # Date selection area
            with Vertical(id="date_container"):
                yield Label("Select Date Range:")

                # Start date row
                with Horizontal(id="start_date_row"):
                    yield Label("From:", classes="date_label")
                    yield Select(
                        options=self.month_options,
                        value=self.start_month,
                        id="start_month_select",
                        classes="date_selector",
                    )
                    yield Select(
                        options=self.day_options,
                        value=self.start_day,
                        id="start_day_select",
                        classes="date_selector",
                    )
                    yield Select(
                        options=self.year_options,
                        value=self.start_year,
                        id="start_year_select",
                        classes="date_selector",
                    )

                # End date row
                with Horizontal(id="end_date_row"):
                    yield Label("To:", classes="date_label")
                    yield Select(
                        options=self.month_options,
                        value=self.end_month,
                        id="end_month_select",
                        classes="date_selector",
                    )
                    yield Select(
                        options=self.day_options,
                        value=self.end_day,
                        id="end_day_select",
                        classes="date_selector",
                    )
                    yield Select(
                        options=self.year_options,
                        value=self.end_year,
                        id="end_year_select",
                        classes="date_selector",
                    )

                # Quick select buttons
                with Horizontal(id="quick_buttons_row"):
                    yield Button(
                        "1 Day",
                        id="quick_1day",
                        classes="quick_select_btn",
                        variant="default",
                    )
                    yield Button(
                        "1 Week",
                        id="quick_1week",
                        classes="quick_select_btn",
                        variant="default",
                    )
                    yield Button(
                        "30 Days",
                        id="quick_30days",
                        classes="quick_select_btn",
                        variant="default",
                    )

            # Buttons
            with Horizontal(id="button_row"):
                yield Button("Fetch History", id="fetch_btn", variant="primary")
                yield Button("Close", id="close_btn", variant="error")

            # Status
            yield Static(
                "Select date range and click 'Fetch History'", id="status_label"
            )

            # Results container (hidden initially, shown after fetch)
            with Vertical(id="results_container"):
                with Horizontal(id="results_button_row"):
                    yield Button("Export CSV", id="export_btn", variant="success")
                    yield Button("Back", id="back_btn", variant="default")
                yield DataTable(id="history_table")

        yield Footer()

    def on_mount(self) -> None:
        """Initialize the table when screen is mounted."""
        table = self.query_one("#history_table", DataTable)
        table.cursor_type = "row"
        table.zebra_stripes = True

        # Initially empty - will populate after fetch
        logger.info(f"ExecutionHistoryScreen mounted with {len(self.agents)} agents")

    def on_select_changed(self, event: Select.Changed) -> None:
        """Handle date selection changes."""
        select_id = event.select.id

        if select_id == "start_month_select":
            self.start_month = event.value
            logger.debug(f"Start month changed to: {self.start_month}")
        elif select_id == "start_day_select":
            self.start_day = event.value
            logger.debug(f"Start day changed to: {self.start_day}")
        elif select_id == "start_year_select":
            self.start_year = event.value
            logger.debug(f"Start year changed to: {self.start_year}")
        elif select_id == "end_month_select":
            self.end_month = event.value
            logger.debug(f"End month changed to: {self.end_month}")
        elif select_id == "end_day_select":
            self.end_day = event.value
            logger.debug(f"End day changed to: {self.end_day}")
        elif select_id == "end_year_select":
            self.end_year = event.value
            logger.debug(f"End year changed to: {self.end_year}")

    def _set_quick_date_range(self, days: int) -> None:
        """Set the date range based on quick select button."""
        today = datetime.now().date()
        start_date = today - timedelta(days=days)

        # Update internal values
        self.start_month = f"{start_date.month:02d}"
        self.start_day = f"{start_date.day:02d}"
        self.start_year = str(start_date.year)

        self.end_month = f"{today.month:02d}"
        self.end_day = f"{today.day:02d}"
        self.end_year = str(today.year)

        # Update the Select widgets
        try:
            self.query_one("#start_month_select", Select).value = self.start_month
            self.query_one("#start_day_select", Select).value = self.start_day
            self.query_one("#start_year_select", Select).value = self.start_year

            self.query_one("#end_month_select", Select).value = self.end_month
            self.query_one("#end_day_select", Select).value = self.end_day
            self.query_one("#end_year_select", Select).value = self.end_year

            self.app.notify(
                f"Date range set to last {days} day(s)",
                severity="information",
                timeout=2,
            )
            logger.info(f"Quick select: Set date range to last {days} days")
        except Exception as e:
            logger.error(f"Failed to update date selects: {e}")

    def _show_date_selection(self) -> None:
        """Show the date selection view and hide results."""
        try:
            self.query_one("#date_container").styles.display = "block"
            self.query_one("#button_row").styles.display = "block"
            self.query_one("#status_label").styles.display = "block"
            self.query_one("#results_container").styles.display = "none"
        except Exception as e:
            logger.error(f"Failed to show date selection: {e}")

    def _show_results(self) -> None:
        """Hide date selection view and show results."""
        try:
            self.query_one("#date_container").styles.display = "none"
            self.query_one("#button_row").styles.display = "none"
            self.query_one("#status_label").styles.display = "none"
            self.query_one("#results_container").styles.display = "block"
        except Exception as e:
            logger.error(f"Failed to show results: {e}")

    def on_button_pressed(self, event: Button.Pressed) -> None:
        """Handle button clicks."""
        if event.button.id == "fetch_btn":
            self._fetch_execution_history()
        elif event.button.id == "export_btn":
            self._export_to_csv()
        elif event.button.id == "close_btn":
            self.app.pop_screen()
        elif event.button.id == "back_btn":
            self._show_date_selection()
        elif event.button.id == "quick_1day":
            self._set_quick_date_range(days=1)
        elif event.button.id == "quick_1week":
            self._set_quick_date_range(days=7)
        elif event.button.id == "quick_30days":
            self._set_quick_date_range(days=30)

    def _fetch_execution_history(self) -> None:
        """Fetch execution history for all selected agents."""
        status_label = self.query_one("#status_label", Static)
        status_label.update("⏳ Fetching execution history...")

        # Disable buttons during fetch
        fetch_btn = self.query_one("#fetch_btn", Button)
        export_btn = self.query_one("#export_btn", Button)
        fetch_btn.disabled = True
        export_btn.disabled = True

        api = self.app.api
        all_history = []

        try:
            # Construct dates from dropdowns
            start_date_str = f"{self.start_year}-{self.start_month}-{self.start_day}"
            end_date_str = f"{self.end_year}-{self.end_month}-{self.end_day}"

            # Validate dates
            try:
                start_dt = datetime.strptime(start_date_str, "%Y-%m-%d")
                end_dt = datetime.strptime(end_date_str, "%Y-%m-%d")
            except ValueError as e:
                status_label.update(f"❌ Invalid date: {str(e)}")
                fetch_btn.disabled = False
                export_btn.disabled = False
                self.app.notify(f"Invalid date selected: {str(e)}", severity="error")
                return

            if start_dt > end_dt:
                status_label.update("❌ Error: Start date must be before end date")
                fetch_btn.disabled = False
                export_btn.disabled = False
                return

            # Fetch history for each agent
            for i, agent in enumerate(self.agents):
                try:
                    status_label.update(
                        f"⏳ Fetching history for {agent.hostname} ({i+1}/{len(self.agents)})..."
                    )

                    # Call API - note the API expects 'dateto' first, then 'datefrom'
                    history = api.history_execution(
                        today=end_date_str,
                        date_selected=start_date_str,
                        agent_name=agent.hostname,
                    )

                    if history:
                        # Add agent hostname to each record for identification
                        for record in history:
                            record["agent_hostname"] = agent.hostname
                        all_history.extend(history)
                        logger.info(
                            f"Fetched {len(history)} records for {agent.hostname}"
                        )
                    else:
                        logger.info(f"No history found for {agent.hostname}")

                except Exception as e:
                    logger.error(f"Failed to fetch history for {agent.hostname}: {e}")
                    self.app.notify(
                        f"Warning: Failed to fetch history for {agent.hostname}",
                        severity="warning",
                    )

            # Convert to DataFrame
            if all_history:
                status_label.update(
                    "⏳ Enriching execution data with hash information..."
                )

                # Normalize field names (handle API typos)
                for record in all_history:
                    if "policver" in record and "policyver" not in record:
                        record["policyver"] = record.pop("policver")

                # Convert dict records to ExecutionHistoryRecord objects
                execution_records = []
                for record in all_history:
                    try:
                        execution_records.append(ExecutionHistoryRecord(**record))
                    except TypeError as e:
                        logger.warning(f"Failed to create ExecutionHistoryRecord: {e}")
                        # If it fails, just keep the dict
                        continue

                # Enrich with hash data if we have ExecutionHistoryRecord objects
                if execution_records:
                    try:
                        enriched_records = ExecutionHistoryRecord.enrich_with_hashes(
                            api, execution_records
                        )
                        logger.info(
                            f"Enriched {len(enriched_records)} records with hash data"
                        )

                        # Convert back to DataFrame
                        self.execution_data = pd.DataFrame(
                            [r.__dict__ for r in enriched_records]
                        )

                        # Flatten hash_obj if present
                        if (
                            not self.execution_data.empty
                            and "hash_obj" in self.execution_data.columns
                        ):
                            hash_df = self.execution_data["hash_obj"].apply(
                                lambda h: (
                                    h.to_dict() if h and hasattr(h, "to_dict") else {}
                                )
                            )
                            self.execution_data = pd.concat(
                                [
                                    self.execution_data.drop(columns=["hash_obj"]),
                                    hash_df,
                                ],
                                axis=1,
                            )
                    except Exception as e:
                        logger.warning(f"Failed to enrich with hashes: {e}")
                        # Fall back to plain DataFrame
                        self.execution_data = pd.DataFrame(all_history)
                else:
                    # If we couldn't create any ExecutionHistoryRecord objects, just use raw data
                    self.execution_data = pd.DataFrame(all_history)

                self._populate_table()
                self._show_results()  # Switch to results view
                self.app.notify(
                    f"Successfully loaded {len(self.execution_data)} records",
                    severity="information",
                )
            else:
                status_label.update(
                    "ℹ️ No execution history found for selected agents/dates"
                )
                self.app.notify("No execution history found", severity="information")
                self.execution_data = pd.DataFrame()

        except Exception as e:
            logger.error(f"Error fetching execution history: {e}")
            status_label.update(f"❌ Error: {str(e)}")
            self.app.notify(f"Failed to fetch history: {str(e)}", severity="error")

        finally:
            # Re-enable buttons
            fetch_btn.disabled = False
            export_btn.disabled = False

    def _populate_table(self) -> None:
        """Populate the DataTable with execution history data."""
        table = self.query_one("#history_table", DataTable)
        table.clear(columns=True)

        if self.execution_data.empty:
            return

        # Define preferred column order (your specified order)
        preferred_order = [
            "policyname",
            "policyver",
            "hostname",
            "username",
            "publisher",
            "filename",
            "pprocess",
            "gprocess",
            "sha256",
            "commandline",
            "agent_hostname",  # Our custom field
        ]

        # Get available columns in preferred order, then add any remaining columns
        available_cols = []
        for col in preferred_order:
            if col in self.execution_data.columns:
                available_cols.append(col)

        # Add any remaining columns not in preferred order
        for col in self.execution_data.columns:
            if col not in available_cols:
                available_cols.append(col)

        # Add columns to table
        for col in available_cols:
            table.add_column(col, key=col)

        # Add rows
        for idx, row in self.execution_data.iterrows():
            row_data = []
            for col in available_cols:
                value = row[col]
                # Convert to string, handle None/NaN
                if pd.isna(value):
                    row_data.append("")
                else:
                    row_data.append(str(value))
            table.add_row(*row_data, key=str(idx))

        logger.info(f"Populated table with {len(self.execution_data)} rows")

    def _export_to_csv(self) -> None:
        """Export the current execution data to CSV."""
        if self.execution_data.empty:
            self.app.notify("No data to export", severity="warning")
            return

        try:
            # Create filename with timestamp
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"execution_history_{timestamp}.csv"
            filepath = os.path.join(self.working_dir, filename)

            # Export to CSV
            self.execution_data.to_csv(filepath, index=False, encoding="utf-8-sig")

            self.app.notify(
                f"✅ Exported {len(self.execution_data)} records to: {filepath}",
                severity="information",
                timeout=5,
            )
            logger.info(f"Exported execution history to: {filepath}")

        except Exception as e:
            logger.error(f"Failed to export CSV: {e}")
            self.app.notify(f"Failed to export CSV: {str(e)}", severity="error")

    def action_export_csv(self) -> None:
        """Keybinding action to export CSV."""
        self._export_to_csv()

    def action_close_screen(self) -> None:
        """Close this screen and return to previous."""
        self.app.pop_screen()