"""Safe firmware extraction and heuristic analysis."""

from __future__ import annotations

import hashlib
import re
from pathlib import Path

from sentra.firmware.extractors import ExtractionRequest, extract_firmware_bundle
from sentra.models import (
    CertificateFinding,
    CertificateMetadata,
    ComponentEvidence,
    ConfigHint,
    EvidenceFinding,
    ExtractedArtifact,
    ExtractionMetadata,
    FirmwareEvidence,
    InferenceHint,
    ObservedFact,
    ServiceHint,
    SoftwareComponent,
)
from sentra.vulns import enrich_firmware_vulnerabilities
from sentra.vulns.offline import normalize_product_name

_COMPONENT_TOKENS = {
    "busybox",
    "curl",
    "dnsmasq",
    "dropbear",
    "goahead",
    "lighttpd",
    "mosquitto",
    "nginx",
    "openssl",
    "uhttpd",
}
_VENDOR_HINTS = {
    "acme": "Acme Vision",
    "atlas": "Atlas",
    "meridian": "Meridian",
    "northwind": "Northwind Devices",
    "sentinel": "Sentinel",
}
_PRODUCT_HINTS = {
    "camera": "Camera",
    "dvr": "DVR",
    "gateway": "Gateway",
    "nvr": "NVR",
    "plc": "PLC",
    "printer": "Printer",
    "router": "Router",
}
_SERVICE_TOKEN_MAP = {
    "boa": "http",
    "coap": "coap",
    "dropbear": "ssh",
    "httpd": "http",
    "lighttpd": "http",
    "modbus": "modbus",
    "mosquitto": "mqtt",
    "nginx": "http",
    "rtsp": "rtsp",
    "snmp": "snmp",
    "sshd": "ssh",
    "telnet": "telnet",
    "uhttpd": "http",
}
_CONFIG_SUFFIXES = {".conf", ".cfg", ".env", ".ini", ".json", ".properties", ".xml", ".yaml", ".yml"}
_CERT_SUFFIXES = {".pem", ".crt", ".cer"}
_CONFIG_KEYS = {
    "broker": ("service", 0.72),
    "host": ("service", 0.68),
    "listen": ("service", 0.7),
    "mqtt_port": ("service", 0.78),
    "password": ("credentials", 0.9),
    "port": ("service", 0.66),
    "server": ("service", 0.65),
    "tls_cert": ("certificate", 0.74),
    "tls_key": ("certificate", 0.7),
    "username": ("credentials", 0.88),
    "web_root": ("service", 0.64),
}
_LIKELY_SERVICE_PORTS = {
    "21": "ftp",
    "22": "ssh",
    "23": "telnet",
    "53": "dns",
    "80": "http",
    "161": "snmp",
    "443": "https",
    "554": "rtsp",
    "1883": "mqtt",
    "5683": "coap",
    "502": "modbus",
}
_VERSION_PATTERN = re.compile(r"\b(?P<name>[A-Za-z][A-Za-z0-9._+-]{2,32})[ /_-]v?(?P<version>\d+(?:\.\d+){1,3})\b")
_KEY_VALUE_PATTERN = re.compile(r"^\s*([A-Za-z0-9_.-]{2,64})\s*[:=]\s*(.+?)\s*$")
_CERT_VALUE_PATTERN = re.compile(r"(?P<key>CN|O|OU|issuer)=([^\s,;/]+)", re.IGNORECASE)


