#    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.pathfunctions as pathf
import utils.pretty as ct
from AirlockTools import tryToReadCSV

#Standard Libary Imports:
import ast
import gc
import os
import re

#3rd Party Imports:
import pandas as pd

def split_filepaths_grouped(df, col="filename", group_parts=4, min_parts=4):
    def clean_split(path):
        parts = os.path.normpath(path).split(os.sep)
        # Remove leading empty strings caused by UNC paths
        parts = [p for p in parts if p]
        return parts

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

    # Filter out paths with fewer than `min_parts` components
    df = df[split_paths.apply(lambda parts: len(parts) >= min_parts)].copy()
    split_paths = split_paths[df.index]  # Update split_paths to match filtered df

    df["group_key"] = split_paths.apply(lambda parts: os.sep.join(parts[:group_parts]))
    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 mask_from_csv(df, csv_path, filepath_col):
    """
    Reads reviewed CSV of groups, keeps only files in approved groups.
    """
    review_df = pd.read_csv(csv_path)

    def parse_paths(val):
        if isinstance(val, str):
            try:
                # Try to parse as a list
                parsed = ast.literal_eval(val)
                # If it's not a list, wrap it
                return parsed if isinstance(parsed, list) else [parsed]
            except (ValueError, SyntaxError):
                # If parsing fails, treat it as a single path
                return [val]
        return [val]

    review_df[filepath_col] = review_df[filepath_col].apply(parse_paths)

    # Flatten all approved file paths into a set for masking
    approved_files = set()
    for paths in review_df[filepath_col]:
        approved_files.update(paths)

    # Keep only rows in df that are in approved_files
    masked_df = df[df[filepath_col].isin(approved_files)].copy()
    remainder = df[~df[filepath_col].isin(approved_files)].copy()
    return remainder

def filter_and_drop(approved, eligiblepaths, min_hashes):
    """
    Filters eligiblepaths to rows where all hashes are in approved,
    then drops rows with fewer than min_hashes hashes.
    """
    approved_hashes = set(approved['sha256'])

    def all_hashes_approved(row):
        return all(h in approved_hashes for h in row['sha256'])

    filtered = eligiblepaths[eligiblepaths.apply(all_hashes_approved, axis=1)]
    filtered = filtered[filtered['sha256'].apply(len) >= min_hashes]
    
    return filtered

def inspect_parquet(path):
    try:
        df = pd.read_parquet(path)
        print(f"✅ Successfully read: {path}")
        print(f"📄 Columns: {df.columns.tolist()}")
        print(f"🔢 Rows: {len(df)}")
        return df
    except Exception as e:
        print(f"❌ Error reading {path}: {e}")
        return pd.DataFrame()


def regulator(paths, case_insensitive=True):
    """
    Build a regex pattern that matches any of the given Windows path fragments.
    """
    escaped = [re.escape(p) for p in paths]
    pattern = "(?:" + "|".join(escaped) + ")"
    if case_insensitive:
        pattern = "(?i)" + pattern  # Add inline case-insensitive flag
    print(f"Regulator is providing: {pattern}")
    return pattern

def generatePathReview(first_policy, second_policy, badpathparts, min_files_for_path):

    if not os.path.exists(f"parquet\\all_approved_hashes_{first_policy}_{second_policy}.parquet"):

        df1 = tryToReadCSV(f"approved\\hashes_rep_unknown_{first_policy}_{second_policy}.csv")
        df2 = tryToReadCSV(f"approved\\hashes_rep_good_{first_policy}_{second_policy}.csv")

        all_approved_hashes = pd.concat([df1 , df2], ignore_index=True).sort_values(by=['filename'])



        print(ct.colorText(f"Approved hash lists have been combined","green"))

        all_approved_hashes.to_parquet(f"parquet\\all_approved_hashes_{first_policy}_{second_policy}.parquet", index=False)
        del all_approved_hashes
        gc.collect()

    if not os.path.exists(f"parquet\\path_needs_approved_{first_policy}_{second_policy}.parquet"):
        all_approved_hashes = pd.read_parquet(f"parquet\\all_approved_hashes_{first_policy}_{second_policy}.parquet")
        print(ct.colorText(f"Beginning calculating longest common filepaths for path exceptions","green"))
        
        haslcp = pathf.split_filepaths_grouped(all_approved_hashes)
        haslcp.drop_duplicates()

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

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

        # Make a real DataFrame copy before modifying
        lcp_not_forbidden = haslcp[~forbidden_lcfp].copy()

        #For the review, drop down to only the columns we care, and then group by the commmon file path, consolidating and dropping dupes
        lcp_not_forbidden_review = lcp_not_forbidden[['longestcfp', 'middle', 'filename_only', 'file_extension', 'sha256']]				

        # Count unique sha256 per longestcfp
        unique_sha_counts = lcp_not_forbidden_review.groupby('longestcfp')['sha256'].nunique().reset_index()
        unique_sha_counts.columns = ['longestcfp', 'unique_sha256_count']

        # Merge the count back into the original DataFrame
        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]
        
        lcp_not_forbidden_review.to_parquet(f"parquet\\path_needs_approved_{first_policy}_{second_policy}.parquet",index=False)
        lcp_not_forbidden_review.to_csv(f"needs_approved\\path_needs_approved_{first_policy}_{second_policy}.csv",index=False)
        ct.style_dataframe_dark(lcp_not_forbidden_review,f"needs_approved\\path_needs_approved_{first_policy}_{second_policy}.html", True)

        del lcp_not_forbidden
        del unique_sha_counts
        del lcp_not_forbidden_review