"""Built-in protocol analyzers and lightweight enrichment helpers."""

from __future__ import annotations

from dataclasses import dataclass
import re

from sentra.models import CertificateMetadata, Endpoint, RawEvidence, ServiceEvidence, SoftwareComponent
from sentra.vulns.offline import normalize_product_name


HTTP_PORTS = {80, 443, 8080, 8443}
TLS_PORTS = {443, 8443, 8883}
MODBUS_PORTS = {502}
SNMP_PORTS = {161}


@dataclass(frozen=True)
class EnrichmentResult:
    service: ServiceEvidence
    endpoints: list[Endpoint]
    certificates: list[CertificateMetadata]
    software_components: list[SoftwareComponent]


def append_evidence(
    service: ServiceEvidence,
    *,
    source: str,
    kind: str,
    value: str,
    confidence: float,
    attributes: dict[str, object] | None = None,
) -> ServiceEvidence:
    evidence = list(service.evidence)
    evidence.append(
        RawEvidence(
            source=source,
            kind=kind,
            value=value,
            confidence=confidence,
            attributes=dict(sorted((attributes or {}).items())),
        )
    )
    evidence.sort(key=lambda item: (item.source, item.kind, item.value))
    return service.model_copy(update={"evidence": evidence})


class HttpBannerAnalyzer:
    name = "http-banner"

    def analyze(self, host: str, service: ServiceEvidence) -> ServiceEvidence | None:
        if service.port not in HTTP_PORTS:
            return None
        http_probe = _normalize_mapping(service.metadata.get("http_probe"))
        banner = _first_non_empty(
            http_probe.get("server_header"),
            http_probe.get("status_line"),
            service.banner,
            f"{host}:{service.port}",
        )
        product, version = _parse_product_version(banner)
        metadata = {
            **service.metadata,
            "http_probe": http_probe,
        }
        enriched = service.model_copy(
            update={
                "service": "https" if service.port in {443, 8443} else "http",
                "banner": banner,
                "product": product or service.product,
                "normalized_product": normalize_product_name(product) if product else service.normalized_product,
                "version": version or service.version,
                "confidence": _bounded_confidence(service.confidence, 0.14 if http_probe else 0.1),
                "fingerprint_source": self.name,
                "metadata": metadata,
            }
        )
        if banner:
            enriched = append_evidence(
                enriched,
                source=self.name,
                kind="http-banner",
                value=banner,
                confidence=enriched.confidence,
                attributes={"host": host, "port": service.port},
            )
        if http_probe.get("status_code"):
            enriched = append_evidence(
                enriched,
                source=self.name,
                kind="http-status",
                value=str(http_probe["status_code"]),
                confidence=enriched.confidence,
                attributes={"server_header": http_probe.get("server_header", "")},
            )
        if http_probe.get("location"):
            enriched = append_evidence(
                enriched,
                source=self.name,
                kind="http-location",
                value=str(http_probe["location"]),
                confidence=enriched.confidence,
                attributes={"host": host},
            )
        return enriched


class MqttAnalyzer:
    name = "mqtt-identifier"

    def analyze(self, host: str, service: ServiceEvidence) -> ServiceEvidence | None:
        mqtt_probe = _normalize_mapping(service.metadata.get("mqtt_probe"))
        if service.port not in {1883, 8883} and not mqtt_probe:
            return None
        product, version = _parse_product_version(service.banner or service.product or "mosquitto")
        product = product or service.product or "mosquitto"
        service_name = "mqtts" if service.port == 8883 else "mqtt"
        enriched = service.model_copy(
            update={
                "service": service_name,
                "product": product,
                "normalized_product": normalize_product_name(product),
                "version": version or service.version,
                "confidence": _bounded_confidence(service.confidence, 0.18 if mqtt_probe else 0.12),
                "fingerprint_source": self.name,
                "metadata": {
                    **service.metadata,
                    "mqtt_probe": mqtt_probe,
                },
            }
        )
        if mqtt_probe.get("connack_rc") is not None:
            enriched = append_evidence(
                enriched,
                source=self.name,
                kind="mqtt-connack",
                value=f"rc={mqtt_probe['connack_rc']}",
                confidence=enriched.confidence,
                attributes={"session_present": bool(mqtt_probe.get("session_present", False))},
            )
        return enriched


