#    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 asyncio
import logging
import os
import os.path
from typing import List, Optional

import dotenv
import pandas as pd

from models.execution import ExecutionHistoryRecord, Hash
from models.policy import Allowlist, Policy
from services.API import AirlockAPIWrapper
from services.TaskQueue import AsyncTaskQueue, run_sync_task_in_thread
from utils.configmanager import get_protected_value, load_env, load_env_json
from utils.Selector import Selector
from utils.utils import (
    colorText,
    formatHTML,
    regulator,
)

logger = logging.getLogger(__name__)

dotenv.load_dotenv()

async def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:
    df = await api.policy_find_all()
    policies = [Policy(**row.to_dict()) for _, row in df.iterrows()]

    logger.debug("Prompting for Policies")
    print(colorText("Please select policy/policies", "white"))

    selected = await Selector.select_objects(policies, allow_multiple, prompt_each=True)

    if selected is None:
        return []

    logger.debug("Returning selected policies")
    return selected if isinstance(selected, list) else [selected]


async def selectAllowlists(api: AirlockAPIWrapper, policy="all", allow_multiple=True) -> List[Allowlist]:
    if policy == "all":
        df = await api.allowlist_find_all()
    else:
        df = await api.policy_list_allowlists(policy[0].groupid) # pyright: ignore[reportAttributeAccessIssue]

    allowlists = [Allowlist(**row.to_dict()) for _, row in df.iterrows()]

    logger.debug("Prompting for Allowlist(s)")
    print(colorText("Please select allowlist(s)", "white"))

    selected = await Selector.select_objects(allowlists, allow_multiple, prompt_each=True)

    if selected is None:
        return []

    logger.debug(f"Returning {selected}")
    return selected if isinstance(selected, list) else [selected]


async def sortHashes(
    api: AirlockAPIWrapper,
    queue: AsyncTaskQueue,
    selected_policies: List[Policy],
    type=[1, 2, 6, 7],
    history_days: Optional[int] = None
):
    if history_days is None:
        history_days = await Selector.select_value(
            prompt="Enter how many days of history to pull (1–150): ",
            value_type=int,
            valid_range=(1, 150),
        )

    logger.debug(f"{history_days} day selected for history")

    if history_days is None:
        logging.warning("No history range selected. Aborting.")
        return

    executions = []
    hashes = []
    working_dir = await load_env("WORKING_DIR")

    # Pull execution histories for each policy
    policy_executions = await ExecutionHistoryRecord.from_policies(
        api, selected_policies, type_=type, history_days=history_days
    )

    logger.debug(f"Policy_executions is {policy_executions}")
    executions.extend(policy_executions)
    logger.debug(f"Executions contains {executions}")

    if executions:
        sha_list = [record.sha256 for record in executions]
        hash_df = await api.hash_query(sha_list)
        hashes = [
            Hash(sha256=row["sha256"], **row["data"])
            for _, row in hash_df.iterrows()
        ]

    if hashes:
        unique_hashes = Hash.deduplicate(hashes)
        needs_review, approved, unapproved = await Hash.categorize_hashes(hashes=unique_hashes)

        categories = {
            "needs_review": needs_review,
            "approved": approved,
            "unapproved": unapproved,
        }

        for label, category in categories.items():
            csv_path = f"{working_dir}/Needs_Review/Review_First/{selected_policies[0].name}_{label}_executions.csv"
            html_path = f"{working_dir}/Needs_Review/HTML/{selected_policies[0].name}_{label}.html"

            df = await ExecutionHistoryRecord.enrich_with_hashes(executions, category)
            
            asyncio.create_task(queue.enqueue(
                f"DF TO CSV {selected_policies[0].name}_{label}",
                run_sync_task_in_thread,
                df.to_csv,
                csv_path,
                index=False,
                encoding='utf-8'
            ))


            asyncio.create_task(queue.enqueue(
                f"DF TO HTML {selected_policies[0].name}_{label}", 
                run_sync_task_in_thread, 
                formatHTML, 
                df, 
                html_path
            ))

    print("sortHashes completed successfully.")