def analyze_firmware(path: str, workspace: Path | None = None) -> FirmwareEvidence:
    firmware_path = Path(path)
    if firmware_path.exists() and firmware_path.is_dir():
        raise ValueError(f"Firmware path must be a file, not a directory: {firmware_path}")

    payload = firmware_path.read_bytes() if firmware_path.exists() else path.encode("utf-8")
    digest = hashlib.sha256(payload).hexdigest()
    decoded = payload.decode("utf-8", errors="ignore")
    strings = _extract_strings(decoded)

    extraction = _extract_bundle(firmware_path, digest, workspace) if firmware_path.exists() else None
    extracted_artifacts = extraction.artifacts if extraction else []
    extraction_metadata = _build_extraction_metadata(extraction)
    extracted_metadata = extraction_metadata.model_dump(mode="python")

    extracted_root = extraction.extracted_root if extraction and extraction.extracted_root else None
    config_hints = _collect_config_hints(extracted_root)
    certificate_findings = _collect_certificate_findings(extracted_root)
    certificates = _to_certificate_metadata(certificate_findings)
    service_hints = _collect_service_hints(decoded, extracted_root, config_hints)
    software_components, component_evidence = _collect_components(strings, extracted_root, extracted_artifacts)
    suspected_vendors, suspected_products = _collect_vendor_product_hints(decoded, extracted_root, config_hints)

    observed_facts = _build_observed_facts(extraction_metadata, config_hints, certificates, extracted_artifacts, software_components)
    inference_hints = _build_inference_hints(suspected_vendors, suspected_products, service_hints, software_components)
    findings = _build_findings(decoded, config_hints, certificate_findings, service_hints, software_components)

    artifact = FirmwareEvidence(
        path=str(firmware_path),
        sha256=digest,
        size_bytes=len(payload),
        extracted_strings=strings,
        suspected_vendors=suspected_vendors,
        suspected_products=suspected_products,
        software_components=software_components,
        component_evidence=component_evidence,
        extracted_artifacts=extracted_artifacts,
        extracted_metadata=extracted_metadata,
        extraction_metadata=extraction_metadata,
        certificates=certificates,
        certificate_findings=certificate_findings,
        config_hints=config_hints,
        service_hints=service_hints,
        observed_facts=observed_facts,
        inference_hints=inference_hints,
        findings=findings,
    )
    return enrich_firmware_vulnerabilities(artifact, workspace=workspace)


def _extract_bundle(firmware_path: Path, digest: str, workspace: Path | None):
    if workspace is None:
        return None
    destination_root = workspace / ".sentra-fw" / digest[:12]
    request = ExtractionRequest(source_path=firmware_path, destination_root=destination_root)
    return extract_firmware_bundle(request)


def _build_extraction_metadata(extraction) -> ExtractionMetadata:  # type: ignore[no-untyped-def]
    if extraction is None:
        return ExtractionMetadata()
    return ExtractionMetadata(
        provider=extraction.provider,
        detected_type=str(extraction.metadata.get("detected_type", "unknown")),
        extracted_root=str(extraction.extracted_root) if extraction.extracted_root else "",
        member_count=int(extraction.metadata.get("member_count", 0)),
        total_bytes=int(extraction.metadata.get("total_bytes", 0)),
        max_files=int(extraction.metadata.get("max_files", 0)),
        max_total_bytes=int(extraction.metadata.get("max_total_bytes", 0)),
        truncated=bool(extraction.metadata.get("truncated", False)),
        depth_limit_hit=bool(extraction.metadata.get("depth_limit_hit", False)),
        nested_archives_seen=int(extraction.metadata.get("nested_archives_seen", 0)),
        extraction_notes=sorted(str(item) for item in extraction.metadata.get("extraction_notes", [])),
    )


def _extract_strings(decoded: str) -> list[str]:
    candidates = [token.strip() for token in decoded.replace("\x00", " ").split()]
    strings = [token for token in candidates if len(token) >= 4]
    return strings[:32]


def _collect_config_hints(extracted_root: Path | None) -> list[ConfigHint]:
    if extracted_root is None:
        return []

    hints: dict[tuple[str, str, str], ConfigHint] = {}
    for path in _iter_extracted_files(extracted_root):
        if path.suffix.lower() not in _CONFIG_SUFFIXES:
            continue
        content = path.read_text(encoding="utf-8", errors="ignore")
        for line in content.splitlines()[:40]:
            match = _KEY_VALUE_PATTERN.match(line)
            if not match:
                continue
            key = match.group(1).strip().lower()
            value = match.group(2).strip().strip("\"'")
            category, confidence = _CONFIG_KEYS.get(key, ("config", 0.58))
            normalized_path = str(path.relative_to(extracted_root)).replace("\\", "/")
            hint = ConfigHint(
                path=normalized_path,
                key=key,
                value_preview=value[:80],
                category=category,
                confidence=confidence,
                inferred=False,
            )
            hints[(hint.path, hint.key, hint.value_preview)] = hint
    return sorted(hints.values(), key=lambda item: (item.path, item.key, item.value_preview))[:32]


