"""Deobfuscation layer for MCP Security Audit.

Pre-processes source code BEFORE AST parsing to detect and strip
obfuscated payloads: zero-width characters, Unicode tag characters
(ASCII smuggling), variation selectors, BiDi control characters,
HTML comments in strings, and Base64-encoded blocks.

Inspired by AgentSeal's deobfuscation approach.
Does NOT modify files on disk — operates only in memory.
"""

from __future__ import annotations

import base64
import re
from dataclasses import dataclass, field
from typing import Sequence


# ---------------------------------------------------------------------------
# Deobfuscation finding model
# ---------------------------------------------------------------------------

@dataclass
class DeobfuscationFinding:
    """A single deobfuscation finding — evidence of obfuscation detected."""

    line: int
    description: str
    original_bytes: bytes
    category: str  # zero_width | unicode_tag | variation_selector | bidi | html_comment | base64

    def to_dict(self) -> dict:
        return {
            "line": self.line,
            "description": self.description,
            "original_bytes": self.original_bytes.hex(),
            "category": self.category,
        }


# ---------------------------------------------------------------------------
# Character sets
# ---------------------------------------------------------------------------

# Zero-width characters (invisible Unicode)
ZERO_WIDTH_CHARS: frozenset[int] = frozenset({
    0x200B,  # ZERO WIDTH SPACE
    0x200C,  # ZERO WIDTH NON-JOINER
    0x200D,  # ZERO WIDTH JOINER
    0xFEFF,  # ZERO WIDTH NO-BREAK SPACE (BOM)
    0x00AD,  # SOFT HYPHEN
    0x2060,  # WORD JOINER
    0x180E,  # MONGOLIAN VOWEL SEPARATOR
})

# Unicode tag characters (U+E0001-U+E007F) — used for ASCII smuggling
TAG_CHAR_START: int = 0xE0001
TAG_CHAR_END: int = 0xE007F

# Variation selectors (U+FE00-U+FE0F)
VARIATION_SELECTOR_START: int = 0xFE00
VARIATION_SELECTOR_END: int = 0xFE0F

# BiDi control characters
BIDI_CHARS: frozenset[int] = frozenset({
    0x200E,  # LEFT-TO-RIGHT MARK
    0x200F,  # RIGHT-TO-LEFT MARK
    0x202A,  # LEFT-TO-RIGHT EMBEDDING
    0x202B,  # RIGHT-TO-LEFT EMBEDDING
    0x202C,  # POP DIRECTIONAL FORMATTING
    0x202D,  # LEFT-TO-RIGHT OVERRIDE
    0x202E,  # RIGHT-TO-LEFT OVERRIDE
    0x2066,  # LEFT-TO-RIGHT ISOLATE
    0x2067,  # RIGHT-TO-LEFT ISOLATE
    0x2068,  # FIRST STRONG ISOLATE
    0x2069,  # POP DIRECTIONAL ISOLATE
})

# HTML comment pattern inside strings
_HTML_COMMENT_RE = re.compile(r"<!--.*?-->", re.DOTALL)

# Base64 block detection — matches long Base64-looking blocks in string literals
# At least 40 chars of valid base64 with optional padding
_BASE64_BLOCK_RE = re.compile(
    r"""(?:['"])"""                      # opening quote
    r"""("""
    r"""[A-Za-z0-9+/]{40,}"""           # base64 content (min 40 chars)
    r"""={0,2}"""                        # optional padding
    r""")"""
    r"""(?:['"])""",                     # closing quote
)


# ---------------------------------------------------------------------------
# Core deobfuscation functions
# ---------------------------------------------------------------------------

def _is_suspicious_char(cp: int) -> str | None:
    """Classify a codepoint as a suspicious category, or None if clean.

    Args:
        cp: Unicode codepoint.

    Returns:
        Category string or None.
    """
    if cp in ZERO_WIDTH_CHARS:
        return "zero_width"
    if TAG_CHAR_START <= cp <= TAG_CHAR_END:
        return "unicode_tag"
    if VARIATION_SELECTOR_START <= cp <= VARIATION_SELECTOR_END:
        return "variation_selector"
    if cp in BIDI_CHARS:
        return "bidi"
    return None


def _decode_tag_chars(chars: Sequence[str]) -> str:
    """Decode Unicode tag characters to their ASCII equivalents.

    Tag characters U+E0001-U+E007F map to ASCII by subtracting 0xE0000.

    Args:
        chars: Sequence of tag characters.

    Returns:
        Decoded ASCII string (for reporting).
    """
    decoded: list[str] = []
    for ch in chars:
        cp = ord(ch)
        if TAG_CHAR_START <= cp <= TAG_CHAR_END:
            ascii_cp = cp - 0xE0000
            if 0x20 <= ascii_cp <= 0x7E:
                decoded.append(chr(ascii_cp))
    return "".join(decoded)


