from collections import Counter
from collections.abc import Iterable
from datetime import datetime
from typing import Any, TypedDict

from django.utils.translation import gettext_lazy as _

from octopoes.models.ooi.dns.zone import Hostname
from octopoes.models.ooi.findings import Finding, FindingType, KATFindingType, RiskLevelSeverity
from octopoes.models.ooi.network import IPAddressV4, IPAddressV6
from reports.report_types.definitions import Report


class FindingsData(TypedDict):
    finding_types: list[FindingType]
    findings: list[Finding]


class VulnerabilityReport(Report):
    id = "vulnerability-report"
    name = _("Vulnerability Report")
    description: str = _("Vulnerabilities found are grouped for each system.")
    plugins = {
        "required": {"dns-records", "nmap", "webpage-analysis"},
        "optional": {"nmap-udp", "nmap-ports", "shodan"},
    }
    input_ooi_types = {Hostname, IPAddressV4, IPAddressV6}
    template_path = "vulnerability_report/report.html"
    label_style = "5-light"

    def collect_data(self, input_oois: Iterable[str], valid_time: datetime) -> dict[str, dict[str, Any]]:
        result = {}
        all_findings_for_all_oois = self.get_findings(input_oois, valid_time)
        readable_hostnames_by_input_ooi = Report.hostnames_to_human_readable(self.to_hostnames(input_oois, valid_time))
        # TODO: these caches should not be necessary if we rewrite the inner loops or use new bulk queries there as well
        history_cache = {}
        origin_cache = {}

        for input_ooi in input_oois:
            data = {}
            all_findings = all_findings_for_all_oois.get(input_ooi, {})

            for ip, findings_data in all_findings.items():
                vulnerabilities: dict[str, dict[str, Any]] = {}
                total_criticals = 0
                finding_types_ids = [finding_type.id for finding_type in findings_data["finding_types"]]
                occurrences = Counter(finding_types_ids)
                total_findings = sum(occurrences.values())
                terms = list(occurrences)
                recommendations = []
                first_seen = ""

                # remove duplicate findings and finding types
                finding_types = list(set(findings_data["finding_types"]))
                findings = list(set(findings_data["findings"]))

                for finding_type in finding_types:
                    filtered_findings = {}

                    for finding in findings:
                        if finding.finding_type.tokenized.id != finding_type.id:
                            continue

                        if finding.primary_key not in history_cache:
                            history_cache[finding.reference] = self.octopoes_api_connector.get_history(
                                finding.reference
                            )

                        time_history = [transaction.valid_time for transaction in history_cache[finding.reference]]

                        if time_history:
                            first_seen = str(time_history[0])

                        if finding.reference not in origin_cache:
                            origin_cache[finding.reference] = self.octopoes_api_connector.list_origins(
                                result=finding.ooi, valid_time=valid_time
                            )

                        origins = origin_cache[finding.reference]
                        sources = ", ".join([origin.method for origin in origins])

                        evidence = origins[0].task_id if origins else "-"

                        filtered_findings[finding.human_readable] = {
                            str(_("Source")): sources,
                            str(_("First seen")): first_seen,
                            str(_("Last seen")): "-",
                            str(_("Evidence")): evidence,
                        }

                    cvss: dict[str, str | float | None] = {"class": "-", "score": None, "risk_level": "-"}

                    if finding_type.risk_severity:
                        cvss = {
                            "class": str(finding_type.risk_severity.value).lower(),
                            "score": finding_type.risk_score,
                            "risk_level": " ".join(
                                [str(finding_type.risk_severity.value).title(), str(finding_type.risk_score)]
                            ),
                        }

                    vulnerabilities[finding_type.id] = {
                        "cvss": cvss,
                        "occurrences": occurrences[finding_type.id],
                        "description": finding_type.description or "-",
                        "advice": finding_type.recommendation or "-",
                        "findings": filtered_findings,
                    }

                    if finding_type.risk_severity == RiskLevelSeverity.CRITICAL:
                        total_criticals += 1

                    if finding_type.recommendation:
                        recommendations.append(finding_type.recommendation)

                sorted_vulnerabilities = sorted(
                    vulnerabilities.items(), key=lambda x: x[1].get("cvss", {}).get("score", 0) or 0, reverse=True
                )

                data[ip] = {
                    "hostnames": readable_hostnames_by_input_ooi[input_ooi],
                    "vulnerabilities": dict(sorted_vulnerabilities),
                    "summary": {
                        "total_findings": total_findings,
                        "total_criticals": total_criticals,
                        "terms": terms,
                        "recommendations": recommendations,
                    },
                }

            result[input_ooi] = data

        return result

    def get_findings(self, input_oois: Iterable[str], valid_time: datetime) -> dict:
        ips_by_input_ooi = self.to_ips(input_oois, valid_time)
        all_ips = list({ip for key, ips in ips_by_input_ooi.items() for ip in ips})

        port_finding_types = self.group_finding_types_by_source(
            self.octopoes_api_connector.query_many(
                "IPAddress.<address[is IPPort].<ooi[is Finding].finding_type", valid_time, all_ips
            )
        )
        software_instance_finding_types = self.group_finding_types_by_source(
            self.octopoes_api_connector.query_many(
                "IPAddress.<address [is ResolvedHostname].hostname.<netloc [is HostnameHTTPURL]"
                ".<ooi [is SoftwareInstance].<ooi [is Finding].finding_type",
                valid_time,
                all_ips,
            )
        )
        port_findings = self.group_finding_types_by_source(
            self.octopoes_api_connector.query_many(
                "IPAddress.<address[is IPPort].<ooi[is Finding]", valid_time, all_ips
            )
        )
        software_instance_findings = self.group_finding_types_by_source(
            self.octopoes_api_connector.query_many(
                "IPAddress.<address [is ResolvedHostname].hostname.<netloc [is HostnameHTTPURL]"
                ".<ooi [is SoftwareInstance].<ooi [is Finding]",
                valid_time,
                all_ips,
            )
        )

        result = {}
        for input_ooi, ip_references in ips_by_input_ooi.items():
            findings_data = {}
            all_finding_types = []
            aggregated_findings = []

            for ip in ip_references:
                found_finding_types = port_finding_types.get(ip, []) + software_instance_finding_types.get(ip, [])

                for finding_type in found_finding_types:
                    if finding_type.reference.class_type != KATFindingType:
                        all_finding_types.append(finding_type)

                findings = port_findings.get(ip, []) + software_instance_findings.get(ip, [])

                for finding in findings:
                    if finding.finding_type.class_type != KATFindingType:
                        aggregated_findings.append(finding)

                findings_data[ip] = {"finding_types": all_finding_types, "findings": aggregated_findings}

            result[input_ooi] = findings_data

        return result
