#    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/>.

import datetime
import logging

import dotenv
import pandas as pd

from flows.prepPolicy import selectPolicies
from services.API import AirlockAPIWrapper
from services.policyhandler import getPolicyInfo
from utils.selector import Selector
from utils.utils import colorText, get_sanitized_input, load_env

logger = logging.getLogger(__name__)


dotenv.load_dotenv()


def findQuietAgents(api: AirlockAPIWrapper):
    working_dir = load_env("WORKING_DIR")
    # Get policy selection and agent list
    selected_policy = selectPolicies(api, False)
    if selected_policy:
        agents = api.agents_find_by_group(selected_policy[0].groupid)

        # Prompt user for history range
        history_days = Selector.select_value(
            prompt="Enter how many days of history to pull (1–150): ",
            value_type=int,
            valid_range=(1, 150),
        )
        required_quiet = Selector.select_value(
            prompt="Enter how many days without an untrusted execution before these are considered ready for enforcement? (1–365): ",
            value_type=int,
            valid_range=(1, 150),
        )

        confirm = Selector.confirm(f"Do you wish to proceed to pull history for {selected_policy[0].name}? Y/N :  ")
        # Get execution history as a DataFrame
        if confirm:
            policy_exec_history = getPolicyInfo(
            api, selected_policy[0], [1, 2, 6, 7], history_days
            )
        
       
            if policy_exec_history.empty:
                logging.info("No execution history found for the selected policy and time range.")
                get_sanitized_input("Press enter to continue")
                return


            # Convert 'datetime' column to timezone-aware datetime objects
            policy_exec_history["datetime"] = pd.to_datetime(
                policy_exec_history["datetime"], format="%Y-%m-%dT%H:%M:%SZ", utc=True
            )

            # Get current UTC time
            now = datetime.datetime.now(datetime.timezone.utc)

            # Calculate days ago
            policy_exec_history["days_ago"] = policy_exec_history["datetime"].apply(
                lambda dt: (now - dt).days
            )

            # Count total executions per hostname
            hostname_counts = policy_exec_history["hostname"].value_counts()

            # Map execution counts to agents
            agents["execution_count"] = agents["hostname"].map(hostname_counts).fillna(0).astype(int)

            # Find most recent execution per hostname
            most_recent_exec = policy_exec_history.sort_values(by="days_ago").drop_duplicates(
                subset="hostname", keep="first"
            )

            # Map most recent execution age to agents
            agents["days_since"] = agents["hostname"].map(
                most_recent_exec.set_index("hostname")["days_ago"]
            )

            # Check for enforcement readiness
            agents["required_quiet"] = required_quiet
            agents["enforce_ready"] = agents["days_since"].apply(
                lambda x: True if pd.isna(x) or x > required_quiet else False
            )

            # Sort agents by execution count and hostname
            agents = agents.sort_values(by=["execution_count", "hostname"], ascending=[True, True])

            # Save to CSV
            filename = f"{working_dir}\\{selected_policy[0].name}_agents_last_{history_days}_days.csv"
            logging.debug(f"Saving CSV to {filename}")
            print(colorText(f"Saving CSV to {filename}", "green"))
            agents.to_csv(filename, index=False)

            # Summary statistics
            total_agents = len(agents)
            ready_agents = agents["enforce_ready"].sum()
            not_ready_agents = total_agents - ready_agents
            ready_percentage = (ready_agents / total_agents) * 100

            # Print results
            
        
            message = (
                f"Total agents: {total_agents}\n"
                f"Agents marked as 'enforce_ready': {ready_agents}\n"
                f"Agents not ready: {not_ready_agents}\n"
                f"Percentage ready for enforcement: {ready_percentage:.2f}%"
            )
            logger.debug(message)
            colorText(message,"green")
    get_sanitized_input("Press enter to continue")