import pandas as pd
import requests
import os
import json


def aggregateHashes(executions_json) -> pd.DataFrame:
    """
    Takes the executions,  aggregates all the data with sha256 as primary, then returns aggregated dataframe
    """
    """
   old version 
   data = executions_json.json()

    df = pd.DataFrame(data["response"]["exechistories"])
    print(df)
    if df.empty:
        return df 
    
    #Aggregate by sha256 - keep all entries in lists
    agg_df = df.groupby("sha256").agg(lambda x: list(x)).reset_index()
    
    return agg_df
    """


    data = executions_json.json()
    df = pd.DataFrame(data["response"]["exechistories"])

    if df.empty:
         return df

    # Aggregate by sha256, deduplicate lists, and preserve order
    agg_df = df.groupby("sha256").agg(lambda x: list(dict.fromkeys(x))).reset_index()

    # Add a column for the number of unique hostnames
    agg_df["num_devices"] = agg_df["hostname"].apply(len)

    # Sort by num_devices in descending order
    agg_df = agg_df.sort_values("num_devices", ascending=False)

    return agg_df

def augmentAggregatedHashes(url, agg_df: pd.DataFrame) -> pd.DataFrame:
    """
    Takes output of aggregatedHashes, queries API for those hashes, flattens response while keeping one row per hash,
    aggregate applications and baselines into lists, then merges results back into agg_df to create a 
    """
    endpoint = url + 'v1/hash/query'
    payload = agg_df['sha256'].tolist()
    headers = {"X-APIKey": os.getenv('APIKEY')}

    response = requests.post(endpoint, headers=headers, json=payload, verify=False)
    data = response.json()  
    results = data.get("response", {}).get("results", [])

    rows = []
    for res in results:
        row = {"sha256": res.get("sha256"), "result": res.get("result")}

        if "data" in res:
            d = res["data"]
            for key in ["filename", "filepath", "description", "filesize", "md5",
                        "productname", "productversion", "publisher", "createtime", "modtime",
                        "sha128", "sha384", "sha512", "datetime"]:
                row[key] = d.get(key)

            row["applications"] = d.get("applications", [])
            row["baselines"] = d.get("baselines", [])

            reputation = d.get("reputation", {})
            for k, v in reputation.items():
                row[f"reputation_{k}"] = v

        rows.append(row)

    df_api = pd.DataFrame(rows)

    aug_df = agg_df.merge(df_api, on="sha256", how="left")

    return aug_df

def categorize_hashes(aug_df: pd.DataFrame, threat_tolerance: int, untrusted_publishers: list):
    """
    Categorize hashes into needsreview, approved, and remaining based on publisher and threat level.
    
    """
    if untrusted_publishers is None:
        untrusted_publishers = []

    # Flatten threatlevel from nested reputation dict
    df = aug_df.copy()  
    df["threatlevel"] = df["reputation"].apply(lambda x: x.get("threatlevel") if pd.notnull(x) else None)

    # Masks for each category
    mask_needsreview = (df["publisher"] == "Not Signed") & (df["threatlevel"] > threat_tolerance)
    mask_approved = (df["publisher"] != "Not Signed") & (~df["publisher"].isin(untrusted_publishers))

    # Create DataFrames for each category
    needsreview_df = df[mask_needsreview].drop(columns=["threatlevel"])
    approved_df = df[mask_approved].drop(columns=["threatlevel"])
    remaining_df = df[~(mask_needsreview | mask_approved)].drop(columns=["threatlevel"])

    return needsreview_df, approved_df, remaining_df

def approve_hashes(approved_df: pd.DataFrame):
    pass