import os
import shutil
import tempfile
import json
import yaml

import pytest

from vdb.lib import db6, search
from vdb.lib.custom import CustomSource

test_tmp_dir = os.getenv(
    "TEST_VDB_HOME", os.path.join(tempfile.gettempdir(), "vdb6-custom-tests")
)

db6.get(":memory:", ":memory:")


@pytest.fixture(autouse=True)
def setup_teardown():
    """
    Fixture to ensure a clean environment for every test.
    """
    if os.path.exists(test_tmp_dir):
        shutil.rmtree(test_tmp_dir, ignore_errors=True)
    os.makedirs(test_tmp_dir, exist_ok=True)

    db6.clear_all()
    search.CUSTOM_DATA_CACHE = []
    if hasattr(search, "CUSTOM_OVERRIDE_INDEX"):
        search.CUSTOM_OVERRIDE_INDEX = set()
    yield
    if os.path.exists(test_tmp_dir):
        shutil.rmtree(test_tmp_dir, ignore_errors=True)


def create_custom_file(filename, content):
    """Helper to create a temporary custom vulnerability file"""
    file_path = os.path.join(test_tmp_dir, filename)
    with open(file_path, "w", encoding="utf-8") as f:
        if filename.endswith(".yaml") or filename.endswith(".yml"):
            yaml.dump(content, f)
        elif filename.endswith(".json"):
            json.dump(content, f)
    return file_path


@pytest.fixture
def private_cve_yaml():
    return {
        "dataType": "CVE_RECORD",
        "dataVersion": "5.1",
        "cveMetadata": {
            "cveId": "PRIVATE-2024-001",
            "assignerOrgId": "00000000-0000-4000-8000-000000000000",
            "state": "PUBLISHED",
            "datePublished": "2024-01-01T00:00:00Z",
            "dateUpdated": "2024-01-01T00:00:00Z",
        },
        "containers": {
            "cna": {
                "providerMetadata": {"orgId": "00000000-0000-4000-8000-000000000000"},
                "descriptions": [{"lang": "en", "value": "Private vuln"}],
                "references": [{"url": "http://internal"}],
                "affected": [
                    {
                        "vendor": "internal",
                        "product": "my-lib",
                        "packageName": "my-lib",
                        "packageURL": "pkg:pypi/my-lib",
                        "versions": [
                            {
                                "version": "1.0.0",
                                "status": "affected",
                                "versionType": "semver",
                                "lessThan": "2.0.0",
                            }
                        ],
                    }
                ],
            }
        },
    }


@pytest.fixture
def override_cve_yaml():
    """Defines a CVE as UNAFFECTED for a specific version range"""
    return {
        "dataType": "CVE_RECORD",
        "dataVersion": "5.1",
        "cveMetadata": {
            "cveId": "CVE-2023-9999",
            "assignerOrgId": "00000000-0000-4000-8000-000000000000",
            "state": "PUBLISHED",
        },
        "containers": {
            "cna": {
                "providerMetadata": {"orgId": "00000000-0000-4000-8000-000000000000"},
                "descriptions": [{"lang": "en", "value": "Override test"}],
                "references": [{"url": "http://internal"}],
                "affected": [
                    {
                        "vendor": "public",
                        "product": "public-lib",
                        "packageName": "public-lib",
                        "packageURL": "pkg:npm/public-lib",
                        "versions": [
                            {
                                "version": "1.2.3",
                                "status": "unaffected",
                                "versionType": "semver",
                            }
                        ],
                    }
                ],
            }
        },
    }


def test_custom_source_load(private_cve_yaml):
    """Test loading and normalizing a custom YAML file"""
    create_custom_file("private.yaml", private_cve_yaml)

    cs = CustomSource(test_tmp_dir)
    results = cs.load_in_memory()

    assert len(results) == 1
    item = results[0]
    assert item["cve_id"] == "PRIVATE-2024-001"
    assert item["purl_prefix"] == "pkg:pypi/my-lib"
    assert "vers:pypi/>=1.0.0|<2.0.0" in item["vers"]
    assert item["source_data"]


def test_overlay_search_private_cve(private_cve_yaml):
    """Test searching for a private CVE that only exists in custom data"""
    create_custom_file("private.yaml", private_cve_yaml)
    search.load_custom_data(test_tmp_dir)

    results = search.search_by_purl_like("pkg:pypi/my-lib@1.5.0", with_data=False)
    assert len(results) == 1
    assert results[0]["cve_id"] == "PRIVATE-2024-001"

    results = search.search_by_purl_like("pkg:pypi/my-lib@2.1.0", with_data=False)
    assert len(results) == 0

    results = search.search_by_cve("PRIVATE-2024-001", with_data=True)
    assert len(results) == 1
    assert (
        results[0]["source_data"].root.containers.cna.descriptions.root[0].value
        == "Private vuln"
    )


