#    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.configmanager import load_env
from utils.selector import Selector
from utils.utils import colorText, get_sanitized_input

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")