async def buildPathsandPublishers(queue: AsyncTaskQueue, split):
    working_dir = await load_env("WORKING_DIR")
    path1 = f"{working_dir}/Approved/approved_executions.csv"
    path2 = f"{working_dir}/Approved/needs_review_executions.csv"

    df1 = await asyncio.to_thread(pd.read_csv, path1) if os.path.exists(path1) else pd.DataFrame()
    if df1.empty:
        logger.warning(f"File not found or empty: {path1}")

    df2 = await asyncio.to_thread(pd.read_csv, path2) if os.path.exists(path2) else pd.DataFrame()
    if df2.empty:
        logger.warning(f"File not found or empty: {path2}")

    if df1.empty and df2.empty:
        logger.warning("Both DataFrames are empty. Skipping sort.")
        all_approved_hashes = pd.DataFrame()
        logger.debug(all_approved_hashes.head())
        return

    all_approved_hashes = pd.concat([df1, df2], ignore_index=True)
    if "filename_exec" in all_approved_hashes.columns:
        all_approved_hashes = all_approved_hashes.sort_values(by="filename_exec")
    else:
        logger.warning("'filename_exec' column not found in concatenated DataFrame.")

    primary_path_exclusions = await calculatePath(all_approved_hashes, split)
    remaining_hashes = all_approved_hashes[~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"])]
    secondary_path_exclusions = await calculatePath(remaining_hashes, split)
    remaining_hashes = remaining_hashes[~remaining_hashes["sha256"].isin(secondary_path_exclusions["sha256"])]

    dataframes = {
        "primary_Paths": primary_path_exclusions,
        "secondary_Paths": secondary_path_exclusions,
        "hashes_to_add": remaining_hashes,
    }

    for name, df in dataframes.items():
        logger.debug(f"DataFrame headers for {name}: {list(df.columns)}")
        sort_column = "filename_exec" if name == "hashes_to_add" else "longestcfp"
        df.sort_values(by=sort_column, inplace=True)
        csv_path = f"{working_dir}/Needs_Review/Review_Second/{name}.csv"
        html_path = f"{working_dir}/Needs_Review/HTML/{name}.html"
        await asyncio.to_thread(df.to_csv, csv_path, index=False)
        await asyncio.to_thread(formatHTML, df, html_path)

    publist = all_approved_hashes[all_approved_hashes["publisher_hash"] != "Not Signed"].drop_duplicates(subset=["publisher_hash"])
    pattern = regulator(await load_env_json("BAD_PUBLISHERS", "[]"))
    publist = publist[~publist["publisher_hash"].str.contains(pattern, na=False)]
    publist = publist[["publisher_hash"]]
    publist.sort_values(by="publisher_hash", inplace=True)
    pub_csv_path = f"{working_dir}/Needs_Review/Review_Second/publishers.csv"
    await asyncio.to_thread(publist.to_csv, pub_csv_path, index=False)

    print("buildPathsandPublishers completed asynchronously.")


async def buildPreflights():
    working_dir = await load_env("WORKING_DIR")

    hash_path = f"{working_dir}/Approved/hashes_to_add.csv"
    path1 = f"{working_dir}/Approved/primary_Paths.csv"
    path2 = f"{working_dir}/Approved/secondary_Paths.csv"
    publishers_path = f"{working_dir}/Approved/publishers.csv"

    df1 = await asyncio.to_thread(pd.read_csv, path1) if os.path.exists(path1) else pd.DataFrame()
    if df1.empty:
        logger.warning(f"File not found or empty: {path1}")

    df2 = await asyncio.to_thread(pd.read_csv, path2) if os.path.exists(path2) else pd.DataFrame()
    if df2.empty:
        logger.warning(f"File not found or empty: {path2}")

    approved_hashes = await asyncio.to_thread(pd.read_csv, hash_path) if os.path.exists(hash_path) else pd.DataFrame()
    if approved_hashes.empty:
        logger.warning(f"File not found or empty: {hash_path}")

    approved_publishers = await asyncio.to_thread(pd.read_csv, publishers_path) if os.path.exists(publishers_path) else pd.DataFrame()
    if approved_publishers.empty:
        logger.warning(f"File not found or empty: {publishers_path}")

    approved_paths = pd.concat([df1, df2], ignore_index=True) if not (df1.empty and df2.empty) else pd.DataFrame()

    dataframes = {
        "approved_paths": approved_paths,
        "approved_hashes": approved_hashes,
        "approved_publishers": approved_publishers
    }

    for name, df in dataframes.items():
        logger.debug(f"DataFrame headers for {name}: {list(df.columns)}")
        if name == "approved_paths":
            df.sort_values(by="longestcfp", inplace=True)
        elif name == "approved_hashes":
            df.sort_values(by="filename_exec", inplace=True)
        elif name == "approved_publishers":
            df.sort_values(by="publisher_hash", inplace=True)

        csv_path = f"{working_dir}/Preflight/{name}.csv"
        html_path = f"{working_dir}/Preflight/HTML/{name}.html"

        await asyncio.to_thread(df.to_csv, csv_path, index=False)
        await asyncio.to_thread(formatHTML, df, html_path)

    print("buildPreflights completed asynchronously.")


async def splitFilepathsGrouped(df, col="filename"):
    path_task = asyncio.create_task(get_protected_value("PATH_EXCLUSION_CONST", int))
    min_files_task = asyncio.create_task(get_protected_value("MIN_FILES_FOR_PATH", int))

    path_exclusion_constant = await path_task
    min_files_for_path = await min_files_task
    
    def clean_split(path):
        if not isinstance(path, (str, bytes, os.PathLike)):
            return []
        parts = os.path.normpath(path).split(os.sep)
        return [p for p in parts if p]

    non_string_entries = df[~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike)))]
    if not non_string_entries.empty:
        logger.warning(f"Non-string entries found in column '{col}':")
        logger.debug(non_string_entries)

    df = df.copy()
    split_paths = df[col].apply(clean_split)
    df = df[split_paths.apply(lambda parts: len(parts) >= min_files_for_path)].copy()
    split_paths = split_paths[df.index]

    df["group_key"] = split_paths.apply(lambda parts: os.sep.join(parts[:path_exclusion_constant]))
    grouped = df.groupby("group_key")
    new_rows = []

    for _, group_df in grouped:
        paths = group_df[col].tolist()
        split_parts = [clean_split(p) for p in paths]

        def longest_common_prefix(paths):
            if not paths:
                return []
            prefix = paths[0]
            for path in paths[1:]:
                prefix = [a for a, b in zip(prefix, path) if a == b]
                if not prefix:
                    break
            return prefix

        common_prefix = longest_common_prefix(split_parts)
        prefix_str = os.sep.join(common_prefix)

        for i, parts in enumerate(split_parts):
            filename = parts[-1]
            middle = os.sep.join(parts[len(common_prefix):-1]) if len(parts) > len(common_prefix) + 1 else ""
            row = group_df.iloc[i].copy()
            row["longestcfp"] = prefix_str
            row["middle"] = middle
            row["filename_only"] = filename
            row["file_extension"] = os.path.splitext(filename)[1].lower()
            new_rows.append(row)

    return pd.DataFrame(new_rows).drop(columns=["group_key"])


