#    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 os
import re
import time

import dotenv
import numpy as np
import pandas as pd

from models.agent import Agent
from services.agenthandler import findAllAgents, moveAgentToRelatedPolicy, selectAgents
from services.API import AirlockAPIWrapper
from utils.configmanager import get_protected_json, load_env, load_env_json
from utils.setup import get_base_directory
from utils.utils import colorText, get_sanitized_input

logger = logging.getLogger(__name__)

dotenv.load_dotenv()


def getLocalApprovals(api: AirlockAPIWrapper):
    base_dir = get_base_directory
    result = api.otp_find_awaiting()
    local_approval = pd.DataFrame(result["response"]["otpusage"])
    if os.path.exists(f"{base_dir}\\cache\\newest_local_approval.parquet"):
        previous_run = pd.read_parquet(f"{base_dir}\\cache\\newest_local_approval.parquet")
        previous_run.to_parquet(
            f"{base_dir}\\cache\\last_local_approval.parquet", index=False
        )
        os.remove(f"{base_dir}\\cache\\newest_local_approval.parquet")

    # Only keep rows presumably created by the generate local approval function
    local_approval = local_approval[
        local_approval["purpose"].str.startswith("🎫 Local Approval 🎫")
    ]

    local_approval["batchid"] = local_approval["purpose"].apply(
        lambda x: (match := re.search(r"batch:(\S+)", str(x))) and match.group(1)
    )

    if not local_approval.empty:
        local_approval.to_parquet(
            f"{base_dir}\\cache\\newest_local_approval.parquet", index=False
        )

    return local_approval


def scheduleAddingLAHashes(api: AirlockAPIWrapper):

    policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD","{}")
    bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]")
    pups = load_env_json("PUPS", "[]")
    threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE", cast_type = int)

    try:
        register_function("add_hash", returnFromLocalApproval)
        register_function("move_device", moveAgentToRelatedPolicy)
    except Exception as e:
        logger.warning(f"Failed to register functions: {e}")
        return

    try:
        approvals_df = getNewLocalApprovals(api)
        if approvals_df.empty:
            logger.debug("No new local approvals found. Nothing to schedule.")
            return
        batches = approvals_df.groupby("batchid")
    except Exception as e:
        logger.warning(f"Failed to retrieve or group local approvals: {e}")
        return

    for batchid, batch_df in batches:
        try:
            duration_minutes = int(batch_df["duration"].iloc[0])
            start_time = datetime.datetime.now()
            run_time = start_time + datetime.timedelta(minutes=duration_minutes)
            early_time = start_time + datetime.timedelta(minutes=np.floor(duration_minutes * 0.95))

            early_timestamp = early_time.timestamp()
            run_timestamp = run_time.timestamp()

            # Schedule add_hash job
            try:
                run_once_job(
                    f"add_hash_{batchid}",
                    "add_hash",
                    early_timestamp,
                    [
                        api,
                        batch_df,
                        policy_relationship_map,
                        bad_publisher_list,
                        pups,
                        threat_tolerance_constant,
                    ],
                    None,
                )
                logger.debug(f"Scheduled add_hash for batch {batchid} at {early_time}")
            except Exception:
                logger.debug("Failed to schedule add_hash for batch {batchid}: {e}")

            # Schedule move_device jobs
            devices = batch_df["agentid"].drop_duplicates().tolist()
            agents = []

            for device in devices:
                rows = api.agent_find_by_hostname(device).iterrows()
                agents += [Agent(**row["data"]) for _, row in rows]

            for agent in agents:
                try:
                    run_once_job(
                        f"move_device_{agent.hostame}_{batchid}",
                        "move_device",
                        run_timestamp,
                        [api, agent, policy_relationship_map],
                        "enforcement",
                    )

                    print(
                        f"Scheduled move_device for device {agent.hostname} in batch {batchid} at {run_time}"
                    )
                except Exception as e:
                    print(
                        f"Failed to schedule move_device for device {agent.hostname} in batch {batchid}: {e}"
                    )

        except Exception as e:
            logger.warning(f"Failed to process batch {batchid}: {e}")


def returnFromLocalApproval(api, device_df, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant
):
    """
      # Get unique policy names from device list
      policies_in_devicelist = sorted(device_df['policy_name'].unique().tolist())

      # Create inverse map to go from Audit to Enforcement
      inverse_map = {v: k for k, v in policy_relationship_map.items()}

      # Fetch all policies
      all_policies = [Policy(row['groupid'], row['hidden'], row['name'], row['parent']) for _, row in api.policy_find_all().iterrows()]

      # Define policy types
      policy_types = [1, 2, 6, 7]

    #TODO finish logic for adding hashes
    """
    working_dir = load_env("WORKING_DIR")
    policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD","{}")
    bad_publisher_list = load_env_json("BAD_PUBLISHER", "[]")
    pups = load_env_json("PUPS", "[]")
    threat_tolerance_constant = load_env("VT_THREAT_TOLERANCE")
    print(f"{working_dir}, {policy_relationship_map}, {bad_publisher_list}, {pups}, {threat_tolerance_constant}")

