from typing import List

from sator_core.models.patch import PatchReferences
from sator_core.models.vulnerability.references import VulnerabilityReferences

from sator_core.ports.driven.persistence.storage import StoragePersistencePort
from sator_core.ports.driven.repositories.vulnerability import VulnerabilityRepositoryPort
from sator_core.ports.driving.resolution.references.vulnerability import VulnerabilityReferencesResolutionPort


class VulnerabilityReferencesResolution(VulnerabilityReferencesResolutionPort):
    def __init__(
            self, vulnerability_repositories: List[VulnerabilityRepositoryPort], storage_port: StoragePersistencePort
    ):
        self.vulnerability_repositories = vulnerability_repositories
        self.storage_port = storage_port

    def search_vulnerability_references(self, vulnerability_id: str) -> VulnerabilityReferences | None:
        vulnerability_references = self.storage_port.load(VulnerabilityReferences, vulnerability_id)

        if not vulnerability_references:
            vulnerability_references = self._get_vulnerability_references(vulnerability_id)

            if vulnerability_references:
                self.storage_port.save(vulnerability_references, vulnerability_id)

        # Vulnerability repositories can also provide patch references, so we need to check for them and save if found.
        patch_references = self.storage_port.load(PatchReferences, vulnerability_id)

        if not patch_references:
            patch_references = self._get_patch_references(vulnerability_id)

            if patch_references:
                self.storage_port.save(patch_references, vulnerability_id)

        return vulnerability_references

    def _get_vulnerability_references(self, vulnerability_id: str) -> VulnerabilityReferences | None:
        vulnerability_references = VulnerabilityReferences(
            vulnerability_id=vulnerability_id
        )

        for port in self.vulnerability_repositories:
            references = port.get_references(vulnerability_id)

            if references:
                vulnerability_references.extend(references)

        return vulnerability_references if len(vulnerability_references) > 0 else None

    def _get_patch_references(self, vulnerability_id: str) -> PatchReferences | None:
        patch_references = PatchReferences(
            vulnerability_id=vulnerability_id
        )

        for port in self.vulnerability_repositories:
            references = port.get_patch_references(vulnerability_id)

            if references:
                patch_references.extend(references)

        return patch_references if len(patch_references) > 0 else None
