"""
AgentDrop Shield -- Layer 1: Format Validation.

Validates file content against declared extensions using magic-byte
signatures, detects extension-content mismatches, and identifies
polyglot files that are valid as multiple formats simultaneously.

Uses only Python stdlib -- no libmagic dependency.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import Optional

from .models import ScanConfig, ThreatLevel

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


# ── Magic byte signatures ────────────────────────────────────────────
# (offset, magic_bytes, detected_mime)
# Ordered from most specific / longest match to shortest.

_MAGIC_TABLE: list[tuple[int, bytes, str]] = [
    # Archives
    (0, b"PK\x03\x04", "application/zip"),
    (0, b"PK\x05\x06", "application/zip"),       # empty archive
    (0, b"\x1f\x8b", "application/gzip"),
    (0, b"Rar!\x1a\x07", "application/x-rar"),
    (0, b"\xfd7zXZ\x00", "application/x-xz"),
    (0, b"7z\xbc\xaf\x27\x1c", "application/x-7z-compressed"),
    (0, b"BZh", "application/x-bzip2"),

    # Documents
    (0, b"%PDF-", "application/pdf"),
    (0, b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1", "application/msword"),  # OLE2 (doc/xls/ppt)

    # Images
    (0, b"\x89PNG\r\n\x1a\n", "image/png"),
    (0, b"\xff\xd8\xff", "image/jpeg"),
    (0, b"GIF87a", "image/gif"),
    (0, b"GIF89a", "image/gif"),
    (0, b"RIFF", "image/webp"),  # further check for WEBP at offset 8
    (0, b"BM", "image/bmp"),
    (0, b"II\x2a\x00", "image/tiff"),  # little-endian TIFF
    (0, b"MM\x00\x2a", "image/tiff"),  # big-endian TIFF

    # Audio / Video
    (0, b"ID3", "audio/mpeg"),
    (0, b"\xff\xfb", "audio/mpeg"),
    (0, b"\xff\xf3", "audio/mpeg"),
    (0, b"OggS", "audio/ogg"),
    (4, b"ftyp", "video/mp4"),

    # Executables
    (0, b"MZ", "application/x-executable-pe"),
    (0, b"\x7fELF", "application/x-executable-elf"),
    (0, b"\xfe\xed\xfa", "application/x-executable-macho"),
    (0, b"\xcf\xfa\xed\xfe", "application/x-executable-macho"),

    # Web / markup
    (0, b"<!DOCTYPE", "text/html"),
    (0, b"<html", "text/html"),
    (0, b"<?xml", "application/xml"),
    (0, b"<svg", "image/svg+xml"),
]

# ── Extension to expected MIME mapping ───────────────────────────────

EXTENSION_MIME_MAP: dict[str, set[str]] = {
    # Text
    ".txt": {"text/plain"},
    ".md": {"text/plain", "text/markdown"},
    ".csv": {"text/plain", "text/csv"},
    ".tsv": {"text/plain", "text/tab-separated-values"},
    ".log": {"text/plain"},

    # Data
    ".json": {"text/plain", "application/json"},
    ".yaml": {"text/plain", "application/x-yaml"},
    ".yml": {"text/plain", "application/x-yaml"},
    ".xml": {"text/plain", "application/xml"},
    ".toml": {"text/plain"},

    # Documents
    ".pdf": {"application/pdf"},
    ".doc": {"application/msword"},
    ".docx": {"application/zip"},   # OOXML is a ZIP
    ".xls": {"application/msword"},  # OLE2
    ".xlsx": {"application/zip"},
    ".pptx": {"application/zip"},

    # Images
    ".png": {"image/png"},
    ".jpg": {"image/jpeg"},
    ".jpeg": {"image/jpeg"},
    ".gif": {"image/gif"},
    ".webp": {"image/webp"},
    ".bmp": {"image/bmp"},
    ".tiff": {"image/tiff"},
    ".tif": {"image/tiff"},
    ".svg": {"image/svg+xml", "application/xml", "text/plain"},

    # Archives
    ".zip": {"application/zip"},
    ".gz": {"application/gzip"},
    ".tar": {"application/x-tar"},
    ".tar.gz": {"application/gzip"},
    ".tgz": {"application/gzip"},
    ".rar": {"application/x-rar"},
    ".7z": {"application/x-7z-compressed"},
    ".bz2": {"application/x-bzip2"},

    # Code (all should be plain text)
    ".py": {"text/plain"},
    ".js": {"text/plain"},
    ".ts": {"text/plain"},
    ".html": {"text/plain", "text/html"},
    ".htm": {"text/plain", "text/html"},
    ".css": {"text/plain"},
    ".sql": {"text/plain"},
    ".sh": {"text/plain"},
    ".bat": {"text/plain"},
    ".rs": {"text/plain"},
    ".go": {"text/plain"},
    ".java": {"text/plain"},
    ".c": {"text/plain"},
    ".cpp": {"text/plain"},
    ".h": {"text/plain"},
    ".rb": {"text/plain"},
}

# Dangerous MIME types that should always be flagged
_DANGEROUS_MIMES: set[str] = {
    "application/x-executable-pe",
    "application/x-executable-elf",
    "application/x-executable-macho",
}


@dataclass
class FormatValidationResult:
    """Outcome of the Layer 1 format validation."""

    detected_mime: Optional[str]
    declared_extension: str
    type_mismatch: bool = False
    is_polyglot: bool = False
    parseable_as: list[str] = field(default_factory=list)
    threat_level: ThreatLevel = ThreatLevel.CLEAN
    warnings: list[str] = field(default_factory=list)
    details: dict = field(default_factory=dict)


class FormatValidator:
    """Layer 1 -- magic-byte validation, mismatch detection, polyglot detection.

    All detection is based on the raw bytes; no external libraries are used.
    """

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

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

    def validate(self, data: bytes, filename: str) -> FormatValidationResult:
        """Validate file format integrity.

        Args:
            data: Raw file bytes.
            filename: Original filename (used for extension).

        Returns:
            FormatValidationResult with mismatch / polyglot findings.
        """
        ext = self._get_extension(filename)
        detected_mime = self._detect_mime(data)
        result = FormatValidationResult(
            detected_mime=detected_mime,
            declared_extension=ext,
        )
        result.details["filename"] = filename
        result.details["detected_mime"] = detected_mime

        # 1. Extension-content mismatch
        if ext and detected_mime:
            expected = EXTENSION_MIME_MAP.get(ext, set())
            if expected and detected_mime not in expected:
                result.type_mismatch = True
                result.threat_level = max(result.threat_level, ThreatLevel.MEDIUM)
                result.warnings.append(
                    f"Extension {ext!r} expects {expected} but content "
                    f"detected as {detected_mime!r}"
                )
                logger.warning(
                    "Format mismatch: %s declared as %s but detected %s",
                    filename, ext, detected_mime,
                )

        # 2. Dangerous content types
        if detected_mime in _DANGEROUS_MIMES:
            result.threat_level = max(result.threat_level, ThreatLevel.HIGH)
            result.warnings.append(
                f"Detected executable format: {detected_mime}"
            )

        # 3. Polyglot detection
        parseable = self._detect_polyglot(data)
        result.parseable_as = parseable
        result.details["parseable_as"] = parseable
        if len(parseable) > 1:
            result.is_polyglot = True
            result.threat_level = max(result.threat_level, ThreatLevel.HIGH)
            result.warnings.append(
                f"Polyglot file -- valid as: {', '.join(parseable)}"
            )
            logger.warning("Polyglot detected in %s: %s", filename, parseable)

        # 4. Null-extension or double-extension tricks
        if self._has_double_extension(filename):
            result.threat_level = max(result.threat_level, ThreatLevel.MEDIUM)
            result.warnings.append(
                f"Suspicious double extension in filename: {filename!r}"
            )

        return result

    # ------------------------------------------------------------------
    # Magic byte detection
    # ------------------------------------------------------------------

    @staticmethod
    def _detect_mime(data: bytes) -> Optional[str]:
        """Detect MIME type from magic bytes.  Returns None if unknown."""
        if not data:
            return None

        for offset, magic, mime in _MAGIC_TABLE:
            end = offset + len(magic)
            if len(data) >= end and data[offset:end] == magic:
                # Special case: RIFF container -- check for WEBP
                if magic == b"RIFF" and len(data) >= 12:
                    if data[8:12] != b"WEBP":
                        continue  # Not WEBP, skip
                return mime

        # Heuristic: if it looks like valid UTF-8 text, call it text/plain
        try:
            data[:8192].decode("utf-8")
            return "text/plain"
        except UnicodeDecodeError:
            pass

        return None

    # ------------------------------------------------------------------
    # Polyglot detection
    # ------------------------------------------------------------------

    @staticmethod
    def _detect_polyglot(data: bytes) -> list[str]:
        """Check if data is simultaneously valid as multiple formats.

        We do lightweight header probes only -- no full parsing.
        """
        formats: list[str] = []

        if not data:
            return formats

        # PDF
        if b"%PDF-" in data[:1024]:
            formats.append("pdf")

        # ZIP (can start after a preamble)
        if b"PK\x03\x04" in data[:65536]:
            formats.append("zip")

        # HTML
        lower_head = data[:4096].lower()
        if b"<html" in lower_head or b"<!doctype html" in lower_head:
            formats.append("html")

        # JPEG
        if data[:3] == b"\xff\xd8\xff":
            formats.append("jpeg")

        # PNG
        if data[:8] == b"\x89PNG\r\n\x1a\n":
            formats.append("png")

        # PE executable
        if data[:2] == b"MZ":
            formats.append("pe_executable")

        # ELF executable
        if data[:4] == b"\x7fELF":
            formats.append("elf_executable")

        # SVG (XML-based)
        if b"<svg" in lower_head:
            formats.append("svg")

        # JavaScript / script tags in what looks like non-HTML
        if b"<script" in lower_head and "html" not in formats:
            formats.append("script_tag")

        return formats

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

    @staticmethod
    def _get_extension(filename: str) -> str:
        """Extract lowercased extension including compound ones like .tar.gz."""
        lower = filename.lower()
        if lower.endswith(".tar.gz"):
            return ".tar.gz"
        if lower.endswith(".tar.bz2"):
            return ".tar.bz2"
        dot = lower.rfind(".")
        if dot == -1:
            return ""
        return lower[dot:]

    @staticmethod
    def _has_double_extension(filename: str) -> bool:
        """Detect suspicious double extensions like 'report.pdf.exe'."""
        parts = filename.rsplit(".", maxsplit=3)
        if len(parts) < 3:
            return False
        # Check if the penultimate "extension" is a known document type
        suspicious_inner = {
            "pdf", "doc", "docx", "xls", "xlsx", "txt", "csv", "jpg",
            "png", "gif", "zip", "html",
        }
        return parts[-2].lower() in suspicious_inner
