from transformers import PretrainedConfig
import torch



class CybersecurityKnowledgeGraphConfig(PretrainedConfig):
    def __init__(
        self,
        event_nugget_model_path : str = "nugget_model_state_dict.pth",
        event_argument_model_path : str = "argument_model_state_dict.pth",
        event_realis_model_path : str = "realis_model_state_dict.pth",
        **kwargs,
    ):
        self.event_nugget_model_path = event_nugget_model_path
        self.event_argument_model_path = event_argument_model_path
        self.event_realis_model_path = event_realis_model_path

        # self.event_nugget_list = utils.event_nugget_list
        # self.event_args_list = utils.event_args_list
        # self.realis_list = utils.realis_list
        # self.arg_2_role = utils.arg_2_role

        self.event_nugget_list = ['O',
         'B-Ransom',
         'I-Ransom',
         'B-DiscoverVulnerability',
         'I-DiscoverVulnerability',
         'B-PatchVulnerability',
         'I-PatchVulnerability',
         'B-Databreach',
         'I-Databreach',
         'B-Phishing',
         'I-Phishing'
        ]

        self.event_args_list = ['O',
        'B-System',
        'I-System',
        'B-Organization',
        'B-Money',
        'I-Money',
        'B-Device',
        'B-Person',
        'I-Person',
        'B-Vulnerability',
        'I-Vulnerability',
        'B-Capabilities',
        'I-Capabilities',
        'I-Organization',
        'B-PaymentMethod',
        'I-PaymentMethod',
        'B-Data',
        'I-Data',
        'B-Number',
        'I-Number',
        'B-Malware',
        'I-Malware',
        'B-PII',
        'I-PII',
        'B-CVE',
        'I-CVE',
        'B-Purpose',
        'I-Purpose',
        'B-File',
        'I-File',
        'I-Device',
        'B-Time',
        'I-Time',
        'B-Software',
        'I-Software',
        'B-Patch',
        'I-Patch',
        'B-Version',
        'I-Version',
        'B-Website',
        'I-Website',
        'B-GPE',
        'I-GPE'
        ]

        self.realis_list =["O",
            "Generic", 
            "Other", 
            "Actual"
        ]
        self.arg_2_role = {
            "File" : ['Tool', 'Trusted-Entity'],
            "Person" : ['Victim', 'Attacker', 'Discoverer', 'Releaser', 'Trusted-Entity', 'Vulnerable_System_Owner'],
            "Capabilities" : ['Attack-Pattern', 'Capabilities', 'Issues-Addressed'],
            "Purpose" : ['Purpose'],
            "Time" : ['Time'],
            "PII" : ['Compromised-Data', 'Trusted-Entity'],
            "Data" : ['Compromised-Data', 'Trusted-Entity'],
            "Organization" : ['Victim', 'Releaser', 'Discoverer', 'Attacker', 'Vulnerable_System_Owner', 'Trusted-Entity'],
            "Patch" : ['Patch'],
            "Software" : ['Vulnerable_System', 'Victim', 'Trusted-Entity', 'Supported_Platform'],
            "Vulnerability" : ['Vulnerability'],
            "Version" : ['Patch-Number', 'Vulnerable_System_Version'],
            "Device" : ['Vulnerable_System', 'Victim', 'Supported_Platform'],
            "CVE" : ['CVE'],
            "Number" : ['Number-of-Data', 'Number-of-Victim'],
            "System" : ['Victim', 'Supported_Platform', 'Vulnerable_System', 'Trusted-Entity'],
            "Malware" : ['Tool'],
            "Money" : ['Price', 'Damage-Amount'],
            "PaymentMethod" : ['Payment-Method'],
            "GPE" : ['Place'],
            "Website" : ['Trusted-Entity', 'Tool', 'Vulnerable_System', 'Victim', 'Supported_Platform'],
        }

        super().__init__(**kwargs)