def _collect_certificate_findings(extracted_root: Path | None) -> list[CertificateFinding]:
    if extracted_root is None:
        return []

    findings: list[CertificateFinding] = []
    for path in _iter_extracted_files(extracted_root):
        content = path.read_text(encoding="utf-8", errors="ignore")
        if path.suffix.lower() not in _CERT_SUFFIXES and "BEGIN CERTIFICATE" not in content:
            continue
        common_name = _extract_certificate_value(content, "CN")
        issuer = _extract_certificate_value(content, "issuer")
        subject = common_name or path.name
        findings.append(
            CertificateFinding(
                path=str(path.relative_to(extracted_root)).replace("\\", "/"),
                subject=subject,
                issuer=issuer or "Firmware Embedded Certificate",
                common_name=common_name or path.stem,
                fingerprint_sha256=hashlib.sha256(content.encode("utf-8", errors="ignore")).hexdigest(),
                confidence=0.83,
                inferred=False,
                rationale="Certificate material was directly observed in extracted firmware content.",
            )
        )
    findings.sort(key=lambda item: (item.path, item.common_name or "", item.subject or ""))
    return findings[:16]


def _to_certificate_metadata(findings: list[CertificateFinding]) -> list[CertificateMetadata]:
    return [
        CertificateMetadata(
            subject=item.subject,
            issuer=item.issuer,
            common_name=item.common_name,
            san=item.san,
            fingerprint_sha256=item.fingerprint_sha256,
            confidence=item.confidence,
        )
        for item in findings
    ]


def _collect_service_hints(
    decoded: str,
    extracted_root: Path | None,
    config_hints: list[ConfigHint],
) -> list[ServiceHint]:
    hints: dict[tuple[str, str], ServiceHint] = {}
    sources: list[tuple[str, str]] = [("firmware.bin", decoded)]
    if extracted_root is not None:
        for path in _iter_extracted_files(extracted_root):
            sources.append(
                (
                    str(path.relative_to(extracted_root)).replace("\\", "/"),
                    path.read_text(encoding="utf-8", errors="ignore"),
                )
            )

    for source_path, content in sources:
        lowered = content.lower()
        for token, service in _SERVICE_TOKEN_MAP.items():
            if token not in lowered:
                continue
            hints[(service, source_path)] = ServiceHint(
                service=service,
                source_path=source_path,
                rationale=f"Detected service token '{token}' in analyzed firmware content.",
                source_kind="content-token",
                confidence=0.8 if source_path != "firmware.bin" else 0.68,
                inferred=True,
            )

    for hint in config_hints:
        service = _infer_service_from_config(hint)
        if service is None:
            continue
        hints[(service, hint.path)] = ServiceHint(
            service=service,
            source_path=hint.path,
            rationale=f"Configuration key '{hint.key}' suggests a {service} service path.",
            source_kind="config-hint",
            confidence=max(hint.confidence - 0.05, 0.55),
            inferred=True,
        )

    return sorted(hints.values(), key=lambda item: (item.service, item.source_path, item.source_kind))[:24]


def _infer_service_from_config(hint: ConfigHint) -> str | None:
    preview = hint.value_preview.lower()
    if hint.key in {"broker", "mqtt_port"} or "mqtt" in preview:
        return "mqtt"
    if hint.key in {"tls_cert", "tls_key"}:
        return "https"
    if hint.key == "port":
        return _LIKELY_SERVICE_PORTS.get(preview)
    if hint.key in {"listen", "web_root"} and ("80" in preview or "http" in preview or "www" in preview):
        return "http"
    if "telnet" in preview:
        return "telnet"
    if "rtsp" in preview:
        return "rtsp"
    if "modbus" in preview:
        return "modbus"
    return None


