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

from bson import ObjectId
import pandas as pd
import tqdm

from models.policy import Policy
from services.API import AirlockAPIWrapper
from utils.configmanager import get_system_json
from utils.setup import get_base_directory
from utils.utils import areYouSure, colorText, get_sanitized_input

logger = logging.getLogger(__name__)


def pullPolicyExechistories(
    api: AirlockAPIWrapper,
    policy: Policy,
    type: list,
    days,
    outputjson: bool,
):

    file_path = f"{get_base_directory()}\\cache\\chunkinator.json"

    # Ensure the file exists
    if not os.path.exists(file_path):
        with open(file_path, "w") as file:
            json.dump({"error": "Success", "response": {"exechistories": []}}, file)
        logger.debug(f"File '{file_path}' has been created.")
    else:
        logger.debug(f"File '{file_path}' already exists.")

    checkpoint = str(skipback(days))
    json_output = {"error": "Success", "response": {"exechistories": []}}

    with tqdm.tqdm(
        file=sys.stdout,
        leave=True,
        total=10000,
        desc=f"Checkpoint Progress: {checkpoint}",
        colour="blue",
        initial=1,
    ) as filebar:
        with tqdm.tqdm(
            file=sys.stdout,
            leave=True,
            total=100,
            desc=f"Total of {policy} Complete: ",
        ) as pbar:
            while True:
                histories = api.history_logging(
                    type=type, checkpoint=checkpoint, policy=[policy.name]
                )

                # Ensure histories is a list of dictionaries
                if not isinstance(histories, list) or not all(
                    isinstance(h, dict) for h in histories
                ):
                    logger.error(
                        "Unexpected response format from API. Expected list of dictionaries."
                    )
                    break

                filebar.total = len(histories)

                if not histories:
                    break

                for index, history_item in enumerate(histories):
                    if (
                        "checkpoint" not in history_item
                        or "datetime" not in history_item
                    ):
                        continue  # Skip malformed entries

                    # Update checkpoint on last item
                    if index == len(histories) - 1:
                        checkpoint = history_item[
                            "checkpoint"
                        ]  # pyright: ignore[reportArgumentType]
                        filebar.desc = f"Checkpoint Progress: {checkpoint}"
                        break

                    try:
                        history_date = datetime.datetime.strptime(
                            history_item["datetime"].replace(
                                " +0000 UTC", ""
                            ),  # pyright: ignore[reportArgumentType]
                            "%Y-%m-%dT%H:%M:%SZ",
                        ).date()
                    except ValueError:
                        continue  # Skip if date format is invalid

                    if (
                        datetime.date.today() - datetime.timedelta(days=days)
                    ) <= history_date:
                        json_output["response"]["exechistories"].append(history_item)

                    filebar.update(1)
                    filebar.refresh()

                # Deduplicate entries
                seen = {}
                if os.path.exists(file_path):
                    with open(file_path, "r") as file:
                        existing_data = json.load(file)
                    combined = (
                        existing_data["response"]["exechistories"]
                        + json_output["response"]["exechistories"]
                    )
                else:
                    combined = json_output["response"]["exechistories"]

                for entry in combined:
                    key = (
                        entry.get("sha256"),
                        entry.get("filename"),
                        entry.get("hostname"),
                    )
                    seen[key] = entry

                deduplicated = list(seen.values())
                with open(file_path, "w") as file:
                    json.dump(
                        {
                            "error": "Success",
                            "response": {"exechistories": deduplicated},
                        },
                        file,
                    )

                json_output["response"]["exechistories"].clear()

                # Update progress bar based on last valid item
                try:
                    last_date = datetime.datetime.strptime(
                        history_item["datetime"].replace(" +0000 UTC", ""),  # type: ignore
                        "%Y-%m-%dT%H:%M:%SZ",
                    ).date()
                    date_diff = datetime.date.today() - last_date
                    percentage_diff = (
                        ((days + 10) - date_diff.days) / (days + 10)
                    ) * 100
                    pbar.n = round(percentage_diff)
                    pbar.set_description_str(f"Total of {policy} Complete: ")
                    pbar.refresh()
                except Exception:
                    pass

                filebar.n = 1

    # Final output
    with open(file_path, "r") as file:
        final_output = json.load(file)
    os.remove(file_path)

    return json.dumps(final_output) if outputjson else None


def getPolicyInfo(api: AirlockAPIWrapper, policy, type, days):
    import airlock_libs

    executionhist_policy = pd.DataFrame()
    exehist = airlock_libs.pull_policy_exec_histories(api, policy.name, str(type), days)
    if exehist is not None:
        data = json.loads(exehist)
        executionhist_policy = pd.DataFrame(data["response"]["exechistories"])
        if not executionhist_policy.empty:
            executionhist_policy = executionhist_policy[
                [
                    "datetime",
                    "sha256",
                    "publisher",
                    "filename",
                    "hostname",
                    "username",
                    "pprocess",
                    "gprocess",
                    "commandline",
                ]
            ]
            executionhist_policy["policy"] = policy  # Add policy column here
            executionhist_policy = executionhist_policy.drop_duplicates(
                subset=["sha256", "filename", "hostname"]
            )
            executionhist_policy = executionhist_policy.sort_values(
                by=["sha256", "filename"]
            )
        logger.debug(f"Staging of Execution history for policy: {policy} is complete")
        print(
            colorText(
                f"Staging of Execution history for policy: {policy} is complete",
                "green",
            )
        )
        del data
        del exehist
        gc.collect()
    return executionhist_policy


def skipback(days):
    """
    Generate a MongoDB ObjectId for a given number of days ago from today.
    """
    adjusted_days = days
    date_days_ago = datetime.datetime.now(datetime.UTC) - datetime.timedelta(
        days=adjusted_days
    )
    timestamp = int(date_days_ago.timestamp())
    hex_timestamp = format(timestamp, "08x")
    objectid_hex = hex_timestamp + "0000000000000000"
    return ObjectId(objectid_hex)


def updateAuditPoliciesFromEnforcementPolices(api: AirlockAPIWrapper):
    policy_relationship_map = get_system_json("POLICY_MAP_ENF_AUD", "{}")
    for enforcement_policy, audit_policy in policy_relationship_map.items():
        api.policy_clone(enforcement_policy, audit_policy)
        api.policy_set_auditmode(audit_policy, "1")


def confirmUpdateAfromE(api: AirlockAPIWrapper):
    areYouSure()
    confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
    if confirmation.strip() == "I AGREE":
        updateAuditPoliciesFromEnforcementPolices(api)
