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

import apsw
import orjson

from vdb.lib import db6, utils
from vdb.lib.config import PLACEHOLDER_EXCLUDE_VERSION, PLACEHOLDER_FIX_VERSION
from vdb.lib.cve_model import CVE, CVE1, Product, Versions, Status
from vdb.lib.utils import load_json, vers_compare


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

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_unaffected(source_data, vers):
    if source_data and source_data.root.containers:
        products: List[Product] = source_data.root.containers.cna.affected.root
        for p in products:
            versions: List[Versions] = p.versions
            if versions:
                for ver in versions:
                    if ver.status == Status.unaffected:
                        if ver.version:
                            return ver.version.root
    if "|" in vers:
        vers = vers.split("|")[-1]
        if "!=" in vers or "<=" not in vers:
            return vers.replace("<", "").replace("!=", "")
    elif "/<" in vers and "/<=" not in vers:
        return vers.split("/<")[-1]
    elif vers.endswith(f"<{PLACEHOLDER_EXCLUDE_VERSION}"):
        return PLACEHOLDER_EXCLUDE_VERSION
    elif vers.endswith(f"<{PLACEHOLDER_FIX_VERSION}"):
        return PLACEHOLDER_FIX_VERSION
    return ""


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 = 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 = 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, version,)
        raw_hits = exec_query(
            index_conn,
            "SELECT DISTINCT cve_id, type, namespace, name, vers, purl_prefix FROM cve_index where purl_prefix = ? AND vers_compare(?, vers)",
            args,
        )
        filtered_list = _process_hits(raw_hits)
        if with_data:
            return list(get_cve_data_batched(db_conn, filtered_list, purl))
        return filtered_list
    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)
    args = tuple(args)
    raw_hits = exec_query(
        index_conn,
        f"SELECT DISTINCT cve_id, type, namespace, name, vers, purl_prefix FROM cve_index where {filter_part}",
        args,
    )
    filtered_list = _process_hits(raw_hits)
    if with_data:
        return list(get_cve_data_batched(db_conn, filtered_list, cve_id))
    return filtered_list


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 = 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()