def _collect_components(
    strings: list[str],
    extracted_root: Path | None,
    extracted_artifacts: list[ExtractedArtifact],
) -> tuple[list[SoftwareComponent], list[ComponentEvidence]]:
    components: dict[str, SoftwareComponent] = {}
    evidence: dict[tuple[str, str, str], ComponentEvidence] = {}

    def register_component(
        token: str,
        source_path: str,
        detail: str,
        confidence: float,
        evidence_type: str,
        version: str | None = None,
    ) -> None:
        normalized = normalize_product_name(token)
        component = components.get(normalized)
        resolved_version = version or (component.version if component else None)
        if component is None or confidence > component.confidence:
            components[normalized] = SoftwareComponent(
                name=token,
                normalized_name=normalized,
                version=resolved_version,
                source="firmware-analysis",
                confidence=confidence,
            )
        evidence[(normalized, source_path, detail)] = ComponentEvidence(
            component=components[normalized],
            source_path=source_path,
            evidence_type=evidence_type,
            detail=detail[:160],
            confidence=confidence,
        )

    for token in strings:
        lowered = token.lower()
        if lowered in _COMPONENT_TOKENS:
            register_component(token, "firmware.bin", f"Observed component token '{token}' in base payload.", 0.72, "string-match")

    if extracted_root is not None:
        for artifact in extracted_artifacts:
            path = extracted_root / Path(artifact.relative_path)
            content = path.read_text(encoding="utf-8", errors="ignore")
            lowered_content = content.lower()
            for token in sorted(_COMPONENT_TOKENS):
                if token in lowered_content:
                    register_component(
                        token,
                        artifact.relative_path,
                        f"Observed '{token}' in extracted file {artifact.relative_path}.",
                        0.82,
                        "content-token",
                    )
            for match in _VERSION_PATTERN.finditer(content):
                token = match.group("name")
                lowered_token = token.lower()
                if lowered_token in _COMPONENT_TOKENS:
                    register_component(
                        token,
                        artifact.relative_path,
                        f"Version-like component marker '{match.group(0)}' matched in extracted content.",
                        0.87,
                        "version-string",
                        version=match.group("version"),
                    )

    component_list = sorted(components.values(), key=lambda item: (item.normalized_name, item.version or "", item.name))
    evidence_list = sorted(evidence.values(), key=lambda item: (item.component.normalized_name, item.source_path, item.detail))
    return component_list, evidence_list[:32]


def _collect_vendor_product_hints(
    decoded: str,
    extracted_root: Path | None,
    config_hints: list[ConfigHint],
) -> tuple[list[str], list[str]]:
    vendors: set[str] = set()
    products: set[str] = set()
    contents = [decoded.lower()]
    names: list[str] = []
    if extracted_root is not None:
        for path in _iter_extracted_files(extracted_root):
            names.append(path.name.lower())
            contents.append(path.read_text(encoding="utf-8", errors="ignore").lower())
    contents.extend(f"{hint.key}={hint.value_preview}".lower() for hint in config_hints)
    merged = "\n".join([*contents, *names])
    for token, vendor in _VENDOR_HINTS.items():
        if token in merged:
            vendors.add(vendor)
    for token, product in _PRODUCT_HINTS.items():
        if token in merged:
            products.add(product)
    return sorted(vendors), sorted(products)


def _build_observed_facts(
    extraction_metadata: ExtractionMetadata,
    config_hints: list[ConfigHint],
    certificates: list[CertificateMetadata],
    extracted_artifacts: list[ExtractedArtifact],
    software_components: list[SoftwareComponent],
) -> list[ObservedFact]:
    facts: list[ObservedFact] = []
    if extraction_metadata.provider:
        facts.append(
            ObservedFact(
                field="extraction_provider",
                value=extraction_metadata.provider,
                source="firmware-extractor",
                confidence=0.96,
                attributes={
                    "member_count": extraction_metadata.member_count,
                    "detected_type": extraction_metadata.detected_type,
                },
            )
        )
    for artifact in extracted_artifacts[:8]:
        facts.append(
            ObservedFact(
                field="extracted_path",
                value=artifact.relative_path,
                source=artifact.source,
                confidence=artifact.confidence,
                attributes={
                    "media_type": artifact.media_type,
                    "depth": artifact.depth,
                    "size_bytes": artifact.size_bytes,
                },
            )
        )
    for hint in config_hints[:8]:
        facts.append(
            ObservedFact(
                field=f"config:{hint.key}",
                value=hint.value_preview,
                source=hint.path,
                confidence=hint.confidence,
                attributes={"category": hint.category},
            )
        )
    for certificate in certificates[:6]:
        facts.append(
            ObservedFact(
                field="certificate_cn",
                value=certificate.common_name or certificate.subject or "certificate",
                source="firmware-certificate",
                confidence=certificate.confidence,
            )
        )
    for component in software_components[:6]:
        facts.append(
            ObservedFact(
                field="software_component",
                value=component.normalized_name,
                source=component.source,
                confidence=component.confidence,
                attributes={"version": component.version},
            )
        )
    facts.sort(key=lambda item: (item.field, item.value, item.source))
    return facts


