#    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 utils.setup import get_base_directory
from models.agent import Agent
from services.agenthandler import findAllAgents, moveAgentToRelatedPolicy, selectAgents
from services.API import AirlockAPIWrapper
from Server.scheduler import (
    register_function,
    run_once_job,
)
from utils.utils import colorText, load_env, load_env_json

logger = logging.getLogger(__name__)

dotenv.load_dotenv()


def getLocalApprovals(api: AirlockAPIWrapper):
    working_dir = load_env("WORKING_DIR")
    result = api.otp_find_awaiting()
    local_approval = pd.DataFrame(result["response"]["otpusage"])
    if os.path.exists(f"{working_dir}\\Scheduling\\newest_local_approval.parquet"):
        previous_run = pd.read_parquet(f"{working_dir}\\Scheduling\\newest_local_approval.parquet")
        previous_run.to_parquet(
            f"{working_dir}\\Scheduling\\last_local_approval.parquet", index=False
        )
        os.remove(f"{working_dir}\\Scheduling\\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"{working_dir}\\Scheduling\\newest_local_approval.parquet", index=False
        )

    return local_approval


def scheduleAddingLAHashes(api: AirlockAPIWrapper):

    policy_relationship_map = load_env_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 = load_env_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(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 = load_env_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