#    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 logging
import os

import dotenv

import services.policyhandler as policyh
from flows.otp import generate, otp_activities_by_agent, revoke
from flows.prepPolicy import (
    buildPathsandPublishers,
    buildPreflights,
    selectAllowlists,
    selectPolicies,
    sortHashes,
    testChange,
)
from flows.quietAgent import findQuietAgents
from services.agenthandler import (
    findAgents,
    moveAgents,
    toggleEnforcement,
)
from services.API import AirlockAPIWrapper
from utils.configmanager import load_env
from utils.utils import (
    areYouSure,
    clear_screen,
    colorText,
    displayIntro,
    get_sanitized_input,
    locked,
    open_directory,
    printEnforceChecklist,
    welcome,
)

logger = logging.getLogger(__name__)

dotenv.load_dotenv()

def menu_main(api: AirlockAPIWrapper):
    working_dir = load_env("WORKING_DIR")
    extras = load_env("EXTRAS")
    while True:
        clear_screen()
        displayIntro()
        welcome()
        print(colorText("1. 🔍  - Device Search", "yellow"))
        print(colorText("2. 🔀  - Move Device(s)", "yellow"))
        print(colorText("3. 🎫  - One Time Pass (OTP)", "yellow"))
        print(colorText("4. 🔇  - Find Quiet Hosts", "yellow"))
        if extras == "POLICYPREP" : print(colorText("5. 🛡️   - Policy Enforcement Tools", "yellow"))
        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))
        print(colorText("F. 📂  - Open Working Directory", "yellow"))
        print(colorText("S. 🛠️  - Settings", "yellow"))
        print(colorText("Q. 🔚  - Quit", "yellow"))

        choice = get_sanitized_input("\nEnter Menu Item: ")
        if choice == "1":
             findAgents(api,False)
        elif choice == "2":
            menu_move(api)
        elif choice == "3":
            menu_otp(api)
        elif choice == "4":
           findQuietAgents(api)
        elif choice == "5":
            if extras == "POLICYPREP": menu_policymanagment(api)
        elif choice.upper() == "F":
            open_directory(working_dir)
        elif choice.upper() == "S":
            menu_settings()
        elif choice.upper() == "Q":
            break
        else:
            print(colorText("Invalid choice. Please try again.", "red"))

def menu_policy_enforce(api: AirlockAPIWrapper): #TODO Need to clean up 6 and 7 into functions
    selected_policies = []
    destination_policy = []
    destination_allowlist = []
    processed_paths = []
    processed_hashes = []
    processed_publishers = []
    working_dir = load_env("WORKING_DIR")

    while True:
        printEnforceChecklist(selected_policies, destination_policy, destination_allowlist)
        choice = get_sanitized_input("\nEnter your choice: ")

        if choice == "1":
            selected_policies = selectPolicies(api,True)

        elif choice == "2":
            print(colorText("Please choose destination_name Policy for Path Exclusions", "white"))

            destination_policy = selectPolicies(api, False)

            print(colorText("Please choose Allowlist for Hashes", "white"))

            destination_allowlist = selectAllowlists(api, destination_policy, False)

        elif choice == "3":
            sortHashes(
                api,
                selected_policies,
                type=[1, 2, 6, 7],
            )

        elif choice == "4":
            if os.path.exists(f"{working_dir}\\Needs_Review\\Review_First\\{selected_policies[0].name}_approved_executions.csv"):
                buildPathsandPublishers(selected_policies, False)
            else:
                print("File not found. Please make sure it's saved correctly and try again.")

        elif choice == "5":
            if os.path.exists(f"{working_dir}\\Approved\\{selected_policies[0].name}_approved_executions.csv") and os.path.exists(
                f"{working_dir}\\Approved\\{selected_policies[0].name}_primary_Paths.csv"
            ):
                buildPreflights(selected_policies)
            else:
                print("File not found. Please make sure it's saved correctly and try again.")

        elif choice == "6":
            if (
                os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv")
                and os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv")
                and destination_policy
                and destination_allowlist
            ):
                processed_paths, processed_hashes, processed_publishers = testChange(selected_policies, destination_policy, destination_allowlist)
            else:
                # Log which condition(s) failed
                missing_items = []
                if not os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_paths.csv"):
                    missing_items.append("approved_paths.csv not found")
                if not os.path.exists(f"{working_dir}\\Preflight\\{selected_policies[0].name}_approved_hashes.csv"):
                    missing_items.append("approved_hashes.csv not found")
                if not destination_policy:
                    missing_items.append("destination_policy is empty or None")
                if not destination_allowlist:
                    missing_items.append("destination_allowlist is empty or None")

                logger.error("Preflight check failed due to the following:")
                for item in missing_items:
                    logger.error(f" - {item}")

        elif choice == "7":
            areYouSure()
            confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
            if (
                processed_paths
                and processed_hashes
                and processed_publishers
                and destination_policy
                and destination_allowlist
                and confirmation.strip() == "I AGREE"
            ):
                print(colorText("Proceeding with the code...", "yellow"))
                api.hash_add_to_allowlist(destination_allowlist[0].applicationid, processed_hashes)
                api.policy_add_path_exclusions(destination_policy[0].groupid, processed_paths)
                if processed_publishers:
                    api.policy_add_publishers(destination_policy[0].groupid, processed_publishers)

                locked()
                
            else:
                logger.error("Confirmation block failed. Reasons:")
                if not processed_publishers or processed_hashes or processed_paths:
                    logger.error(" - Test not performed.")
                if not destination_policy:
                    logger.error(" - `destination_policy` is missing or invalid.")
                if not destination_allowlist:
                    logger.error(" - `destination_allowlist` is missing or invalid.")
                if confirmation.strip() != "I AGREE":
                    logger.error(" - User did not confirm with 'I AGREE'. Received: '%s'", confirmation.strip())

        elif choice.upper() == "F":
            open_directory(working_dir)
        elif choice.upper() == "S":
            menu_settings()
        elif choice.upper() == "B":
            break

       
        else:
            print(colorText("Invalid choice. Please try again.", "red"))

