import logging
from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List, Optional, Set

from sec_certs.sample.certificate import Certificate
from sec_certs.serialization.json import ComplexSerializableType


class DependencyType(Enum):
    DIRECT = "direct"
    INDIRECT = "indirect"


@dataclass
class DependencyCVE(ComplexSerializableType):
    direct_dependency_cves: Optional[Set[str]] = field(default=None)
    indirect_dependency_cves: Optional[Set[str]] = field(default=None)


Certificates = Dict[str, Certificate]
Vulnerabilities = Dict[str, Dict[str, Optional[Set[str]]]]


class DependencyVulnerabilityFinder:
    """
    The class assigns vulnerabilities to each certificate instance caused by dependencies among certificate instances.
    Adheres to sklearn BaseEstimator interface.
    """

    def __init__(self):
        self.vulnerabilities: Vulnerabilities = {}
        self.certificates: Certificates = {}

    def _overwrite_previous_state(self, certificates: Certificates) -> None:
        self.vulnerabilities = {}
        self.certificates = certificates

    def _get_dataset_cert_ids_occurrences(self) -> Dict[str, int]:
        cert_id_occurrences: Dict[str, int] = {}

        for dgst in self.certificates:
            cert_id = self.certificates[dgst].heuristics.cert_id

            if cert_id is None:
                continue

            cert_id_occurrences[cert_id] = cert_id_occurrences.get(cert_id, 0) + 1

        return cert_id_occurrences

    def _get_cert_dependency_cves(self, dgst: str, dependency_type: DependencyType) -> Optional[Set[str]]:
        dependency_type_dict = {
            DependencyType.DIRECT: self.certificates[dgst].heuristics.report_references.directly_referenced_by,
            DependencyType.INDIRECT: self.certificates[dgst].heuristics.report_references.indirectly_referenced_by,
        }

        references = dependency_type_dict[dependency_type]

        if not references:
            return None

        vulnerabilities = set()
        dataset_cert_id_occurrences = self._get_dataset_cert_ids_occurrences()

        for cert_id in references:
            if cert_id is None:
                continue

            cert_id_occurrences = dataset_cert_id_occurrences.get(cert_id)

            if cert_id_occurrences is None or cert_id_occurrences >= 2:
                continue

            for dgst in self.certificates:
                cert_obj = self.certificates[dgst]
                cert_obj_cves = cert_obj.heuristics.related_cves

                if cert_obj.heuristics.cert_id == cert_id and cert_obj_cves:
                    vulnerabilities.update(cert_obj_cves)

        return vulnerabilities if vulnerabilities else None

    def fit(self, certificates: Certificates) -> Vulnerabilities:
        """
        Method assigns each certificate vulnerabilities caused by dependencies among certificates

        :param Certificates certificates: Dictionary of certificates with digests
        :return Vulnerabilities: Dictionary of vulnerabilities of certificate instances
        """
        self._overwrite_previous_state(certificates)

        cert_id_occurrences = self._get_dataset_cert_ids_occurrences()
        thrown_away_cert_counter = 0

        for dgst in self.certificates:
            cert_id = self.certificates[dgst].heuristics.cert_id

            if cert_id is None:
                continue

            if cert_id_occurrences[cert_id] >= 2:
                thrown_away_cert_counter += 1
                continue

            self.vulnerabilities[dgst] = {}
            self.vulnerabilities[dgst][DependencyType.DIRECT.value] = self._get_cert_dependency_cves(
                dgst, DependencyType.DIRECT
            )
            self.vulnerabilities[dgst][DependencyType.INDIRECT.value] = self._get_cert_dependency_cves(
                dgst, DependencyType.INDIRECT
            )

        if thrown_away_cert_counter > 0:
            logging.warning("There were total of %s certificates skipped due to duplicity", thrown_away_cert_counter)

        return self.vulnerabilities

    def predict_single_cert(self, dgst: str) -> DependencyCVE:
        """
        Method returns vulnerabilities for certificate digest

        :param str dgst: Digest of certificate
        :return DependencyCVE: DependencyCVE object of certificate
        """
        if not self.vulnerabilities.get(dgst):
            return DependencyCVE(direct_dependency_cves=None, indirect_dependency_cves=None)

        return DependencyCVE(
            self.vulnerabilities[dgst][DependencyType.DIRECT.value],
            self.vulnerabilities[dgst][DependencyType.INDIRECT.value],
        )

    def predict(self, dgst_list: List[str]) -> Dict[str, DependencyCVE]:
        """
        Method returns vulnerabilities for a list of certificate digests

        :param List[str] dgst_list: list of certificate digests
        :return Dict[str, DependencyCVE]: Dictionary of DependencyCVE objects for specified certificate digests
        """
        cert_vulnerabilities = {}

        for dgst in dgst_list:
            cert_vulnerabilities[dgst] = self.predict_single_cert(dgst)

        return cert_vulnerabilities