class CoapAnalyzer:
    name = "coap-identifier"

    def analyze(self, host: str, service: ServiceEvidence) -> ServiceEvidence | None:
        coap_probe = _normalize_mapping(service.metadata.get("coap_probe"))
        if service.port != 5683 and not coap_probe:
            return None
        enriched = service.model_copy(
            update={
                "service": "coap",
                "confidence": _bounded_confidence(service.confidence, 0.18 if coap_probe else 0.14),
                "fingerprint_source": self.name,
                "metadata": {
                    **service.metadata,
                    "coap_probe": coap_probe,
                },
            }
        )
        if coap_probe.get("resource_count") is not None:
            enriched = append_evidence(
                enriched,
                source=self.name,
                kind="coap-discovery",
                value=f"resources={coap_probe['resource_count']}",
                confidence=enriched.confidence,
                attributes={"preview": str(coap_probe.get("payload_preview", ""))},
            )
        return enriched


class SsdpAnalyzer:
    name = "ssdp-upnp"

    def analyze(self, host: str, service: ServiceEvidence) -> ServiceEvidence | None:
        ssdp_probe = _normalize_mapping(service.metadata.get("ssdp_probe"))
        if service.port != 1900 and not ssdp_probe:
            return None
        headers = _normalize_mapping(ssdp_probe.get("headers"))
        product = _first_non_empty(headers.get("server"), service.product, "upnp")
        enriched = service.model_copy(
            update={
                "service": "ssdp",
                "product": product,
                "normalized_product": normalize_product_name(product),
                "confidence": _bounded_confidence(service.confidence, 0.16 if ssdp_probe else 0.09),
                "fingerprint_source": self.name,
                "metadata": {
                    **service.metadata,
                    "ssdp_probe": ssdp_probe,
                },
            }
        )
        if headers.get("st"):
            enriched = append_evidence(
                enriched,
                source=self.name,
                kind="ssdp-st",
                value=str(headers["st"]),
                confidence=enriched.confidence,
                attributes={"server": str(headers.get("server", ""))},
            )
        if headers.get("location"):
            enriched = append_evidence(
                enriched,
                source=self.name,
                kind="ssdp-location",
                value=str(headers["location"]),
                confidence=enriched.confidence,
                attributes={"usn": str(headers.get("usn", ""))},
            )
        return enriched


class MdnsAnalyzer:
    name = "mdns"

    def analyze(self, host: str, service: ServiceEvidence) -> ServiceEvidence | None:
        mdns_probe = _normalize_mapping(service.metadata.get("mdns_probe"))
        if service.port != 5353 and not mdns_probe:
            return None
        service_name = "mdns"
        payload_preview = str(mdns_probe.get("payload_preview", service.banner or "")).strip()
        labels = [str(item) for item in mdns_probe.get("labels", [])]
        enriched = service.model_copy(
            update={
                "service": service_name,
                "confidence": _bounded_confidence(service.confidence, 0.16 if mdns_probe else 0.07),
                "fingerprint_source": self.name,
                "metadata": {
                    **service.metadata,
                    "mdns_probe": mdns_probe,
                },
            }
        )
        if payload_preview:
            enriched = append_evidence(
                enriched,
                source=self.name,
                kind="mdns-response",
                value=payload_preview[:120],
                confidence=enriched.confidence,
                attributes={"host": host},
            )
        if labels:
            enriched = append_evidence(
                enriched,
                source=self.name,
                kind="mdns-labels",
                value=", ".join(labels[:4]),
                confidence=enriched.confidence,
                attributes={"label_count": len(labels)},
            )
        return enriched


class TlsCertificateAnalyzer:
    name = "tls-certificate"

    def analyze(self, host: str, service: ServiceEvidence) -> ServiceEvidence | None:
        tls_probe = _normalize_mapping(service.metadata.get("tls_probe"))
        certificate = _normalize_mapping(service.metadata.get("certificate"))
        if service.port not in TLS_PORTS and not tls_probe and not certificate:
            return None
        combined_certificate = {
            "subject": certificate.get("subject") or tls_probe.get("subject") or f"CN={host}",
            "issuer": certificate.get("issuer") or tls_probe.get("issuer") or "Local Discovery Heuristic",
            "san": certificate.get("san") or tls_probe.get("san") or [host],
            "fingerprint_sha256": certificate.get("fingerprint_sha256") or tls_probe.get("fingerprint_sha256"),
        }
        metadata = {
            **service.metadata,
            "tls_probe": tls_probe,
            "certificate": combined_certificate,
        }
        enriched = service.model_copy(
            update={
                "service": "https" if service.port in {443, 8443} else "mqtts",
                "confidence": _bounded_confidence(service.confidence, 0.2 if tls_probe else 0.13),
                "fingerprint_source": self.name,
                "metadata": metadata,
            }
        )
        enriched = append_evidence(
            enriched,
            source=self.name,
            kind="tls-certificate",
            value=str(combined_certificate["subject"]),
            confidence=enriched.confidence,
            attributes={
                "issuer": str(combined_certificate["issuer"]),
                "fingerprint_sha256": str(combined_certificate.get("fingerprint_sha256", "")),
            },
        )
        return enriched


