from __future__ import annotations

import asyncio
import hashlib
import io
import json
import struct
from collections.abc import AsyncIterator
from dataclasses import dataclass
from datetime import UTC, datetime
from email.utils import parsedate_to_datetime
from pathlib import Path
from typing import Protocol

from opentelemetry import trace
from PIL import Image, UnidentifiedImageError
from pillow_heif import register_heif_opener

from app.contracts import Attachment

register_heif_opener()
tracer = trace.get_tracer("message-safety.dependencies")


class ObjectChanged(RuntimeError):
    pass


class DependencyFailure(RuntimeError):
    pass


class ObjectReader(Protocol):
    async def stream(self, attachment: Attachment) -> AsyncIterator[bytes]: ...


class Antivirus(Protocol):
    async def scan(
        self, chunks: AsyncIterator[bytes], *, scan_timeout: float
    ) -> AntivirusScanResult: ...

    async def status(self) -> AntivirusStatus: ...


@dataclass(frozen=True)
class AntivirusScanResult:
    threat: str | None
    signatures_version: str


@dataclass(frozen=True)
class AntivirusStatus:
    engine_version: str
    signatures_version: str
    databases_date: datetime


@dataclass(frozen=True)
class DetectorManifest:
    version: str
    supported: frozenset[str]
    max_size: int

    @classmethod
    def load(cls, path: Path) -> DetectorManifest:
        raw = path.read_bytes()
        data = json.loads(raw)
        version = "sha256:" + hashlib.sha256(raw).hexdigest()
        return cls(
            version, frozenset(data["supported_mime_types"]), data["hard_limits"]["max_size_bytes"]
        )


def validate_metadata(
    attachment: Attachment, manifest: DetectorManifest, enabled: set[str]
) -> str | None:
    if attachment.mime_type not in manifest.supported or attachment.mime_type not in enabled:
        return "file.unsupported_mime"
    if attachment.size_bytes > manifest.max_size:
        return "file.size_limit"
    return None


def detect_format(data: bytes, declared: str) -> str | None:
    matches: list[str] = []
    if data.startswith(b"\xff\xd8\xff") and data.endswith(b"\xff\xd9"):
        matches.append("image/jpeg")
    if data.startswith(b"\x89PNG\r\n\x1a\n") and b"IEND" in data[-64:]:
        matches.append("image/png")
    if len(data) >= 12 and data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        matches.append("image/webp")
    if len(data) >= 12 and data[4:8] == b"ftyp":
        brand = data[8:12]
        if brand in {b"heic", b"heix", b"hevc", b"hevx"}:
            matches.append("image/heic")
        if brand in {b"mif1", b"msf1"}:
            matches.append("image/heif")
    if data.startswith(b"%PDF-") and b"%%EOF" in data[-1024:]:
        matches.append("application/pdf")
    if len(matches) != 1:
        return "file.polyglot_or_ambiguous"
    if matches[0] != declared:
        return "file.format_mismatch"
    if declared == "application/pdf":
        lowered = data.lower()
        if b"/encrypt" in lowered:
            return "file.encrypted_content"
        if any(
            token in lowered
            for token in (b"/javascript", b"/openaction", b"/launch", b"/xfa", b"/embeddedfile")
        ):
            return "file.active_content"
    else:
        try:
            with Image.open(io.BytesIO(data)) as image:
                width, height = image.size
                if width > 10_000 or height > 10_000 or width * height > 25_000_000:
                    return "file.parser_limit"
                image.verify()
        except (UnidentifiedImageError, OSError, ValueError):
            return "file.format_mismatch"
    return None


async def collect_and_hash(
    reader: ObjectReader, attachment: Attachment, *, max_size: int
) -> tuple[bytes, bytes]:
    digest = hashlib.sha256()
    body = bytearray()
    async for chunk in reader.stream(attachment):
        if len(body) + len(chunk) > max_size:
            raise ObjectChanged("object exceeds bounded size")
        digest.update(chunk)
        body.extend(chunk)
    expected = bytes.fromhex(attachment.checksum.removeprefix("sha256:"))
    if len(body) != attachment.size_bytes or digest.digest() != expected:
        raise ObjectChanged("authoritative object metadata mismatch")
    return bytes(body), digest.digest()


