from multiprocessing.pool import ThreadPool

from sdcclient import SdScanningClient


def get_vulnerability_details(client: SdScanningClient, id):
    ok, res = client.get_vulnerability_details(id)
    if not ok:
        raise Exception(res)

    return res


def add_vulnerability_exception_bundle(client: SdScanningClient, name, comment):
    ok, res = client.add_vulnerability_exception_bundle(name=name, comment=comment)
    if not ok:
        if res == "Conflict":
            res = "This bundle already exists"
        raise Exception(res)

    return res


def list_vulnerability_exception_bundles(client: SdScanningClient):
    ok, res = client.list_vulnerability_exception_bundles()
    if not ok:
        raise Exception(res)

    pool = ThreadPool()
    res = pool.map(lambda b: get_vulnerability_exception_bundle(client, b["id"]), res)

    return res


def add_vulnerability_exception(client: SdScanningClient, bundle, cve, note, expiration_date):
    ok, res = client.add_vulnerability_exception(bundle, cve, note, expiration_date)
    if not ok:
        raise Exception(res)

    return res


def delete_vulnerability_exception(client: SdScanningClient, bundle, cve):
    ok, res = client.get_vulnerability_exception_bundle(bundle)
    if not ok:
        raise Exception(res)

    found_cves = [c for c in res["items"] if c["trigger_id"] == cve]
    if len(found_cves) == 0:
        raise Exception(f"CVE not found in bundle {bundle}")

    ok, res = client.delete_vulnerability_exception(bundle, found_cves[0]["id"])
    if not ok:
        raise Exception(res)

    return res


def update_vulnerability_exception(client: SdScanningClient, bundle, id, cve=None, enabled=None, note="__unset__",
                                   expiration_date="__unset__"):
    ok, res = client.get_vulnerability_exception_bundle(bundle)
    if not ok:
        raise Exception(res)

    vulns_found = [vuln for vuln in res["items"] if vuln["id"] == id]
    if len(vulns_found) == 0:
        raise Exception(f"Vulnerability exception not found for bundle {res['name']} with ID: {id}")

    vuln = vulns_found[0]
    if cve is None:
        cve = vuln["trigger_id"]
    if enabled is None:
        enabled = vuln["enabled"]
    if note == "__unset__":
        note = vuln["notes"]
    if expiration_date == "__unset__":
        expiration_date = vuln["expiration_date"]

    ok, res = client.update_vulnerability_exception(bundle, id, cve, enabled, note, expiration_date)
    if not ok:
        raise Exception(res)

    return res


def get_vulnerability_exception_bundle(client: SdScanningClient, bundle):
    ok, res = client.get_vulnerability_exception_bundle(bundle)
    if not ok:
        raise Exception(res)
    return res


def delete_vulnerability_exception_bundle(client: SdScanningClient, bundle):
    ok, res = client.delete_vulnerability_exception_bundle(bundle)
    if not ok:
        raise Exception(res)
    return res


def download_vulnerability_report(client: SdScanningClient, runtime, os, non_os, out_file):
    scope = "runtime" if runtime else "static"

    header = ""
    os_data = ""
    non_os_data = ""

    if (not os and not non_os) or os:
        ok, res = client.download_cve_report_csv(vuln_type="os", scope_type=scope)
        if not ok:
            raise Exception(res)
        os_data = "\n".join(res.splitlines()[1:]) + "\n"
        header = res.splitlines()[0] + "\n"

    if (not os and not non_os) or non_os:
        ok, res = client.download_cve_report_csv(vuln_type="non-os", scope_type=scope)
        if not ok:
            raise Exception(res)
        non_os_data = "\n".join(res.splitlines()[1:]) + "\n"
        header = res.splitlines()[0] + "\n"

    out_file.write(header)
    out_file.write(os_data)
    out_file.write(non_os_data)