class SnmpAnalyzer:
    name = "snmp-metadata"

    def analyze(self, host: str, service: ServiceEvidence) -> ServiceEvidence | None:
        snmp_probe = _normalize_mapping(service.metadata.get("snmp_probe"))
        if service.port not in SNMP_PORTS and not snmp_probe:
            return None
        sys_descr = _first_non_empty(snmp_probe.get("sysdescr"), service.banner, service.product, "snmp-agent")
        product, version = _parse_product_version(sys_descr)
        enriched = service.model_copy(
            update={
                "service": "snmp",
                "banner": sys_descr or service.banner,
                "product": product or service.product,
                "normalized_product": normalize_product_name(product) if product else service.normalized_product,
                "version": version or service.version,
                "confidence": _bounded_confidence(service.confidence, 0.19 if snmp_probe else 0.11),
                "fingerprint_source": self.name,
                "metadata": {**service.metadata, "snmp_probe": snmp_probe},
            }
        )
        if sys_descr:
            enriched = append_evidence(
                enriched,
                source=self.name,
                kind="snmp-sysdescr",
                value=sys_descr,
                confidence=enriched.confidence,
                attributes={"sysname": str(snmp_probe.get("sysname", ""))},
            )
        return enriched


class ModbusAnalyzer:
    name = "modbus-tcp"

    def analyze(self, host: str, service: ServiceEvidence) -> ServiceEvidence | None:
        modbus_probe = _normalize_mapping(service.metadata.get("modbus_probe"))
        if service.port not in MODBUS_PORTS and not modbus_probe:
            return None
        identity = _first_non_empty(modbus_probe.get("device_id"), service.banner, "modbus-device")
        product, version = _parse_product_version(identity)
        enriched = service.model_copy(
            update={
                "service": "modbus",
                "banner": identity or service.banner,
                "product": product or service.product,
                "normalized_product": normalize_product_name(product) if product else service.normalized_product,
                "version": version or service.version,
                "confidence": _bounded_confidence(service.confidence, 0.18 if modbus_probe else 0.1),
                "fingerprint_source": self.name,
                "metadata": {**service.metadata, "modbus_probe": modbus_probe},
            }
        )
        if identity:
            enriched = append_evidence(
                enriched,
                source=self.name,
                kind="modbus-identity",
                value=identity,
                confidence=enriched.confidence,
                attributes={"unit_id": int(modbus_probe.get("unit_id", 0)) if str(modbus_probe.get("unit_id", "")).isdigit() else 0},
            )
        return enriched


def build_service_enrichment(host: str, service: ServiceEvidence) -> EnrichmentResult:
    endpoints: list[Endpoint] = []
    certificates: list[CertificateMetadata] = []
    software: list[SoftwareComponent] = []

    if service.service in {"http", "https"}:
        scheme = "https" if service.service == "https" else "http"
        endpoints.append(Endpoint(url=f"{scheme}://{host}:{service.port}/", confidence=service.confidence))
    if service.service == "mqtt":
        endpoints.append(Endpoint(url=f"mqtt://{host}:{service.port}", kind="broker", confidence=service.confidence))
    if service.service == "mqtts":
        endpoints.append(Endpoint(url=f"mqtts://{host}:{service.port}", kind="broker", confidence=service.confidence))
    if service.service == "coap":
        endpoints.append(Endpoint(url=f"coap://{host}:{service.port}/.well-known/core", kind="resource", confidence=service.confidence))
    if service.service == "ssdp":
        location = _normalize_mapping(service.metadata.get("ssdp_probe")).get("headers", {})
        if isinstance(location, dict) and location.get("location"):
            endpoints.append(
                Endpoint(url=str(location["location"]), kind="descriptor", confidence=service.confidence)
            )
    if service.service == "modbus":
        endpoints.append(Endpoint(url=f"modbus://{host}:{service.port}", kind="endpoint", confidence=service.confidence))
    if service.metadata.get("certificate"):
        certificate = _normalize_mapping(service.metadata["certificate"])
        certificates.append(
            CertificateMetadata(
                subject=str(certificate.get("subject", "")) or None,
                issuer=str(certificate.get("issuer", "")) or None,
                common_name=host,
                san=[str(item) for item in certificate.get("san", [])],
                fingerprint_sha256=str(certificate.get("fingerprint_sha256", "")) or None,
                confidence=service.confidence,
            )
        )
    if service.product or service.normalized_product:
        software.append(
            SoftwareComponent(
                name=service.product or service.service,
                normalized_name=service.normalized_product or normalize_product_name(service.product or service.service),
                version=service.version,
                source=service.fingerprint_source,
                confidence=service.confidence,
            )
        )

    return EnrichmentResult(service=service, endpoints=endpoints, certificates=certificates, software_components=software)


