import json

from pathlib import Path
from cve_driller.driller import CVEDriller


if __name__ == "__main__":
    driller = CVEDriller(verbose=True)
    cve_id = "CVE-2018-7187"

    input_path = f"data/inputs/CWE-78/comp_config_condition/{cve_id}.txt"
    #description = open(input_path, 'r').read()
    description = "The PlanSo Forms WordPress plugin through 2.6.3 does not escape the title of its Form before outputting it in attributes, allowing high privilege users such as admin to set XSS payload in it, even when the unfiltered_html is disallowed, leading to an Authenticated Stored Cross-Site Scripting issue."
    detail_type, details = driller(description)
    # CVE-2018-16460

    if details:
        print(json.dumps(details, indent=4))
    exit()
    input_path = Path("data/inputs")
    output_path = Path("data/outputs")

    dataset_path = Path("data/dataset.jsonl")

    with dataset_path.open(mode="a") as df:
        for cwe in input_path.iterdir():
            if not cwe.is_dir():
                continue

            for detail_type in cwe.iterdir():
                if not detail_type.is_dir():
                    continue

                for filename in detail_type.iterdir():
                    if filename.suffix != ".txt":
                        print(f"Skipping {filename}, does not have a .txt extension")
                        continue

                    with filename.open(mode="r") as f:
                        prompt = f.read()

                    output_file = output_path / cwe.name / detail_type.name / f"{filename.stem}.json"

                    if not output_file.exists():
                        print(f"Skipping {filename}, does not exist")
                        continue

                    with output_file.open(mode="r") as f:
                        completion = json.load(f)

                    jsonl = {"messages": [
                        {"role": "user", "content": prompt},
                        {"role": "assistant", "content": json.dumps(completion)}]
                    }

                    df.write(json.dumps(jsonl) + "\n")

    #if details:
    #    output_path = f"data/outputs/{detail_type.value}/{cve_id}.json"

    #    with open(output_path, 'w') as f:
    #        json.dump(details, f, indent=4)
