#    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 pandas as pd
import os
from itertools import chain
import ast




def split_path(path):
    parts = []
    while True:
        head, tail = os.path.split(path)
        if tail:
            parts.insert(0, tail)
            path = head
        else:
            if head:
                parts.insert(0, head)
            break
    return parts

def local_common_pass(paths, min_parts=3):
    results = {}
    paths_sorted = sorted(paths)
    for i, path in enumerate(paths_sorted):
        candidates = []

        if i > 0:
            try:
                candidates.append(os.path.commonpath([path, paths_sorted[i-1]]))
            except ValueError:
                # different drives, skip
                pass
        if i < len(paths_sorted) - 1:
            try:
                candidates.append(os.path.commonpath([path, paths_sorted[i+1]]))
            except ValueError:
                # different drives, skip
                pass

        best = path
        best_len = 0
        for c in candidates:
            parts = split_path(c)
            if len(parts) >= min_parts and len(parts) > best_len:
                best = c
                best_len = len(parts)
        results[path] = best
    return results


def add_longest_common_two_local(df, col="filename_x", new_col="longestcfp", min_parts=3):
    dirs_series = df[col].astype(str).apply(os.path.dirname)
    first_pass = local_common_pass(dirs_series.tolist(), min_parts)
    second_pass = local_common_pass(list(first_pass.values()), min_parts)
    df[new_col] = dirs_series.map(lambda d: second_pass[first_pass[d]])
    return df


def export_groups_for_review(df, col="filename_x", group_col="longestcfp", csv_path="filegroups_review.csv"):
    """
    Compute longest common paths, group filepaths, write CSV for review.
    """
    df = add_longest_common_two_local(df, col=col, new_col=group_col)
    grouped = df.groupby(group_col)[col].apply(list).reset_index()
    grouped.to_csv(csv_path, index=False)
    print(f"Grouped file list saved to: {csv_path}")
    return grouped, df


def mask_from_csv(df, csv_path, filepath_col="filename_x", group_col="longestcfp"):
    """
    Reads reviewed CSV of groups, keeps only files in approved groups.
    """
    review_df = pd.read_csv(csv_path)
    # Convert string representation of lists back to actual lists
    review_df[filepath_col] = review_df[filepath_col].apply(ast.literal_eval)
    
    # 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()
    return masked_df

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