def moveToLocalApproval(api: AirlockAPIWrapper):
    possible_durations = [15, 60, 360, 1440, 10080]
    duration_selected = None

    print(colorText("Please select a duration:", "white"))
    for i, option in enumerate(possible_durations, start=1):
        print(f"{i}. {option}")

    try:

        choice = int(get_sanitized_input("Enter the number of your choice:"))
        if 1 <= choice <= len(possible_durations):
            duration_selected = possible_durations[choice - 1]
            print(colorText(f"You selected: {duration_selected}", "yellow"))
            logger.debug(f"You selected: {duration_selected}")
        else:
            print(colorText("❌ Invalid choice.", "red"))
            logger.debug("Invalid Input")
            return
    except ValueError:
        print(colorText("❌ Invalid input. Please enter a number.", "red"))
        logger.debug("Invalid Input")
        return

    agents = selectAgents(api)
    batch = int(time.time())

    if not agents:
        print(colorText("❌ No agents found or error retrieving agents.", "red"))
        logger.debug("No agents found or error retrieving agents")
        return

    for agent in agents:
        try:
            addLocalApproval(api, batch, duration_selected, agent.agentid)
            moveAgentToRelatedPolicy(api, agent, "audit")
        except Exception as e:
            print(colorText(f"❌ Error processing agent {agent.hostname}: {e}", "red"))


def addLocalApproval(api: AirlockAPIWrapper, batchid, duration_selected, agentid):
    
    purpose = f"🎫 Local Approval 🎫 - {duration_selected} mins - batch:{batchid} Client:{agentid}"
    api.otp_generate(agentid, duration_selected, purpose)
 


def monitorAuditStatus(api: AirlockAPIWrapper):
    current_agents = findAllAgents(api)
    last_agents = []
    if not last_agents:
        last_agents = current_agents
    policy_relationship_map = get_protected_json("POLICY_MAP_ENF_AUD","{}")

    # Reverse map for audit → enforcement
    reverse_policy_map = {v: k for k, v in policy_relationship_map.items()}
    known_transitions = set(policy_relationship_map.items()) | set(reverse_policy_map.items())

    # Index last_agents by hostname for quick lookup
    last_agent_map = {agent.hostname: agent for agent in last_agents}

    # Result buckets
    newly_added = []
    same_policy = []
    moved_to_audit = []
    moved_to_enforcement = []
    unusual_move = []

    for current in current_agents:
        previous = last_agent_map.get(current.hostname)

        if not previous:
            newly_added.append(current)
            continue

        if current.groupid == previous.groupid:
            same_policy.append(current)
        elif (previous.groupid, current.groupid) in known_transitions:
            moved_to_audit.append(current)
        elif (current.groupid, previous.groupid) in known_transitions:
            moved_to_enforcement.append(current)
        else:
            unusual_move.append(current)

    # Return all five DataFrames
    return newly_added, same_policy, moved_to_audit, moved_to_enforcement, unusual_move


def getNewLocalApprovals(api: AirlockAPIWrapper):
    
    working_dir = load_env("WORKING_DIR")                   
    current_la = getLocalApprovals(api)

    # Load old approval list
    old_la_path = f"{working_dir}\\Scheduling\\last_local_approval.parquet"
    if os.path.exists(old_la_path):
        old_la = pd.read_parquet(old_la_path)
    else:
        old_la = pd.DataFrame(columns=current_la.columns)

    # Create composite keys
    current_la["key"] = current_la["clientid"].astype(str) + "_" + current_la["granted"].astype(str)
    old_la["key"] = old_la["clientid"].astype(str) + "_" + old_la["granted"].astype(str)

    # Find new entries
    new_entries = current_la[~current_la["key"].isin(old_la["key"])]

    # Convert 'granted' to datetime and filter by last 10 minutes
    new_entries["granted"] = pd.to_datetime(new_entries["granted"], utc=True, errors="coerce")
    ten_minutes_ago = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(minutes=10)
    recent_entries = new_entries[new_entries["granted"] > ten_minutes_ago]

    # Save current approvals for next run
    current_la.drop(columns=["key"], inplace=True)
    current_la.to_parquet(old_la_path, index=False)

    return recent_entries