import streamlit as st
from transformers import AutoModelForTokenClassification
from annotated_text import annotated_text
import numpy as np
import os, joblib

from utils import get_idxs_from_text

model = AutoModelForTokenClassification.from_pretrained("CyberPeace-Institute/Cybersecurity-Knowledge-Graph", trust_remote_code=True)

role_classifiers = {}
folder_path = '/arg_role_models'
for filename in os.listdir(os.getcwd() + folder_path):
    if filename.endswith('.joblib'):
        file_path = os.getcwd() + os.path.join(folder_path, filename)
        clf = joblib.load(file_path)
        arg = filename.split(".")[0]
        role_classifiers[arg] = clf

def annotate(name):
    tokens = [item["token"] for item in output]
    tokens = [token.replace(" ", "") for token in tokens]
    text = model.tokenizer.decode([item["id"] for item in output])
    idxs = get_idxs_from_text(text, tokens)
    labels = [item[name] for item in output]

    annotated_text_list = []
    last_label = ""
    cumulative_tokens = "" 
    last_id = 0
    for idx, label in zip(idxs, labels):
        to_label = label
        label_short = to_label.split("-")[1] if "-" in to_label else to_label
        if last_label == label_short:
            cumulative_tokens += text[last_id : idx["end_idx"]]
            last_id = idx["end_idx"]
        else:
            if last_label != "":
                if last_label == "O":
                    annotated_text_list.append(cumulative_tokens)
                else:
                    annotated_text_list.append((cumulative_tokens, last_label))
            last_label = label_short
            cumulative_tokens = idx["word"]
            last_id = idx["end_idx"]
    if last_label == "O":
        annotated_text_list.append(cumulative_tokens)
    else:  
        annotated_text_list.append((cumulative_tokens, last_label))
    annotated_text(annotated_text_list)

def get_arg_roles(output):
    args = [(idx, item["argument"], item["token"]) for idx, item in enumerate(output) if item["argument"]!= "O"]
        
    entities = []
    current_entity = None
    for position, label, token in args:
        if label.startswith('B-'):
            if current_entity is not None:
                entities.append(current_entity)
            current_entity = {'label': label[2:], 'text': token.replace(" ", ""), 'start': position, 'end': position}
        elif label.startswith('I-'):
            if current_entity is not None:
                current_entity['text'] += ' ' + token.replace(" ", "")
                current_entity['end'] = position
    for entity in entities:
        context = model.tokenizer.decode([item["id"] for item in output[max(0, entity["start"] - 15) : min(len(output), entity["end"] + 15)]])
        entity["context"] = context
    
    for entity in entities:
        if len(model.arg_2_role[entity["label"]]) > 1:
            sent_embed = model.embed_model.encode(entity["context"])
            arg_embed = model.embed_model.encode(entity["text"])
            embed = np.concatenate((sent_embed, arg_embed))
            arg_clf = role_classifiers[entity["label"]]
            role_id = arg_clf.predict(embed.reshape(1, -1))
            role = model.arg_2_role[entity["label"]][role_id[0]]
            entity["role"] = role
        else:
            entity["role"] = model.arg_2_role[entity["label"]][0]
    
    for item in output:
        item["role"] = "O"
    for entity in entities:
        for i in range(entity["start"], entity["end"] + 1):
            output[i]["role"] = entity["role"]
    return output

st.title("Create Knowledge Graphs from Cyber Incidents")

text_input = st.text_area("Enter your text here", height=100)

if text_input or st.button('Apply'):
    output = model(text_input)
    st.subheader("Event Nuggets")
    annotate("nugget")
    st.subheader("Event Arguments")
    annotate("argument")
    st.subheader("Realis of Event Nuggets")
    annotate("realis")
    output = get_arg_roles(output)
    st.subheader("Role of the Event Arguments")
    annotate("role")