def _parse_product_version(banner: str) -> tuple[str | None, str | None]:
    tokens = banner.replace("/", " ").split()
    if not tokens:
        return None, None
    product = tokens[0]
    version = tokens[1] if len(tokens) > 1 and any(char.isdigit() for char in tokens[1]) else None
    return product, version


def parse_http_response(payload: bytes) -> dict[str, object]:
    text = payload.decode("utf-8", errors="ignore")
    lines = [line.strip() for line in text.splitlines() if line.strip()]
    if not lines:
        return {}
    status_line = lines[0]
    status_match = re.search(r"HTTP/\d+\.\d+\s+(\d+)", status_line)
    headers: dict[str, str] = {}
    for line in lines[1:]:
        if ":" in line:
            key, value = line.split(":", 1)
            headers[key.strip().lower()] = value.strip()
    return {
        "status_line": status_line,
        "status_code": int(status_match.group(1)) if status_match else None,
        "server_header": headers.get("server"),
        "location": headers.get("location"),
        "powered_by": headers.get("x-powered-by"),
        "headers": dict(sorted(headers.items())),
    }


def parse_mqtt_connack(payload: bytes) -> dict[str, object]:
    if len(payload) < 4 or payload[0] != 0x20:
        return {}
    return {
        "session_present": bool(payload[2] & 0x01),
        "connack_rc": int(payload[3]),
    }


def parse_coap_payload(payload: bytes) -> dict[str, object]:
    text = payload.decode("utf-8", errors="ignore")
    resources = sorted({item.strip() for item in text.replace(",", "\n").splitlines() if "/" in item})
    return {
        "resource_count": len(resources),
        "payload_preview": text[:160],
        "resources": resources[:10],
    }


def parse_ssdp_response(payload: bytes) -> dict[str, object]:
    text = payload.decode("utf-8", errors="ignore")
    headers: dict[str, str] = {}
    for line in text.splitlines():
        if ":" in line:
            key, value = line.split(":", 1)
            headers[key.strip().lower()] = value.strip()
    return {
        "payload_preview": text[:160],
        "headers": dict(sorted(headers.items())),
    }


def parse_mdns_response(payload: bytes) -> dict[str, object]:
    text = payload.decode("utf-8", errors="ignore")
    labels = sorted({chunk for chunk in re.split(r"[\x00\s]+", text) if "." in chunk or "_" in chunk})[:10]
    return {
        "payload_preview": text[:160],
        "labels": labels,
    }


def parse_snmp_payload(payload: bytes) -> dict[str, object]:
    text = payload.decode("utf-8", errors="ignore")
    printable = " ".join(token for token in re.split(r"[\x00-\x1f]+", text) if token.strip())
    tokens = [token.strip() for token in re.split(r"[;,\n]+", printable) if token.strip()]
    sysdescr = next((token for token in tokens if "sysdescr" in token.lower()), printable[:160]).strip()
    sysname = next((token for token in tokens if "sysname" in token.lower()), "")
    return {
        "payload_preview": printable[:160],
        "sysdescr": sysdescr.replace("sysDescr", "").replace("=", " ").strip() if sysdescr else "",
        "sysname": sysname.replace("sysName", "").replace("=", " ").strip() if sysname else "",
    }


def parse_modbus_response(payload: bytes) -> dict[str, object]:
    if len(payload) < 8:
        return {}
    unit_id = payload[6]
    function_code = payload[7]
    ascii_tail = "".join(chr(byte) if 32 <= byte <= 126 else " " for byte in payload[8:]).strip()
    return {
        "unit_id": unit_id,
        "function_code": function_code,
        "device_id": ascii_tail[:120],
    }


def _normalize_mapping(value: object) -> dict[str, object]:
    if isinstance(value, dict):
        return dict(sorted(value.items()))
    return {}


def _first_non_empty(*values: object) -> str:
    for value in values:
        if isinstance(value, str) and value.strip():
            return value.strip()
    return ""


def _bounded_confidence(base: float, bonus: float) -> float:
    return round(min(0.99, max(base, base + bonus)), 2)

