#    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 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 utils.configmanager import get_protected_value, load_env, load_env_json
from utils.selector import Selector
from utils.utils import (
    colorText,
    formatHTML,
    import_to_dataframe,
    regulator,
)

logger = logging.getLogger(__name__)

dotenv.load_dotenv()



def selectPolicies(api: AirlockAPIWrapper, allow_multiple=True) -> List[Policy]:

    policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
    logger.debug("Prompting for Policies")
    print(colorText("Please select policy/policies", "white"))
    selected = Selector.select_objects(policies, allow_multiple, prompt_each=True)

    if selected is None:
        return []

    # Normalize to always return a list
    logger.debug("Returning {selected.dict}")
    return selected if isinstance(selected, list) else [selected]


def selectAllowlists(api: AirlockAPIWrapper, policy = all, allow_multiple=True) -> List[Allowlist]:
    if policy == "all": allowlists = [Allowlist(**row.to_dict()) for _, row in api.allowlist_find_all().iterrows()]
    else:  allowlists = [Allowlist(**row.to_dict()) for _, row in api.policy_list_allowlists(policy[0].groupid).iterrows()]
    logger.debug("Prompting for Allowlist(s)")
    print(colorText("Please select allowlist(s)", "white"))
    selected = Selector.select_objects(allowlists, allow_multiple, prompt_each=True)

    if selected is None:
        return []

    # Normalize to always return a list
    logger.debug(f"Returning {selected}")
    return selected if isinstance(selected, list) else [selected]



def sortHashes(
    api: AirlockAPIWrapper,
    selected_policies: List[Policy],
    type=[1, 2, 6, 7],
    history_days: Optional[int] = None
):

 
    if history_days is None:
        history_days = 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 = load_env("WORKING_DIR")
    # Pull execution histories for each policy
 
    policy_executions = 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:
        hashes = [Hash(sha256=row["sha256"], **row["data"]) for _, row in api.hash_query([record.sha256 for record in executions]).iterrows()
        ]

    if hashes:
        unique_hashes = Hash.deduplicate(hashes)

        needs_review, approved, unapproved = 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\\{label}.html"

            ExecutionHistoryRecord.enrich_with_hashes_and_export(
                executions, category, f"{working_dir}\\Needs_Review\\Review_First", label=label
            )
            df = import_to_dataframe(csv_path)
            formatHTML(df, html_path)

def buildPathsandPublishers(split):
    working_dir = load_env("WORKING_DIR")
    df1 = pd.DataFrame()
    df2 = pd.DataFrame()
    all_approved_hashes = pd.DataFrame()
    path1 = f"{working_dir}\\Approved\\approved_executions.csv"
    path2 = f"{working_dir}\\Approved\\needs_review_executions.csv"

    if os.path.exists(path1):
        df1 = pd.read_csv(path1)
    else:
        logger.warning(f"File not found: {path1}")

    if os.path.exists(path2):
        df2 = pd.read_csv(path2)
    else:
        logger.warning(f"File not found: {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)
    else:
        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("Warning: 'filename_exec' column not found in concatenated DataFrame.")

    if not all_approved_hashes.empty:
        primary_path_exclusions = calculatePath(
            all_approved_hashes,
            split,
        )
        remaining_hashes = all_approved_hashes[
            ~all_approved_hashes["sha256"].isin(primary_path_exclusions["sha256"])
        ]
        secondary_path_exclusions = 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,
        }
        logger.debug("Preparing to sort dataframes")
        for name, df in dataframes.items():
            logger.debug(f" DataFrame headers: {list(df.columns)}")
            if name == "hashes_to_add":   df.sort_values(by="filename_exec", inplace=True)
            else: df.sort_values(by="longestcfp", inplace=True)
          
            df.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\{name}.csv", index=False)
            formatHTML(df, f"{working_dir}\\Needs_Review\\HTML\\{name}.html")

        if not all_approved_hashes.empty:
            # Drop all not signed, only keep unique values
            publist = all_approved_hashes[
                all_approved_hashes["publisher_hash"] != "Not Signed"
            ].drop_duplicates(subset=["publisher_hash"])
            # Remove Bad publisher if somehow they made it this far
            pattern = regulator(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)
            publist.to_csv(f"{working_dir}\\Needs_Review\\Review_Second\\publishers.csv", index=False)

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

    df1 = pd.DataFrame()
    df2 = pd.DataFrame()
    approved_hashes = pd.DataFrame()
    approved_publishers = pd.DataFrame()

    hash = 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 = f"{working_dir}\\Approved\\publishers.csv"

    if os.path.exists(hash):
        approved_hashes = pd.read_csv(hash)

    else:
        logger.warning(f"File not found: {hash}")

    if os.path.exists(path1):
        df1 = pd.read_csv(path1)
    else:
        logger.warning(f"File not found: {path1}")

    if os.path.exists(path2):
        df2 = pd.read_csv(path2)
    else:
        logger.warning(f"File not found: {path2}")

    if df1.empty and df2.empty:
        logger.warning("Both DataFrames are empty. Skipping sort.")
        approved_paths = pd.DataFrame()
    else:
        approved_paths = pd.concat([df1, df2], ignore_index=True)

    if os.path.exists(publishers):
        approved_publishers = pd.read_csv(publishers)

    else:
        logger.warning(f"File not found: {publishers}")

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

    for name, df in dataframes.items():
        logger.debug(f" DataFrame headers: {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)
        
        df.to_csv(f"{working_dir}\\Preflight\\{name}.csv", index=False)
        formatHTML(df, f"{working_dir}\\Preflight\\HTML\\{name}.html")

def splitFilepathsGrouped(df, col="filename"):
    path_exclusion_constant = get_protected_value("PATH_EXCLUSION_CONST", cast_type= int)
    min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type= int)

    def clean_split(path):
        if not isinstance(path, (str, bytes, os.PathLike)):
            return []
        parts = os.path.normpath(path).split(os.sep)
        parts = [p for p in parts if p]  # Remove empty strings
        return parts

    # Diagnostic: log any non-string entries
    non_string_entries = df[~df[col].apply(lambda x: isinstance(x, (str, bytes, os.PathLike)))]
    if not non_string_entries.empty:
        print(f"[WARNING] Non-string entries found in column '{col}':")
        print(non_string_entries)

    df = df.copy()
    split_paths = df[col].apply(clean_split)

    # Filter out paths with fewer than `min_files_for_path` components
    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"])

def calculatePath(approved_hashes, split):
    if split:
        dfs_by_policy = [group for _, group in approved_hashes.groupby("policy")]
    else:
        dfs_by_policy = [approved_hashes]

    badpathparts = load_env_json("BAD_PATH_PARTS", "[]")
    min_files_for_path = get_protected_value("MIN_FILES_FOR_PATH", cast_type = int)

    processed_dfs = []

    for df in dfs_by_policy:
        haslcp = 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