def test_override_false_positive(override_cve_yaml):
    """
    Test the Priority Logic via CNA override.
    """
    db_conn, index_conn = db6.get()
    index_conn.execute(
        "INSERT INTO cve_index(cve_id, type, namespace, name, vers, purl_prefix) VALUES (?, ?, ?, ?, ?, ?)",
        ("CVE-2023-9999", "npm", "", "public-lib", "vers:npm/*", "pkg:npm/public-lib"),
    )
    results = search.search_by_purl_like("pkg:npm/public-lib@1.2.3")
    assert len(results) == 1
    assert results[0]["cve_id"] == "CVE-2023-9999"

    create_custom_file("override.yaml", override_cve_yaml)
    search.load_custom_data(test_tmp_dir)

    results = search.search_by_purl_like("pkg:npm/public-lib@1.2.3")
    assert len(results) == 0


def test_override_with_adp():
    """
    Test overriding a DB vulnerability using an ADP container (without CNA).
    """
    db_conn, index_conn = db6.get()
    index_conn.execute(
        "INSERT INTO cve_index(cve_id, type, namespace, name, vers, purl_prefix) VALUES (?, ?, ?, ?, ?, ?)",
        ("CVE-2024-ADP-TEST", "npm", "", "adp-lib", "vers:npm/*", "pkg:npm/adp-lib"),
    )

    results = search.search_by_purl_like("pkg:npm/adp-lib@1.0.0")
    assert len(results) == 1
    assert results[0]["cve_id"] == "CVE-2024-ADP-TEST"

    adp_override_data = {
        "dataType": "CVE_RECORD",
        "dataVersion": "5.1",
        "cveMetadata": {
            "cveId": "CVE-2024-ADP-TEST",
            "assignerOrgId": "00000000-0000-4000-8000-000000000000",
            "state": "PUBLISHED",
        },
        "containers": {
            "adp": [
                {
                    "providerMetadata": {
                        "orgId": "00000000-0000-4000-8000-000000000000"
                    },
                    "descriptions": [{"lang": "en", "value": "ADP Override"}],
                    "affected": [
                        {
                            "product": "adp-lib",
                            "packageName": "adp-lib",
                            "packageURL": "pkg:npm/adp-lib",
                            "versions": [
                                {
                                    "version": "1.0.0",
                                    "status": "unaffected",
                                    "versionType": "semver",
                                }
                            ],
                        }
                    ],
                }
            ]
        },
    }
    create_custom_file("adp_override.yaml", adp_override_data)
    search.load_custom_data(test_tmp_dir)

    results = search.search_by_purl_like("pkg:npm/adp-lib@1.0.0")
    assert len(results) == 0


def test_override_partial_fix(override_cve_yaml):
    """
    Test modifying an existing CVE range.
    """
    override_cve_yaml["containers"]["cna"]["affected"][0]["versions"] = [
        {
            "version": "0.0.0",
            "status": "affected",
            "versionType": "semver",
            "lessThan": "1.0.0",
        }
    ]

    db_conn, index_conn = db6.get()
    index_conn.execute(
        "INSERT INTO cve_index(cve_id, type, namespace, name, vers, purl_prefix) VALUES (?, ?, ?, ?, ?, ?)",
        ("CVE-2023-9999", "npm", "", "public-lib", "vers:npm/*", "pkg:npm/public-lib"),
    )

    create_custom_file("fix.yaml", override_cve_yaml)
    search.load_custom_data(test_tmp_dir)

    results = search.search_by_purl_like("pkg:npm/public-lib@0.9.0")
    assert len(results) == 1
    assert results[0]["cve_id"] == "CVE-2023-9999"

    results = search.search_by_purl_like("pkg:npm/public-lib@1.5.0")
    assert len(results) == 0


def test_schema_validation_failure():
    """Test that invalid files are skipped gracefully"""
    bad_data = {"invalid": "schema", "cveMetadata": {}}
    create_custom_file("bad.json", bad_data)

    cs = CustomSource(test_tmp_dir)
    results = cs.load_in_memory()
    assert len(results) == 0


def test_mixed_results_merge():
    """Test that valid DB results persist alongside Custom results"""
    db_conn, index_conn = db6.get()
    index_conn.execute(
        "INSERT INTO cve_index(cve_id, type, namespace, name, vers, purl_prefix) VALUES (?, ?, ?, ?, ?, ?)",
        ("CVE-2020-5555", "npm", "", "public-lib", "vers:npm/*", "pkg:npm/public-lib"),
    )

    cve_data = {
        "dataType": "CVE_RECORD",
        "dataVersion": "5.1",
        "cveMetadata": {
            "cveId": "PRIVATE-NEW",
            "assignerOrgId": "00000000-0000-4000-8000-000000000000",
            "state": "PUBLISHED",
        },
        "containers": {
            "cna": {
                "providerMetadata": {"orgId": "00000000-0000-4000-8000-000000000000"},
                "descriptions": [{"lang": "en", "value": "New"}],
                "affected": [
                    {
                        "product": "public-lib",
                        "packageURL": "pkg:npm/public-lib",
                        "versions": [{"version": "0", "status": "affected"}],
                    }
                ],
            }
        },
    }
    create_custom_file("new.yaml", cve_data)
    search.load_custom_data(test_tmp_dir)

    results = search.search_by_purl_like("pkg:npm/public-lib@1.0.0")
    ids = sorted([r["cve_id"] for r in results])
    assert ids == ["CVE-2020-5555", "PRIVATE-NEW"]
