import os
from typing import Dict, List, Any, Generator

try:
    import yaml
except ImportError:
    yaml = None

try:
    import tomllib
except ImportError:
    try:
        import toml as tomllib
    except ImportError:
        tomllib = None

from vdb.lib import utils
from vdb.lib.cve_model import CVE


class CustomSource:
    """
    Source for loading custom vulnerability data strictly following CVE 5.2 schema
    from local files (JSON, YAML, TOML) for runtime overlay.
    """

    def __init__(self, path: str):
        self.path = path

    def load_in_memory(self) -> List[Dict]:
        """
        Load processed data into memory for runtime overlay search.
        """
        results = []
        if not os.path.exists(self.path):
            return results

        for item in self._yield_items():
            results.append(item)

        return results

    def _yield_items(self) -> Generator[Dict, None, None]:
        """Generator that yields processed CVE items normalized for search"""
        for root, _, files in os.walk(self.path):
            for file in files:
                file_path = os.path.join(root, file)
                try:
                    raw_data = self._parse_file(file_path)
                    if not raw_data:
                        continue

                    items = raw_data if isinstance(raw_data, list) else [raw_data]

                    for item in items:
                        if "containers" in item:
                            if (
                                "adp" in item["containers"]
                                and "cna" not in item["containers"]
                            ):
                                item["containers"]["cna"] = None

                        try:
                            cve_obj = CVE.model_validate(item)
                        except Exception as e:
                            print(f"Schema validation failed for {file_path}: {e}")
                            continue
                        cve_record = cve_obj.root
                        cve_id = (
                            cve_record.cveMetadata.cveId.root
                            if hasattr(cve_record.cveMetadata.cveId, "root")
                            else cve_record.cveMetadata.cveId
                        )
                        # Process ADP
                        if cve_record.containers.adp:
                            for adp in cve_record.containers.adp:
                                if adp.affected:
                                    for product in adp.affected.root:
                                        yield from self._process_product(
                                            cve_id, cve_obj, product
                                        )
                        # Process CNA
                        if (
                            cve_record.containers.cna
                            and cve_record.containers.cna.affected
                        ):
                            for product in cve_record.containers.cna.affected.root:
                                yield from self._process_product(
                                    cve_id, cve_obj, product
                                )
                except Exception as e:
                    print(f"Failed to process {file_path}: {e}")

    def _process_product(self, cve_id, cve_obj, product) -> Generator[Dict, None, None]:
        """Convert a Product object into an in-memory index row"""
        purl_obj = {}
        purl_prefix = ""

        if product.packageURL:
            p_url = (
                product.packageURL.root
                if hasattr(product.packageURL, "root")
                else product.packageURL
            )
            purl_obj = utils.parse_purl(str(p_url))

            if purl_obj:
                ptype = purl_obj.get("type")
                namespace = purl_obj.get("namespace")
                name = purl_obj.get("name")
                purl_prefix = f"pkg:{ptype}/"
                if namespace:
                    purl_prefix = f"{purl_prefix}{namespace.replace('@', '%40')}/"
                purl_prefix = f"{purl_prefix}{name}"

        if not purl_prefix:
            return

        vendor_for_vers = purl_obj.get("type", "generic")
        vers_str = ""
        if product.versions:
            vers_str = utils.to_purl_vers(vendor_for_vers, product.versions)

        fix_version = utils.get_unaffected(cve_obj, vers_str)

        yield {
            "cve_id": cve_id,
            "type": purl_obj.get("type", "custom"),
            "namespace": purl_obj.get("namespace", ""),
            "name": purl_obj.get("name", product.packageName or "unknown"),
            "source_data": cve_obj,
            "vers": vers_str,
            "purl_prefix": purl_prefix,
            "fix_version": fix_version,
            "override_data": None,
        }

    def _parse_file(self, file_path: str) -> Any:
        ext = file_path.split(".")[-1].lower()
        with open(file_path, "r", encoding="utf-8") as fp:
            content = fp.read()
            if ext == "json":
                return utils.load(content)
            elif ext in ("yaml", "yml") and yaml:
                return yaml.safe_load(content)
            elif ext == "toml" and tomllib:
                return tomllib.loads(content)
        return None
