"""
AgentDrop Shield -- Layer 0: Resource Guards.

Prevents resource exhaustion attacks before any content analysis runs.
Detects zip bombs, oversized files, excessive nesting, and decompression
ratio anomalies.  All checks are O(1) or streaming -- never load a full
decompressed payload into memory.
"""

from __future__ import annotations

import io
import logging
import struct
import zipfile
from dataclasses import dataclass, field
from typing import Optional

from .models import ScanConfig, ThreatLevel

logger = logging.getLogger("agentdrop.shield.resource_guard")


@dataclass
class ResourceGuardResult:
    """Outcome of the Layer 0 resource checks."""

    passed: bool
    threat_level: ThreatLevel = ThreatLevel.CLEAN
    reason: Optional[str] = None
    warnings: list[str] = field(default_factory=list)
    details: dict = field(default_factory=dict)


class ResourceGuard:
    """Layer 0 -- file size limits, zip-bomb detection, depth limits.

    This layer runs first and is intentionally cheap.  If a file fails
    here the rest of the pipeline is skipped entirely.
    """

    def __init__(self, config: ScanConfig) -> None:
        self._cfg = config

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def check(self, data: bytes, filename: str) -> ResourceGuardResult:
        """Run all resource-guard checks on *data*.

        Args:
            data: Raw file bytes (before any decompression).
            filename: Original filename for extension-based decisions.

        Returns:
            ResourceGuardResult with pass/fail and diagnostics.
        """
        result = ResourceGuardResult(passed=True)

        # 1. Raw size check
        size = len(data)
        result.details["raw_size_bytes"] = size

        if size > self._cfg.max_file_size_bytes:
            return self._block(
                result,
                ThreatLevel.HIGH,
                f"File size {size:,} bytes exceeds limit "
                f"({self._cfg.max_file_size_bytes:,} bytes)",
            )

        if size == 0:
            result.warnings.append("Empty file (0 bytes)")
            return result

        # 2. Archive-specific checks
        if self._is_zip(data):
            archive_result = self._check_zip(data, result)
            if not archive_result.passed:
                return archive_result

        if self._is_gzip(data):
            gzip_result = self._check_gzip(data, result)
            if not gzip_result.passed:
                return gzip_result

        return result

    # ------------------------------------------------------------------
    # ZIP checks
    # ------------------------------------------------------------------

    @staticmethod
    def _is_zip(data: bytes) -> bool:
        """Check for ZIP magic bytes (PK\\x03\\x04)."""
        return data[:4] == b"PK\x03\x04"

    def _check_zip(self, data: bytes, result: ResourceGuardResult) -> ResourceGuardResult:
        """Inspect a ZIP archive without extracting."""
        try:
            with zipfile.ZipFile(io.BytesIO(data), "r") as zf:
                infos = zf.infolist()

                # Entry count
                entry_count = len(infos)
                result.details["archive_entries"] = entry_count
                if entry_count > self._cfg.max_archive_entries:
                    return self._block(
                        result,
                        ThreatLevel.HIGH,
                        f"ZIP contains {entry_count} entries "
                        f"(limit: {self._cfg.max_archive_entries})",
                    )

                # Aggregated decompression size + ratio
                total_compressed = 0
                total_uncompressed = 0
                nested_archives = 0

                for info in infos:
                    # Path traversal check
                    if info.filename.startswith("/") or ".." in info.filename:
                        return self._block(
                            result,
                            ThreatLevel.CRITICAL,
                            f"ZIP path traversal attempt: {info.filename!r}",
                        )

                    total_compressed += info.compress_size
                    total_uncompressed += info.file_size

                    # Check for nested archives by extension
                    lower = info.filename.lower()
                    if lower.endswith((".zip", ".gz", ".tar", ".tar.gz", ".tgz",
                                       ".bz2", ".xz", ".7z", ".rar")):
                        nested_archives += 1

                result.details["total_compressed"] = total_compressed
                result.details["total_uncompressed"] = total_uncompressed
                result.details["nested_archives"] = nested_archives

                # Decompressed size limit
                if total_uncompressed > self._cfg.max_decompressed_bytes:
                    return self._block(
                        result,
                        ThreatLevel.CRITICAL,
                        f"ZIP decompressed size {total_uncompressed:,} bytes "
                        f"exceeds limit ({self._cfg.max_decompressed_bytes:,} bytes) "
                        "-- possible zip bomb",
                    )

                # Decompression ratio
                if total_compressed > 0:
                    ratio = total_uncompressed / total_compressed
                    result.details["decompression_ratio"] = round(ratio, 2)
                    if ratio > self._cfg.max_decompression_ratio:
                        return self._block(
                            result,
                            ThreatLevel.CRITICAL,
                            f"ZIP decompression ratio {ratio:.1f}:1 exceeds limit "
                            f"({self._cfg.max_decompression_ratio:.0f}:1) "
                            "-- possible zip bomb",
                        )

                # Nesting depth warning (we cannot know true depth without
                # recursive extraction, but presence of archive entries is
                # a useful signal).
                if nested_archives > 0:
                    result.warnings.append(
                        f"ZIP contains {nested_archives} nested archive(s); "
                        f"depth limit is {self._cfg.max_archive_depth}"
                    )

        except zipfile.BadZipFile:
            result.warnings.append("File has ZIP magic bytes but is not a valid ZIP")
        except Exception as exc:
            logger.warning("ZIP inspection failed: %s", exc)
            result.warnings.append(f"ZIP inspection error: {exc}")

        return result

    # ------------------------------------------------------------------
    # GZIP checks
    # ------------------------------------------------------------------

    @staticmethod
    def _is_gzip(data: bytes) -> bool:
        """Check for GZIP magic bytes (\\x1f\\x8b)."""
        return data[:2] == b"\x1f\x8b"

    def _check_gzip(self, data: bytes, result: ResourceGuardResult) -> ResourceGuardResult:
        """Check gzip for decompression-bomb indicators.

        The last 4 bytes of a gzip stream contain the original uncompressed
        size modulo 2**32 (ISIZE field, little-endian).  This is not
        perfectly reliable for files > 4 GB but catches the classic bombs.
        """
        if len(data) < 18:
            result.warnings.append("GZIP file too small to contain valid data")
            return result

        try:
            isize = struct.unpack("<I", data[-4:])[0]
            result.details["gzip_declared_size"] = isize
            compressed_size = len(data)

            if isize > self._cfg.max_decompressed_bytes:
                return self._block(
                    result,
                    ThreatLevel.CRITICAL,
                    f"GZIP declared size {isize:,} bytes exceeds limit "
                    f"({self._cfg.max_decompressed_bytes:,} bytes) "
                    "-- possible gzip bomb",
                )

            if compressed_size > 0:
                ratio = isize / compressed_size
                result.details["gzip_ratio"] = round(ratio, 2)
                if ratio > self._cfg.max_decompression_ratio:
                    return self._block(
                        result,
                        ThreatLevel.CRITICAL,
                        f"GZIP ratio {ratio:.1f}:1 exceeds limit "
                        f"({self._cfg.max_decompression_ratio:.0f}:1) "
                        "-- possible gzip bomb",
                    )
        except struct.error:
            result.warnings.append("Could not read GZIP ISIZE field")

        return result

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _block(
        result: ResourceGuardResult,
        level: ThreatLevel,
        reason: str,
    ) -> ResourceGuardResult:
        """Mark a result as blocked."""
        result.passed = False
        result.threat_level = level
        result.reason = reason
        logger.warning("Resource guard BLOCKED: %s", reason)
        return result