def _strip_invisible_chars(source: str) -> tuple[str, list[DeobfuscationFinding]]:
    """Strip zero-width, tag, variation selector, and BiDi characters.

    Args:
        source: Raw source code.

    Returns:
        Tuple of (cleaned source, list of findings).
    """
    findings: list[DeobfuscationFinding] = []
    lines = source.split("\n")
    cleaned_lines: list[str] = []

    for line_num, line in enumerate(lines, start=1):
        cleaned_chars: list[str] = []
        line_categories: dict[str, list[str]] = {}

        for ch in line:
            category = _is_suspicious_char(ord(ch))
            if category is not None:
                line_categories.setdefault(category, []).append(ch)
            else:
                cleaned_chars.append(ch)

        cleaned_lines.append("".join(cleaned_chars))

        for category, chars in line_categories.items():
            original_bytes = "".join(chars).encode("utf-8")
            count = len(chars)

            if category == "zero_width":
                desc = f"Stripped {count} zero-width character(s)"
            elif category == "unicode_tag":
                decoded = _decode_tag_chars(chars)
                desc = f"Stripped {count} Unicode tag character(s)"
                if decoded:
                    desc += f" (decoded: {decoded!r})"
            elif category == "variation_selector":
                desc = f"Stripped {count} variation selector(s)"
            elif category == "bidi":
                desc = f"Stripped {count} BiDi control character(s)"
            else:
                desc = f"Stripped {count} suspicious character(s)"

            findings.append(DeobfuscationFinding(
                line=line_num,
                description=desc,
                original_bytes=original_bytes,
                category=category,
            ))

    return "\n".join(cleaned_lines), findings


def _strip_html_comments_in_strings(source: str) -> tuple[str, list[DeobfuscationFinding]]:
    """Strip HTML comments found within string literals in source code.

    Args:
        source: Source code (may already be partially cleaned).

    Returns:
        Tuple of (cleaned source, list of findings).
    """
    findings: list[DeobfuscationFinding] = []
    lines = source.split("\n")
    cleaned_lines: list[str] = []

    for line_num, line in enumerate(lines, start=1):
        matches = list(_HTML_COMMENT_RE.finditer(line))
        if matches:
            cleaned_line = _HTML_COMMENT_RE.sub("", line)
            for match in matches:
                original = match.group(0)
                findings.append(DeobfuscationFinding(
                    line=line_num,
                    description=f"Stripped HTML comment in string: {original!r}",
                    original_bytes=original.encode("utf-8"),
                    category="html_comment",
                ))
            cleaned_lines.append(cleaned_line)
        else:
            cleaned_lines.append(line)

    return "\n".join(cleaned_lines), findings


def _detect_base64_blocks(source: str) -> list[DeobfuscationFinding]:
    """Detect Base64-encoded blocks in string literals.

    Does NOT strip them (they may be legitimate), but reports them
    as findings for manual review. Attempts decoding to check validity.

    Args:
        source: Source code.

    Returns:
        List of findings for detected Base64 blocks.
    """
    findings: list[DeobfuscationFinding] = []
    lines = source.split("\n")

    for line_num, line in enumerate(lines, start=1):
        for match in _BASE64_BLOCK_RE.finditer(line):
            b64_content = match.group(1)
            try:
                decoded = base64.b64decode(b64_content, validate=True)
                # Check if decoded content looks like text/code
                try:
                    decoded_str = decoded.decode("utf-8", errors="strict")
                    # Only flag if decoded content has printable chars
                    printable_ratio = sum(
                        1 for c in decoded_str if c.isprintable() or c in "\n\r\t"
                    ) / max(len(decoded_str), 1)
                    if printable_ratio > 0.7:
                        preview = decoded_str[:80].replace("\n", "\\n")
                        desc = (
                            f"Base64 block detected ({len(b64_content)} chars), "
                            f"decodes to: {preview!r}"
                        )
                        findings.append(DeobfuscationFinding(
                            line=line_num,
                            description=desc,
                            original_bytes=b64_content.encode("utf-8"),
                            category="base64",
                        ))
                except UnicodeDecodeError:
                    # Binary content — still suspicious
                    findings.append(DeobfuscationFinding(
                        line=line_num,
                        description=(
                            f"Base64 block detected ({len(b64_content)} chars), "
                            f"decodes to binary data"
                        ),
                        original_bytes=b64_content.encode("utf-8"),
                        category="base64",
                    ))
            except Exception:
                # Not valid base64, skip
                continue

    return findings


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

def deobfuscate(source: str) -> tuple[str, list[DeobfuscationFinding]]:
    """Pre-process source code to detect and strip obfuscated payloads.

    Runs the full deobfuscation pipeline:
    1. Strip zero-width characters
    2. Strip Unicode tag characters (ASCII smuggling)
    3. Strip variation selectors
    4. Strip BiDi control characters
    5. Strip HTML comments in strings
    6. Detect Base64 blocks (report only, no stripping)

    Args:
        source: Raw source code string.

    Returns:
        Tuple of (cleaned_source, list[DeobfuscationFinding]).
        The cleaned source can be safely passed to AST parsing.
        Findings describe what was detected and stripped.
    """
    all_findings: list[DeobfuscationFinding] = []

    # Step 1-4: Strip invisible characters (all handled in one pass)
    cleaned, invisible_findings = _strip_invisible_chars(source)
    all_findings.extend(invisible_findings)

    # Step 5: Strip HTML comments in strings
    cleaned, html_findings = _strip_html_comments_in_strings(cleaned)
    all_findings.extend(html_findings)

    # Step 6: Detect Base64 blocks (on cleaned source)
    base64_findings = _detect_base64_blocks(cleaned)
    all_findings.extend(base64_findings)

    return cleaned, all_findings
