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 = "A Server-side Template Injection (SSTI) vulnerability exists in Nystudio107 Seomatic 3.4.12 in src/helpers/UrlHelper.php via the host header."
    detail_type, details = driller(description)
    # CVE-2021-44618

    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)