def menu_move(api: AirlockAPIWrapper):
    working_dir = load_env("WORKING_DIR")
    while True:
        clear_screen()
        displayIntro()
        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= ↔️  Agent Movement ↔️  =-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))
        print(colorText("1. ✅  - Move to local approval (Placeholder)", "yellow"))
        print(colorText("2. 🔄  - Move to Audit/Enforcement", "yellow"))
        print(colorText("3. 🔀  - Move - Other", "yellow"))
        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))
        print(colorText("F. 📂  - Open Working Directory", "yellow"))
        print(colorText("S. 🛠️  - Settings", "yellow"))
        print(colorText("B. 🔙  - Back", "yellow"))

        choice = get_sanitized_input("Enter your choice: ")

        if choice == "1":
            print("This Feature is still in development")
            get_sanitized_input("Press enter to continue")
        elif choice == "2":
            toggleEnforcement(api)
        elif choice == "3":
            moveAgents(api)
        elif choice.upper() == "F":
            open_directory(working_dir)
        elif choice.upper() == "S":
            menu_settings()
        elif choice.upper() == "B":
            break
    

def menu_otp(api: AirlockAPIWrapper):
    working_dir = load_env("WORKING_DIR")
    while True:
        clear_screen()
        displayIntro()
        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= 🎫  OTP 🎫 =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))
        print(colorText("1. 🔐  -Generate OTPs", "cyan"))
        print(colorText("2. 📊  -OTP Activities By Agent", "cyan"))
        print(colorText("3. ❌  -Revoke OTPs", "cyan"))
        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))
        print(colorText("F. 📂  - Open Working Directory", "yellow"))
        print(colorText("S. 🛠️  - Settings", "yellow"))
        print(colorText("B. 🔙  - Back", "yellow"))

        choice = get_sanitized_input("Enter your choice: ")

        if choice == "1":
            otp_list = generate(api)
            print(colorText("Requested Codes:", "green"))
            for key, value in otp_list.items():
                print(colorText(f"{key} | {value}","green"))
        elif choice == "2":
            otp_activities_by_agent(api)
        elif choice == "3":
            revoke(api)
        elif choice.upper() == "F":
            open_directory(working_dir)
        elif choice.upper() == "S":
            menu_settings()
        elif choice.upper() == "B":
            break
   
def menu_policymanagment(api: AirlockAPIWrapper):
    working_dir = load_env("WORKING_DIR")
    while True:
        clear_screen()
        displayIntro()
        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-= 🛡️  Policy Tools 🛡️  =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))       
        print(colorText("1. 🔒  - Prepare Policy For Enforcement", "yellow"))
        print(colorText("2. 🔄  - Update Audit Policies from Enforcement Policies", "yellow"))
        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))
        print(colorText("F. 📂  - Open Working Directory", "yellow"))
        print(colorText("S. 🛠️  - Settings", "yellow"))
        print(colorText("B. 🔙  - Back", "yellow"))
        choice = get_sanitized_input("\n Enter Menu Item: ")

        if choice == "1":
            menu_policy_enforce(api)
        elif choice == "2":
            areYouSure()
            confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
            if confirmation.strip() == "I AGREE":
                policyh.updateAuditPoliciesFromEnforcementPolices(api)

        elif choice.upper() == "F":
            open_directory(working_dir)
        elif choice.upper() == "S":
            menu_settings()
        elif choice.upper() == "B":
            break
        else:
            print(colorText("Invalid choice. Please try again.", "red"))
def menu_settings():
    while True:
        clear_screen()
        displayIntro()
        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= 🛠️  Settings 🛠️  =-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))
        print(colorText("This Feature is still in development, if you have ideas for options you would like to see, let us know.", "cyan"))
        print(colorText("=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=", "cyan"))

        # print(colorText("2. Sub-option B","cyan"))
        print(colorText("B. 🔙  - Back", "yellow"))
        choice = get_sanitized_input("Enter your choice: ")

        if choice == "1":
            pass #TODO ADD CHANGE WORKDIR CODE

        elif choice.upper() == "B":
            print("Returning to Main Menu...")
            break
        else:
            print("Invalid choice. Please try again.")