import re
from typing import Any, Generator, Dict, List, Tuple

import apsw

from vdb.lib import db6, utils
from vdb.lib.custom import CustomSource
from vdb.lib.cve_model import CVE, CVE1

IS_ADVISORY = re.compile("^[A-Z]{1,7}-")
BATCH_SIZE = 200

CUSTOM_DATA_CACHE: List[Dict] = []
CUSTOM_OVERRIDE_INDEX = set()


def load_custom_data(path: str):
    """Loads custom data into memory"""
    global CUSTOM_DATA_CACHE, CUSTOM_OVERRIDE_INDEX
    cs = CustomSource(path)
    CUSTOM_DATA_CACHE = cs.load_in_memory()
    CUSTOM_OVERRIDE_INDEX = {
        (item["cve_id"], item["purl_prefix"]) for item in CUSTOM_DATA_CACHE
    }


def _merge_results(
    db_results: List[Dict], custom_candidates: List[Dict], search_version: str = None
) -> List[Dict]:
    """
    Merge DB and Custom results.
    """
    final_results = []
    processed_custom_keys = set()
    for item in custom_candidates:
        key = (item["cve_id"], item["purl_prefix"])
        processed_custom_keys.add(key)
        is_match = True
        if search_version and search_version != "*":
            if not utils.vers_compare(search_version, item["vers"]):
                is_match = False
        if is_match:
            hit = item.copy()
            hit["matched_by"] = item["purl_prefix"]
            hit["matching_vers"] = search_version
            final_results.append(hit)
    for item in db_results:
        key = (item["cve_id"], item["purl_prefix"])
        if key in CUSTOM_OVERRIDE_INDEX:
            continue
        final_results.append(item)
    return final_results


def _search_memory(purl_prefix=None, cve_id=None) -> List[Dict]:
    """Pre-filter custom cache by index keys"""
    hits = []
    if not CUSTOM_DATA_CACHE:
        return hits
    search_prefix_lower = purl_prefix.lower() if purl_prefix else None
    for item in CUSTOM_DATA_CACHE:
        if purl_prefix:
            if item.get("purl_prefix", "").lower() == search_prefix_lower:
                hits.append(item)
        elif cve_id:
            if cve_id == item["cve_id"]:
                hits.append(item)
            elif "%" in cve_id and item["cve_id"].startswith(cve_id.replace("%", "")):
                hits.append(item)
    return hits


def _process_hits(raw_hits: List[Tuple]) -> List[Dict[str, Any]]:
    filtered_list = []
    for ahit in raw_hits:
        cve_id = ahit[0]
        vers = ahit[4]
        filtered_list.append(
            {
                "cve_id": cve_id,
                "type": ahit[1],
                "namespace": ahit[2],
                "name": ahit[3],
                "vers": vers,
                "purl_prefix": ahit[-1],
            }
        )
    return filtered_list


def get_cve_data(
    db_conn: apsw.Connection | None, index_hits: List, search_str: str
) -> Generator:
    """Get CVE data for the index results

    Args:
        db_conn: DB Connection or None to create a new one
        index_hits: Hits from one of the search methods
        search_str: Original search string used

    Returns:
        generator: generator for CVE data with original source data as a pydantic model
    """
    if not db_conn:
        db_conn, _ = db6.get(read_only=True)
    for ahit in index_hits:
        results = exec_query(
            db_conn,
            "SELECT DISTINCT cve_id, type, namespace, name, source_data_hash, source_data, override_data, vers, purl_prefix FROM cve_data WHERE cve_id = ? AND vers = ? AND purl_prefix = ? GROUP BY purl_prefix ORDER BY cve_id DESC",
            (ahit["cve_id"], ahit["vers"], ahit["purl_prefix"]),
        )
        for res in results:
            source_data = (
                CVE(root=CVE1.model_validate(apsw.jsonb_decode(res[5]), strict=False))
                if res[5]
                else None
            )
            vers = res[7]
            fix_version = utils.get_unaffected(source_data, vers)
            yield {
                "cve_id": res[0],
                "type": res[1],
                "namespace": res[2],
                "name": res[3],
                "matching_vers": ahit["vers"],
                "matched_by": search_str,
                "source_data_hash": res[4],
                "source_data": source_data,
                "override_data": (apsw.jsonb_decode(res[6]) if res[6] else None),
                "vers": vers,
                "purl_prefix": res[8],
                "fix_version": fix_version,
            }