def _build_inference_hints(
    vendors: list[str],
    products: list[str],
    service_hints: list[ServiceHint],
    software_components: list[SoftwareComponent],
) -> list[InferenceHint]:
    hints: list[InferenceHint] = []
    for vendor in vendors:
        hints.append(
            InferenceHint(
                field="vendor",
                value=vendor,
                source="firmware-analysis",
                rationale="Vendor token matched extracted file names, config content, or payload strings.",
                confidence=0.74,
            )
        )
    for product in products:
        hints.append(
            InferenceHint(
                field="product",
                value=product,
                source="firmware-analysis",
                rationale="Product token matched extracted file names, config content, or payload strings.",
                confidence=0.68,
            )
        )
    for service_hint in service_hints[:10]:
        hints.append(
            InferenceHint(
                field="service",
                value=service_hint.service,
                source=service_hint.source_path,
                rationale=service_hint.rationale,
                confidence=service_hint.confidence,
                supporting_values=[service_hint.source_path],
                attributes={"source_kind": service_hint.source_kind},
            )
        )
    for component in software_components[:8]:
        hints.append(
            InferenceHint(
                field="component-family",
                value=component.normalized_name,
                source="firmware-analysis",
                rationale="Software component evidence suggests this component family is likely present in the firmware.",
                confidence=max(component.confidence - 0.08, 0.5),
                supporting_values=[component.name],
                attributes={"version": component.version or ""},
            )
        )
    hints.sort(key=lambda item: (item.field, item.value, item.source))
    return hints


def _build_findings(
    decoded: str,
    config_hints: list[ConfigHint],
    certificate_findings: list[CertificateFinding],
    service_hints: list[ServiceHint],
    software_components: list[SoftwareComponent],
) -> list[EvidenceFinding]:
    lowered = decoded.lower()
    findings: list[EvidenceFinding] = []
    if "password" in lowered or any(hint.key == "password" for hint in config_hints):
        findings.append(
            EvidenceFinding(
                title="Embedded credential reference",
                severity="high",
                confidence=0.86,
                description="The firmware analysis found credential-related strings or configuration values.",
                tags=["firmware", "credentials"],
            )
        )
    if "telnet" in lowered or any(hint.service == "telnet" for hint in service_hints):
        findings.append(
            EvidenceFinding(
                title="Legacy remote management feature",
                severity="medium",
                confidence=0.76,
                description="Firmware content suggests a legacy remote administration path such as Telnet.",
                tags=["firmware", "telnet"],
            )
        )
    if certificate_findings:
        findings.append(
            EvidenceFinding(
                title="Embedded certificate material",
                severity="medium",
                confidence=0.72,
                description="Extracted firmware content contains certificate material that should be reviewed for reuse and trust scope.",
                tags=["firmware", "certificate"],
            )
        )
    if any(hint.service in {"http", "https", "mqtt", "modbus", "rtsp"} for hint in service_hints):
        findings.append(
            EvidenceFinding(
                title="Recoverable service exposure hints",
                severity="medium",
                confidence=0.69,
                description="Extracted firmware content suggests likely remotely accessible services that should be validated against the live device posture.",
                tags=["firmware", "service-hint"],
            )
        )
    if any(component.normalized_name in {"openssl", "busybox", "dropbear"} for component in software_components):
        findings.append(
            EvidenceFinding(
                title="Recoverable software component evidence",
                severity="medium",
                confidence=0.67,
                description="Firmware analysis recovered software component evidence that can be correlated against patch baselines and local vulnerability feeds.",
                tags=["firmware", "component-evidence"],
            )
        )
    if not findings:
        findings.append(
            EvidenceFinding(
                title="No critical markers detected",
                severity="low",
                confidence=0.55,
                description="No high-confidence firmware indicators were matched by the current analysis pipeline.",
                tags=["firmware", "baseline"],
            )
        )
    findings.sort(key=lambda item: (-_severity_rank(item.severity), item.title))
    return findings


def _iter_extracted_files(extracted_root: Path) -> list[Path]:
    return sorted((path for path in extracted_root.rglob("*") if path.is_file()), key=lambda item: str(item.relative_to(extracted_root)))


def _extract_certificate_value(content: str, key: str) -> str:
    normalized_key = key.lower()
    for match in _CERT_VALUE_PATTERN.finditer(content):
        if match.group("key").lower() == normalized_key:
            return match.group(0).split("=", 1)[1].strip(" ,;/")
    return ""


def _severity_rank(severity: str) -> int:
    return {"critical": 4, "high": 3, "medium": 2, "low": 1}.get(severity.lower(), 0)

