import pandas as pd

from tqdm import tqdm
from pathlib import Path
from cve_driller.driller import CVEDriller

MULTI_SENTENCE_PATTERN = r'\w+\.\s+[A-Z0-9]'

if __name__ == "__main__":
    driller = CVEDriller()
    path = Path("~/Downloads/cve_impact_types.csv").expanduser()
    path_labels = Path("~/projects/nvdutils/bert_impact_type_labels.csv").expanduser()
    df = pd.read_csv(path, engine='python')
    df_labels = pd.read_csv(path_labels, engine='python')
    print(df['impact'].value_counts())
    df = df[df['impact'].isnull()]
    print(len(df))
    df.drop(columns=['impact'], inplace=True)

    df_labeled = df.merge(df_labels, how='left', left_on='cve-id', right_on='cve-id')
    print(df_labeled["impact"].value_counts())
    not_label_df = df_labeled[df_labeled['impact'] == "Not"]
    # Assuming your DataFrame is named df

    # To filter OUT rows matching the pattern:
    # filtered_df = not_label_df[~not_label_df['description'].str.contains(MULTI_SENTENCE_PATTERN, regex=True, na=False)]
    # print(len(filtered_df))
    impact_count = 0
    match_count = 0
    extracted_impact = []

    for i, row in tqdm(not_label_df.iterrows(), total=len(not_label_df)):
        detail_type, details = driller(row['description'])

        if details:
            match_count += 1

            if 'impact' in details:
                impact_count += 1
                extracted_impact.append(details['impact'])

    print(match_count, impact_count)
    print(extracted_impact)