def get_cve_data_batched(
    db_conn: apsw.Connection | None, index_hits: List, search_str: str
) -> Generator:
    """
    Get CVE data for the index results efficiently.
    Uses batching and CTEs to avoid N+1 query performance issues.
    """
    if not index_hits:
        return

    if not db_conn:
        db_conn, _ = db6.get(read_only=True)
    total_hits = len(index_hits)
    for i in range(0, total_hits, BATCH_SIZE):
        batch = index_hits[i : i + BATCH_SIZE]
        query_args = []
        for hit in batch:
            query_args.extend([hit["cve_id"], hit["vers"], hit["purl_prefix"]])
        placeholders = ",".join(["(?,?,?)"] * len(batch))
        sql = f"""
            WITH lookup(cve_id, vers, purl_prefix) AS (
                VALUES {placeholders}
            )
            SELECT
                d.cve_id, 
                d.type, 
                d.namespace, 
                d.name, 
                d.source_data_hash, 
                d.source_data, 
                d.override_data, 
                d.vers, 
                d.purl_prefix 
            FROM cve_data d
            JOIN lookup l 
                ON d.cve_id = l.cve_id 
                AND d.vers = l.vers 
                AND d.purl_prefix = l.purl_prefix
            GROUP BY d.cve_id, d.vers, d.purl_prefix
            ORDER BY d.cve_id DESC;
        """
        results = exec_query(db_conn, sql, tuple(query_args))
        for res in results:
            source_data_str = res[5]
            override_data_str = res[6]
            source_data = (
                CVE(
                    root=CVE1.model_validate(
                        apsw.jsonb_decode(source_data_str), strict=False
                    )
                )
                if source_data_str
                else None
            )
            vers = res[7]
            fix_version = utils.get_unaffected(source_data, vers)
            yield {
                "cve_id": res[0],
                "type": res[1],
                "namespace": res[2],
                "name": res[3],
                "matching_vers": vers,
                "matched_by": search_str,
                "source_data_hash": res[4],
                "source_data": source_data,
                "override_data": (
                    apsw.jsonb_decode(override_data_str) if override_data_str else None
                ),
                "vers": vers,
                "purl_prefix": res[8],
                "fix_version": fix_version,
            }


def search_by_any(any_str: str, with_data: bool = False) -> List:
    """Convenient method to search by a string"""
    if any_str.startswith("pkg:"):
        return search_by_purl_like(any_str, with_data)
    if IS_ADVISORY.search(any_str):
        return search_by_cve(any_str, with_data)
    if any_str.startswith("http"):
        return search_by_url(any_str, with_data)
    return search_by_cpe_like(any_str, with_data)


def search_by_cpe_like(cpe: str, with_data: bool = False) -> List:
    """Search by CPE or colon-separate strings"""
    db_conn, index_conn = db6.get(read_only=True)
    if cpe.startswith("cpe:"):
        vendor, package, version, _ = utils.parse_cpe(cpe)
    elif cpe.count(":") == 2:
        vendor, package, version = cpe.split(":")
    else:
        return []
    # check for vendor name in both namespace and type
    raw_hits = exec_query(
        index_conn,
        "SELECT DISTINCT cve_id, type, namespace, name, vers, purl_prefix FROM cve_index where (namespace = ? OR type = ?) AND name = ? AND vers_compare(?, vers)",
        (vendor, vendor, package, version),
    )
    filtered_list = _process_hits(raw_hits)
    if with_data:
        return list(get_cve_data_batched(db_conn, filtered_list, cpe))
    return filtered_list


