from pydantic import HttpUrl

from vulnsearch.models.vulnerability import Package, Vulnerability


def consolidate_vulnerabilities(
    name: str, version: str, raw_vulnerabilities: list[Vulnerability]
) -> Package | None:
    """
    Unify vulnerabilities from multiple sources, pick the best representative per group,
    enrich it with missing fields from the group, and build a Package.
    """
    if not raw_vulnerabilities:
        return None

    normalized_vulnerabilities: list[Vulnerability] = []

    for group in unified_vulnerabilities(raw_vulnerabilities):
        best = pick_best_representative(group)
        normalized_vulnerabilities.append(enrich_from_group(best, group))

    return Package(
        name=name,
        version=version,
        vulnerabilities=normalized_vulnerabilities,
    )


def unified_vulnerabilities(
    vulnerabilities: list[Vulnerability],
) -> list[list[Vulnerability]]:
    """
    Group vulnerabilities that refer to the same underlying issue (CVE/alias)
    into unified vulnerability groups.
    """
    vulns_dict = {vuln.id: vuln for vuln in vulnerabilities}

    seen_ids: set[str] = set()
    results: list[list[Vulnerability]] = []

    if len(vulnerabilities) == 0:
        return results

    for vuln in vulnerabilities:
        if vuln.id in seen_ids:
            continue

        group: list[Vulnerability] = [vuln]
        seen_ids.add(vuln.id)

        for alias in vuln.aliases:
            if alias in seen_ids:
                continue

            # Find locally
            local = vulns_dict.get(alias)
            if local is not None:
                group.append(local)
                seen_ids.add(alias)
                continue

        results.append(group)
    return results


def pick_best_representative(group: list[Vulnerability]) -> Vulnerability:
    """Pick a single representative vulnerability from a unified group."""

    # absolute priority: NVD CVE
    for v in group:
        if v.source == "NVD" and v.id.startswith("CVE-"):
            return v

    # any CVE
    for v in group:
        if v.id.startswith("CVE-"):
            return v

    # most filled
    for v in group:
        if v.summary and v.affected and v.references and v.severity:
            return v

    return group[0]


def enrich_from_group(
    base: Vulnerability,
    group: list[Vulnerability],
) -> Vulnerability:
    """
    Enrich Vulnerability object with missing fields (summary, severity, references)
    from other vulnerabilities in the same group.
    """
    # summary
    if not base.summary:
        for v in group:
            if v.summary:
                base.summary = v.summary
                break

    # severity
    if base.severity.score == 0:
        for v in group:
            if v.severity.severity != 0:
                base.severity = v.severity
                break

    # references
    if not base.references:
        refs: set[HttpUrl] = set()
        for v in group:
            if v.references:
                refs.update(v.references)
        if refs:
            base.references = refs

    return base