async def calculatePath(approved_hashes, split):
    dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")] if split else [approved_hashes]
    badpathparts = await asyncio.to_thread(load_env_json, "BAD_PATH_PARTS", "[]")
    min_files_for_path = await asyncio.to_thread(get_protected_value, "MIN_FILES_FOR_PATH", int)

    processed_dfs = []

    for df in dfs_by_policy:
        haslcp = await splitFilepathsGrouped(df, "filename_exec")
        haslcp = haslcp.drop_duplicates()

        forbidden = regulator(badpathparts, True)
        forbidden_lcfp = haslcp["longestcfp"].str.contains(forbidden, na=False)

        logger.debug("Removing forbidden filepaths for path exceptions")
        print(colorText("Removing forbidden filepaths for path exceptions", "green"))

        lcp_not_forbidden = haslcp[~forbidden_lcfp].copy()

        lcp_not_forbidden_review = lcp_not_forbidden[[
            "policyname", "longestcfp", "middle", "filename_only", "file_extension", "sha256"
        ]]

        unique_sha_counts = (
            lcp_not_forbidden_review.groupby("longestcfp")["sha256"].nunique().reset_index()
        )
        unique_sha_counts.columns = ["longestcfp", "unique_sha256_count"]

        lcp_not_forbidden_review = lcp_not_forbidden_review.merge(
            unique_sha_counts, on="longestcfp", how="left"
        )
        lcp_not_forbidden_review = lcp_not_forbidden_review[
            lcp_not_forbidden_review["unique_sha256_count"] >= min_files_for_path
        ]

        processed_dfs.append(lcp_not_forbidden_review)

    pathExclusions = pd.concat(processed_dfs, ignore_index=True)
    return pathExclusions