def search_by_purl_like(purl: str, with_data: bool = False) -> List:
    """Search by purl like string"""
    db_conn, index_conn = db6.get(read_only=True)
    purl_obj = utils.parse_purl(purl)
    if purl_obj:
        ptype = purl_obj.get("type")
        namespace = purl_obj.get("namespace")
        name = purl_obj.get("name")
        version = purl_obj.get("version", "*")
        purl_prefix = f"pkg:{ptype}/"
        if namespace:
            purl_prefix = f"{purl_prefix}{namespace.replace('@', '%40')}/"
        # Handle distro names for linux os purls by prefixing distro name to name
        if purl_obj.get("qualifiers") and purl_obj["qualifiers"].get("distro_name"):
            distro_name = purl_obj["qualifiers"].get("distro_name")
            name = f"{distro_name}/{name}"
        purl_prefix = f"{purl_prefix}{name}"
        args = (
            purl_prefix.lower(),
            version,
        )
        raw_hits = exec_query(
            index_conn,
            "SELECT DISTINCT cve_id, type, namespace, name, vers, purl_prefix FROM cve_index where lower(purl_prefix) = ? AND vers_compare(?, vers)",
            args,
        )
        # Filter and merge the db hits with any custom data hits
        db_index_hits = _process_hits(raw_hits)
        custom_candidates = _search_memory(purl_prefix=purl_prefix)
        override_keys = {(c["cve_id"], c["purl_prefix"]) for c in custom_candidates}
        filtered_db_hits = [
            hit
            for hit in db_index_hits
            if (hit["cve_id"], hit["purl_prefix"]) not in override_keys
        ]
        filtered_custom_hits = []
        for hit in custom_candidates:
            if version == "*" or utils.vers_compare(version, hit["vers"]):
                h = hit.copy()
                h["matched_by"] = purl_prefix
                h["matching_vers"] = version
                filtered_custom_hits.append(h)
        if with_data:
            db_results = list(get_cve_data_batched(db_conn, filtered_db_hits, purl))
            return db_results + filtered_custom_hits
        else:
            custom_summaries = [
                {
                    "cve_id": h["cve_id"],
                    "vers": h["vers"],
                    "purl_prefix": h["purl_prefix"],
                    "type": h.get("type"),
                    "namespace": h.get("namespace"),
                    "name": h.get("name"),
                }
                for h in filtered_custom_hits
            ]

            return filtered_db_hits + custom_summaries
    return []


def search_by_cve(
    cve_id: str, with_data: bool = False, with_limit: int | None = None
) -> List:
    """Search by CVE"""
    db_conn, index_conn = db6.get(read_only=True)
    filter_part = "cve_id LIKE ?" if "%" in cve_id else "cve_id = ?"
    filter_part = f"{filter_part} ORDER BY cve_id DESC"
    args = [cve_id]
    if with_limit and isinstance(with_limit, int):
        filter_part = f"{filter_part} LIMIT ?"
        args.append(with_limit)

    raw_hits = exec_query(
        index_conn,
        f"SELECT DISTINCT cve_id, type, namespace, name, vers, purl_prefix FROM cve_index where {filter_part}",
        tuple(args),
    )
    db_index_hits = _process_hits(raw_hits)
    custom_candidates = _search_memory(cve_id=cve_id)
    override_keys = {(c["cve_id"], c["purl_prefix"]) for c in custom_candidates}
    filtered_db_hits = [
        hit
        for hit in db_index_hits
        if (hit["cve_id"], hit["purl_prefix"]) not in override_keys
    ]
    if with_data:
        db_results = list(get_cve_data_batched(db_conn, filtered_db_hits, cve_id))
        return db_results + custom_candidates
    else:
        custom_summaries = [
            {
                "cve_id": h["cve_id"],
                "vers": h["vers"],
                "purl_prefix": h["purl_prefix"],
                "type": h.get("type"),
                "namespace": h.get("namespace"),
                "name": h.get("name"),
            }
            for h in custom_candidates
        ]
        return filtered_db_hits + custom_summaries


def search_by_url(url: str, with_data: bool = False) -> List:
    """Search by URL"""
    purl_obj = utils.url_to_purl(url)
    if not purl_obj:
        return []
    name = purl_obj["name"]
    purl_str = (
        f"""pkg:{purl_obj["type"]}/{purl_obj["namespace"].replace("@", "%40")}/{name}"""
        if purl_obj["namespace"]
        else f"pkg:{purl_obj['type']}/{name}"
    )
    if purl_obj["version"]:
        purl_str = f"{purl_str}@{purl_obj['version']}"
    return search_by_purl_like(purl_str, with_data)


def search_by_cdx_bom(bom_file: str, with_data: bool = False) -> Generator:
    """Search by CycloneDX BOM file"""
    cdx_obj = utils.load_json(bom_file)
    for component in cdx_obj.get("components", []):
        if component.get("purl"):
            yield search_by_purl_like(component["purl"], with_data)
        if component.get("cpe"):
            yield search_by_cpe_like(component["cpe"], with_data)


def latest_malware(with_limit=20, with_data=False) -> Generator:
    """Search for latest malware with CVE ID beginning with MAL-"""
    yield search_by_cve("MAL-%", with_data=with_data, with_limit=with_limit)


def exec_query(conn, query: str, args: Tuple[str, ...]) -> list:
    res = conn.execute(query, args)
    return res.fetchall()
