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

#Local Imports
import utils.clientfunctions as clientf
import utils.pathfunctions as pathf
import utils.policyfunctions as policyf
import utils.utils as ct
from utils.perstscheduler import register_function, run_once_job, recurring_job, reload_jobs, start_scheduler, find_and_prioritize_jobs_by_pid

#Standard Libary Imports:
import json
import os
import re
import time
import datetime

#3rd Party Imports:
import pandas as pd
import numpy as np
import requests

def getLocalApprovals(url):

    endpoint = url + f'/v1/otp/usage'
    payload = {
            "status" : "0"
    }

    headers = {"X-APIKey": os.getenv('APIKEY')}
    payload = json.dumps(payload)

    response = requests.post(endpoint, headers=headers, data=payload, verify=False)
    result = json.loads(response.text)
    local_approval = pd.DataFrame(result["response"]["otpusage"])
    if os.path.exists("Local_Approval\\PARQ\\newest_local_approval.parquet"):
        previous_run = pd.read_parquet("Local_Approval\\PARQ\\newest_local_approval.parquet")
        previous_run.to_parquet("Local_Approval\\PARQ\\last_local_approval.parquet", index=False)
        os.remove("Local_Approval\\PARQ\\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:
        ct.style_dataframe_dark(local_approval, f"Local_Approval\\HTML\\newest_local_approval.html")
        local_approval.to_parquet("Local_Approval\\PARQ\\newest_local_approval.parquet", index=False)
    
    return local_approval
   
def scheduleAddingLAHashes(url, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant):
    try:
        register_function("add_hash", returnFromLocalApproval)
        register_function("move_device", clientf.moveAgentToEnforcement)
    except Exception as e:
        print(f"[ERROR] Failed to register functions: {e}")
        return

    try:
        approvals_df = getNewLocalApprovals(url)
        if approvals_df.empty:
            print("[INFO] No new local approvals found. Nothing to schedule.")
            return
        batches = approvals_df.groupby('batchid')
    except Exception as e:
        print(f"[ERROR] 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,
                    [url, batch_df, policy_relationship_map, bad_publisher_list, pups, threat_tolerance_constant],
                    None
                )
                print(f"[INFO] Scheduled add_hash for batch {batchid} at {early_time}")
            except Exception as e:
                print(f"[ERROR] Failed to schedule add_hash for batch {batchid}: {e}")

            # Schedule move_device jobs
            devices = batch_df['agentid'].drop_duplicates().tolist()
            for device in devices:
                try:
                    run_once_job(
                        f"move_device_{device}_{batchid}",
                        "move_device",
                        run_timestamp,
                        [url, device, policy_relationship_map],
                        None
                    )
                    print(f"[INFO] Scheduled move_device for device {device} in batch {batchid} at {run_time}")
                except Exception as e:
                    print(f"[ERROR] Failed to schedule move_device for device {device} in batch {batchid}: {e}")

        except Exception as e:
            print(f"[ERROR] Failed to process batch {batchid}: {e}")
            
def returnFromLocalApproval(url, 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 = policyf.getPolicyDataframe(url)

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

    # Define output directories
    parq_base_dir = "Local_Approval\\PARQ\\"
    needappr_base_dir = "Local_Approval\\PARQ\\"

    # Ensure output directories exist
    os.makedirs(parq_base_dir, exist_ok=True)
    os.makedirs(needappr_base_dir, exist_ok=True)

    # Build execution history
    policyf.buildExecHistory(
        url,
        policies_in_devicelist,
        parq_base_dir,
        needappr_base_dir,
        policy_types,
        threat_tolerance_constant,
        bad_publisher_list,
        pups,
    )

    # Read hash data
    unknown_hashes = ct.tryToReadCSV(f"{needappr_base_dir}unknown_hashes.csv")
    good_hashes = ct.tryToReadCSV(f"{needappr_base_dir}good_hashes.csv")
    hashes = pd.concat([unknown_hashes, good_hashes], ignore_index=True)

    # Process each policy
    for policy in policies_in_devicelist:
        # Filter for matching policy
        matching_rows = all_policies[all_policies['name'] == policy]
        
        if matching_rows.empty:
            print(f"Warning: No group ID found for policy '{policy}'. Skipping.")
            continue

        # Extract group ID
        policy_id = matching_rows['groupid'].values[0]

        # Map to destination ID
        destination_id = inverse_map.get(policy_id)
        if destination_id is None:
            print(f"Warning: No corresponding enforcement policy found for group ID '{policy_id}'. Skipping.")
            continue

        allowlist = policyf.getDestAllowlist(url, destination_id)
        
        policyf.addHash(url, allowlist, hashes[hashes['group'] == policy_id])
        
def moveToLocalApproval(url, policy_relationship_map):

    possible_durations = [15, 60, 360, 1440, 10080]
    duration_selected = None

    print(ct.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(ct.colorText(f"You selected: {duration_selected}", "yellow"))
        else:
            print(ct.colorText("❌ Invalid choice.", "red"))
            return
    except ValueError:
        print(ct.colorText("❌ Invalid input. Please enter a number.", "red"))
        return

    devicelist = clientf.promptForDevices()
    device_df = clientf.findAgents(url, devicelist, True)

    batch = int(time.time())

    if device_df is None or device_df.empty:
        print(ct.colorText("❌ No agents found or error retrieving agents.", "red"))
        return

    for row in device_df.itertuples(index=False):
        try:
            addLocalApproval(url, batch, duration_selected, row.agentid)
            clientf.moveAgentToAudit(url, row.agentid, policy_relationship_map)
        except Exception as e:
            print(ct.colorText(f"❌ Error processing agent {row.agentid}: {e}", "red"))

def addLocalApproval(url, batchid, duration_selected, agentid):
    purpose = f"🎫 Local Approval 🎫 - {duration_selected} mins - batch:{batchid} Client:{agentid}"
    endpoint = url + '/v1/otp/retrieve'
    payload = {
        "duration": str(duration_selected),
        "agentid": str(agentid),
        "purpose": purpose
    }

    headers = {"X-APIKey": os.getenv('APIKEY')}
    response = requests.post(endpoint, headers=headers, data=json.dumps(payload), verify=False)

    try:
        result = response.json()
        otpcode = result["response"]["otpcode"]
        print(ct.colorText(f"The OTP code is: {otpcode}", "yellow"))
    except Exception as e: 
        print(ct.colorText(f"An unexpected error occurred for agent {agentid}: {str(e)}", "red"))
        
def monitorAuditStatus(url: str, policy_relationship_map: dict):
    
    # Simulated current agent list
    current_agent_list = clientf.findAllAgents(url)

    # Load old agent list
    old_agent_path = "Local_Approval\\PARQ\\last_agent_list.parquet"
    if os.path.exists(old_agent_path):
        old_agent_list = pd.read_parquet(old_agent_path)
    else:
        old_agent_list = pd.DataFrame(columns=current_agent_list.columns)

    # Merge on hostname
    merged = pd.merge(
        old_agent_list[['hostname', 'groupid']],
        current_agent_list[['hostname', 'groupid']],
        on='hostname',
        how='outer',
        suffixes=('_old', '_current'),
        indicator=True
    )

    # Reverse map for 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())

    # 1. Newly added
    newly_added = merged[merged['_merge'] == 'right_only']

    # 2. Same policy
    same_policy = merged[
        (merged['_merge'] == 'both') &
        (merged['groupid_old'] == merged['groupid_current'])
    ]

    # 3. Moved to audit
    moved_to_audit = merged[
        (merged['_merge'] == 'both') &
        (merged['groupid_old'].isin(policy_relationship_map)) &
        (merged['groupid_current'] == merged['groupid_old'].map(policy_relationship_map))
    ]

    # 4. Moved to enforcement
    moved_to_enforcement = merged[
        (merged['_merge'] == 'both') &
        (merged['groupid_old'].isin(reverse_policy_map)) &
        (merged['groupid_current'] == merged['groupid_old'].map(reverse_policy_map))
    ]

    # 5. Unusual moves
    unusual_move = merged[
        (merged['_merge'] == 'both') &
        (merged['groupid_old'] != merged['groupid_current']) &
        merged.apply(lambda row: (row['groupid_old'], row['groupid_current']) not in known_transitions, axis=1)
    ]

    current_agent_list.to_parquet("Local_Approval\\PARQ\\last_agent_list.parquet", index=False)

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

def getNewLocalApprovals(url):
    current_la = getLocalApprovals(url)

    # Load old approval list
    old_la_path = "Local_Approval\\PARQ\\last_la.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