class KeslSocketScanner:
    MAX_HEADER = 4096
    MAX_RESPONSE = 16 * 1024

    def __init__(self, socket_path: Path, timeout: float = 60.0) -> None:
        self.socket_path = socket_path
        self.timeout = timeout

    async def _request(
        self,
        header: dict[str, object],
        body: bytes = b"",
        *,
        request_timeout: float | None = None,
    ) -> dict[str, object]:
        encoded = json.dumps(header, separators=(",", ":")).encode()
        if len(encoded) > self.MAX_HEADER:
            raise DependencyFailure("KESL request header is too large")

        async def operation() -> dict[str, object]:
            reader, writer = await asyncio.open_unix_connection(str(self.socket_path))
            try:
                writer.write(struct.pack(">I", len(encoded)) + encoded + body)
                await writer.drain()
                (size,) = struct.unpack(">I", await reader.readexactly(4))
                if size < 2 or size > self.MAX_RESPONSE:
                    raise DependencyFailure("invalid KESL broker response size")
                value = json.loads((await reader.readexactly(size)).decode())
                if not isinstance(value, dict):
                    raise DependencyFailure("invalid KESL broker response")
                return value
            finally:
                writer.close()
                await writer.wait_closed()

        try:
            return await asyncio.wait_for(operation(), request_timeout or self.timeout)
        except (OSError, TimeoutError, asyncio.IncompleteReadError, json.JSONDecodeError) as exc:
            raise DependencyFailure("KESL unavailable") from exc

    async def scan(
        self, chunks: AsyncIterator[bytes], *, scan_timeout: float
    ) -> AntivirusScanResult:
        body = bytearray()
        async for chunk in chunks:
            body.extend(chunk)
        with tracer.start_as_current_span(
            "message_safety.antivirus.scan",
            attributes={"message_safety.antivirus.engine": "kesl"},
        ):
            value = await self._request(
                {"op": "scan", "size": len(body)},
                bytes(body),
                request_timeout=scan_timeout,
            )
        if value.get("status") != "completed" or value.get("verdict") not in {
            "clean",
            "infected",
        }:
            raise DependencyFailure("KESL scan failed")
        signatures = value.get("signatures_version")
        if not isinstance(signatures, str) or not signatures.startswith("sha256:"):
            raise DependencyFailure("KESL signatures version is missing")
        threat = value.get("threat")
        if value["verdict"] == "infected" and not isinstance(threat, str):
            raise DependencyFailure("KESL infected verdict has no threat")
        return AntivirusScanResult(
            threat=threat if isinstance(threat, str) else None,
            signatures_version=signatures,
        )

    async def status(self) -> AntivirusStatus:
        with tracer.start_as_current_span(
            "message_safety.antivirus.status",
            attributes={"message_safety.antivirus.engine": "kesl"},
        ):
            value = await self._request({"op": "status"})
        if value.get("status") != "ready":
            raise DependencyFailure("KESL is not ready")
        try:
            database_date = datetime.fromisoformat(
                str(value["databases_date"]).replace("Z", "+00:00")
            )
        except ValueError:
            try:
                database_date = parsedate_to_datetime(str(value["databases_date"]))
            except (TypeError, ValueError) as exc:
                raise DependencyFailure("invalid KESL database date") from exc
        if database_date.tzinfo is None:
            database_date = database_date.replace(tzinfo=UTC)
        engine = value.get("engine_version")
        signatures = value.get("signatures_version")
        if not isinstance(engine, str) or not isinstance(signatures, str):
            raise DependencyFailure("incomplete KESL status")
        return AntivirusStatus(engine, signatures, database_date.astimezone(UTC))


async def one_chunk(data: bytes) -> AsyncIterator[bytes]:
    yield data