"""Built-in actions for Tru-Guard."""

import os
import re
from typing import Any, Callable

import pyffx

from trusys.actions.base import BaseAction
from trusys.exceptions import ScanFailedError
from trusys.types import ActionType, ScanResult, Violation

# Default encryption key - should be overridden via TRUSYS_FPE_KEY environment variable
_DEFAULT_FPE_KEY = "tru$ys-d3fault-fp3-k3y-chang3-!n-pr0duct!0n"

# Character alphabets for FF1/FFX encryption
DIGITS = "0123456789"
LOWERCASE = "abcdefghijklmnopqrstuvwxyz"
UPPERCASE = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
ALPHANUMERIC_LOWER = DIGITS + LOWERCASE
ALPHANUMERIC_UPPER = DIGITS + UPPERCASE
ALPHANUMERIC = DIGITS + LOWERCASE + UPPERCASE
HEX_LOWER = DIGITS + "abcdef"
HEX_UPPER = DIGITS + "ABCDEF"

# Entity type to alphabet mapping for Presidio-detected entities
# Maps entity types to (alphabet, min_length) tuples
# min_length is the minimum chars needed for FFX (usually 2)
ENTITY_ALPHABETS: dict[str, tuple[str, int]] = {
    # Numeric entities - encrypt only digits
    "PHONE_NUMBER": (DIGITS, 2),
    "CREDIT_CARD": (DIGITS, 2),
    "US_SSN": (DIGITS, 2),
    "US_BANK_NUMBER": (DIGITS, 2),
    "US_ITIN": (DIGITS, 2),
    "UK_NHS": (DIGITS, 2),
    "DATE_TIME": (DIGITS, 2),
    "IP_ADDRESS": (DIGITS, 2),
    # Alphanumeric entities
    "EMAIL_ADDRESS": (ALPHANUMERIC_LOWER, 2),
    "URL": (ALPHANUMERIC_LOWER, 2),
    "IBAN_CODE": (ALPHANUMERIC_UPPER, 2),
    "US_PASSPORT": (ALPHANUMERIC_UPPER, 2),
    "US_DRIVER_LICENSE": (ALPHANUMERIC_UPPER, 2),
    "MEDICAL_LICENSE": (ALPHANUMERIC_UPPER, 2),
    "CRYPTO": (ALPHANUMERIC, 2),  # Crypto addresses use mixed case
    # Text entities - letters only (preserve spaces/punctuation)
    "PERSON": (LOWERCASE + UPPERCASE, 2),
    "LOCATION": (LOWERCASE + UPPERCASE, 2),
    "NRP": (LOWERCASE + UPPERCASE, 2),  # Nationality, Religious, Political
}

# Default alphabet for unknown entity types
DEFAULT_ALPHABET = ALPHANUMERIC
DEFAULT_MIN_LENGTH = 2


def _mask_text(matched_text: str) -> str:
    """Replace with asterisks. Uses fixed length for privacy."""
    return "****"


def _get_fpe_key() -> bytes:
    """Get the FPE encryption key as bytes."""
    key_str = os.environ.get("TRUSYS_FPE_KEY", _DEFAULT_FPE_KEY)
    # Ensure key is at least 16 bytes for AES
    key_bytes = key_str.encode("utf-8")
    if len(key_bytes) < 16:
        key_bytes = key_bytes.ljust(16, b"\0")
    return key_bytes


def _extract_and_encrypt_chars(
    text: str,
    alphabet: str,
    key: bytes,
    min_length: int = 2,
    decrypt: bool = False,
) -> str:
    """Extract characters matching alphabet, encrypt/decrypt them, and reconstruct.
    
    This preserves the format by keeping non-alphabet characters in place.
    For example, "212-555-5555" with digits alphabet becomes "847-329-4857".
    
    Args:
        text: The text to encrypt/decrypt
        alphabet: The character alphabet to use for FFX
        key: The encryption key as bytes
        min_length: Minimum characters needed for FFX
        decrypt: If True, decrypt instead of encrypt
        
    Returns:
        The encrypted/decrypted text with format preserved
    """
    # Extract chars that are in the alphabet and their positions
    chars_to_encrypt = []
    positions = []
    
    for i, char in enumerate(text):
        if char in alphabet:
            chars_to_encrypt.append(char)
            positions.append(i)
    
    if len(chars_to_encrypt) < min_length:
        # Not enough chars to encrypt, return original
        return text
    
    # Create FFX cipher and encrypt/decrypt the extracted chars
    chars_str = "".join(chars_to_encrypt)
    cipher = pyffx.String(key, alphabet=alphabet, length=len(chars_str))
    
    if decrypt:
        encrypted_chars = cipher.decrypt(chars_str)
    else:
        encrypted_chars = cipher.encrypt(chars_str)
    
    # Reconstruct the string with encrypted chars in original positions
    result = list(text)
    for pos, encrypted_char in zip(positions, encrypted_chars):
        result[pos] = encrypted_char
    
    return "".join(result)


def _encrypt_email(email: str, key: bytes, decrypt: bool = False) -> str:
    """Special handler for email addresses to preserve structure.
    
    Encrypts local part and domain separately, preserving @ and dots.
    """
    if "@" not in email:
        return _extract_and_encrypt_chars(email, ALPHANUMERIC_LOWER, key, decrypt=decrypt)
    
    local_part, domain = email.rsplit("@", 1)
    
    # Encrypt local part (before @)
    encrypted_local = _extract_and_encrypt_chars(
        local_part.lower(), ALPHANUMERIC_LOWER, key, min_length=1, decrypt=decrypt
    )
    
    # Encrypt domain parts separately (preserving dots)
    domain_parts = domain.lower().split(".")
    encrypted_domain_parts = []
    for part in domain_parts:
        if len(part) >= 2:
            encrypted_part = _extract_and_encrypt_chars(
                part, ALPHANUMERIC_LOWER, key, min_length=2, decrypt=decrypt
            )
        else:
            encrypted_part = part  # Keep very short parts as-is
        encrypted_domain_parts.append(encrypted_part)
    
    return f"{encrypted_local}@{'.'.join(encrypted_domain_parts)}"


def _encrypt_url(url: str, key: bytes, decrypt: bool = False) -> str:
    """Special handler for URLs to preserve structure.
    
    Preserves protocol (http/https), slashes, and basic structure.
    """
    # Match protocol prefix
    protocol_match = re.match(r"^(https?://)?(.*)$", url, re.IGNORECASE)
    if not protocol_match:
        return _extract_and_encrypt_chars(url, ALPHANUMERIC_LOWER, key, decrypt=decrypt)
    
    protocol = protocol_match.group(1) or ""
    rest = protocol_match.group(2)
    
    # Split on slashes and encrypt each part
    parts = rest.split("/")
    encrypted_parts = []
    for part in parts:
        if part:
            encrypted_part = _extract_and_encrypt_chars(
                part.lower(), ALPHANUMERIC_LOWER, key, min_length=2, decrypt=decrypt
            )
        else:
            encrypted_part = part
        encrypted_parts.append(encrypted_part)
    
    return protocol + "/".join(encrypted_parts)


def _encrypt_ip_address(ip: str, key: bytes, decrypt: bool = False) -> str:
    """Special handler for IP addresses to preserve octet structure.
    
    Encrypts each octet separately to maintain valid IP format.
    """
    parts = ip.split(".")
    if len(parts) != 4:
        # Not a standard IPv4, use generic encryption
        return _extract_and_encrypt_chars(ip, DIGITS, key, decrypt=decrypt)
    
    encrypted_parts = []
    for part in parts:
        if part.isdigit() and len(part) >= 1:
            # Encrypt the octet as a number (0-255 range preserved approximately)
            encrypted = _extract_and_encrypt_chars(
                part, DIGITS, key, min_length=1, decrypt=decrypt
            )
            # Ensure it's still a valid octet (0-255)
            try:
                val = int(encrypted)
                if val > 255:
                    encrypted = str(val % 256)
            except ValueError:
                pass
            encrypted_parts.append(encrypted)
        else:
            encrypted_parts.append(part)
    
    return ".".join(encrypted_parts)


def _encrypt_text(
    matched_text: str,
    key: bytes | None = None,
    entity_type: str | None = None,
    decrypt: bool = False,
) -> str:
    """Format-preserving encryption (FPE) using FF1/FFX algorithm.
    
    Encrypts text while preserving its format. Uses entity-type-aware
    alphabets to ensure the encrypted output has the same structure
    as the input (e.g., phone numbers stay as phone numbers).
    
    Args:
        matched_text: The text to encrypt
        key: Optional encryption key (uses TRUSYS_FPE_KEY env var if not provided)
        entity_type: Optional Presidio entity type for entity-aware encryption
        decrypt: If True, decrypt instead of encrypt
        
    Returns:
        The encrypted text with format preserved
    """
    if key is None:
        key = _get_fpe_key()
    
    # Handle special entity types with custom handlers
    if entity_type == "EMAIL_ADDRESS":
        return _encrypt_email(matched_text, key, decrypt=decrypt)
    elif entity_type == "URL":
        return _encrypt_url(matched_text, key, decrypt=decrypt)
    elif entity_type == "IP_ADDRESS":
        return _encrypt_ip_address(matched_text, key, decrypt=decrypt)
    
    # Get alphabet for entity type
    if entity_type and entity_type in ENTITY_ALPHABETS:
        alphabet, min_length = ENTITY_ALPHABETS[entity_type]
    else:
        alphabet, min_length = DEFAULT_ALPHABET, DEFAULT_MIN_LENGTH
    
    return _extract_and_encrypt_chars(
        matched_text, alphabet, key, min_length=min_length, decrypt=decrypt
    )


def decrypt_text(
    encrypted_text: str,
    key: bytes | None = None,
    entity_type: str | None = None,
) -> str:
    """Decrypt text that was encrypted with format-preserving encryption.
    
    This is the reverse of _encrypt_text and can be used to recover
    the original value when needed.
    
    Args:
        encrypted_text: The encrypted text to decrypt
        key: Optional decryption key (uses TRUSYS_FPE_KEY env var if not provided)
        entity_type: Optional Presidio entity type for entity-aware decryption
        
    Returns:
        The decrypted original text
    """
    return _encrypt_text(encrypted_text, key=key, entity_type=entity_type, decrypt=True)


def _get_matched_text(content: str, violation: Violation) -> str | None:
    """Resolve matched_text from violation metadata or span."""
    matched_text = violation.metadata.get("matched_text")
    if not matched_text and violation.span:
        start, end = violation.span
        matched_text = content[start:end]
    return matched_text or None


def _apply_replacements_by_spans(
    content: str,
    violations: list[Violation],
    get_replacement: Callable[[Violation], str],
) -> str:
    """Apply replacements using span positions so duplicate or overlapping text is handled correctly.

    Processes spans from right to left (by end index descending) so indices remain valid.
    Skips spans that are entirely contained in an already-replaced span (e.g. "example.com"
    inside "johnsmith1988@example.com") to avoid double replacement.
    """
    with_spans = [(v, v.span) for v in violations if v.span is not None]
    if not with_spans:
        return content
    # Sort by end descending, then start ascending (right-to-left order)
    with_spans.sort(key=lambda x: (x[1][1], -x[1][0]), reverse=True)
    replaced_ranges: list[tuple[int, int]] = []
    result = content
    for violation, (start, end) in with_spans:
        # Skip if this span is entirely inside an already replaced span
        if any(s <= start and end <= e for s, e in replaced_ranges):
            continue
        repl = get_replacement(violation)
        result = result[:start] + repl + result[end:]
        replaced_ranges.append((start, end))
    return result


def _apply_replacements(
    content: str,
    violations: list[Violation],
    get_replacement: Callable[[str], str],
) -> str:
    """Apply replacements for each violation. Uses span-based replacement when all have spans."""
    violations_with_spans = [v for v in violations if v.span is not None]
    if len(violations_with_spans) == len(violations) and violations_with_spans:

        def replacer(v: Violation) -> str:
            mt = _get_matched_text(content, v)
            return get_replacement(mt) if mt else ""

        return _apply_replacements_by_spans(content, violations_with_spans, replacer)
    # Fallback: replace by first occurrence of matched_text (legacy behavior)
    fixed_content = content
    for violation in violations:
        matched_text = _get_matched_text(content, violation)
        if not matched_text:
            continue
        replacement = get_replacement(matched_text)
        fixed_content = fixed_content.replace(matched_text, replacement, 1)
    return fixed_content


def _apply_entity_aware_replacements(
    content: str,
    violations: list[Violation],
    get_replacement: Callable[[str, str | None], str],
) -> str:
    """Apply entity-aware replacements for each violation.
    Uses span-based replacement when all have spans to avoid wrong occurrence with duplicates/overlaps.
    """
    violations_with_spans = [v for v in violations if v.span is not None]
    if len(violations_with_spans) == len(violations) and violations_with_spans:

        def replacer(v: Violation) -> str:
            mt = _get_matched_text(content, v)
            if not mt:
                return ""
            entity_type = v.metadata.get("entity_type")
            return get_replacement(mt, entity_type)

        return _apply_replacements_by_spans(content, violations_with_spans, replacer)
    # Fallback: replace by first occurrence of matched_text
    fixed_content = content
    for violation in violations:
        matched_text = _get_matched_text(content, violation)
        if not matched_text:
            continue
        entity_type = violation.metadata.get("entity_type")
        replacement = get_replacement(matched_text, entity_type)
        fixed_content = fixed_content.replace(matched_text, replacement, 1)
    return fixed_content


class FixAction(BaseAction):
    """Apply the scanner's suggested_fix when the scan fails.

    Only applies the suggested_fix from the scan result. If the scanner did not
    provide a suggested_fix, content is returned unchanged (no modification).
    """

    action_type = ActionType.FIX

    def execute(
        self,
        result: ScanResult,
        content: str,
        **kwargs: Any,
    ) -> tuple[str | None, bool]:
        """Execute the fix action: apply suggested_fix from the scanner if present."""
        if result.suggested_fix:
            return result.suggested_fix, True
        return None, True


class MaskAction(BaseAction):
    """Replace detected violations with asterisks (****)."""

    action_type = ActionType.MASK

    def execute(
        self,
        result: ScanResult,
        content: str,
        **kwargs: Any,
    ) -> tuple[str | None, bool]:
        """Execute the mask action: replace each violation span with ****."""
        violations_with_spans = [v for v in result.violations if v.span or v.metadata.get("matched_text")]
        if not violations_with_spans:
            if result.suggested_fix:
                return result.suggested_fix, True
            return None, True
        modified = _apply_replacements(content, violations_with_spans, _mask_text)
        return (modified, True) if modified != content else (None, True)


class EncryptAction(BaseAction):
    """Replace detected violations with format-preserving encryption (FPE).

    Uses FF1/FFX algorithm via pyffx for NIST-compliant format-preserving encryption.
    Supports entity-type-aware encryption for Presidio-detected PII entities.
    
    Set TRUSYS_FPE_KEY environment variable to specify the encryption key.
    The encrypted output can be decrypted using decrypt_text() with the same key.
    """

    action_type = ActionType.ENCRYPT

    def execute(
        self,
        result: ScanResult,
        content: str,
        **kwargs: Any,
    ) -> tuple[str | None, bool]:
        """Execute the encrypt action: replace each violation with FPE output.
        
        Uses entity-type-aware encryption when entity_type metadata is present
        (e.g., from PII scanner), selecting appropriate alphabets for each type.
        """
        violations_with_spans = [v for v in result.violations if v.span or v.metadata.get("matched_text")]
        if not violations_with_spans:
            if result.suggested_fix:
                return result.suggested_fix, True
            return None, True
        
        # Use entity-aware encryption that leverages Presidio entity types
        def encrypt_with_entity_type(matched_text: str, entity_type: str | None) -> str:
            return _encrypt_text(matched_text, entity_type=entity_type)
        
        modified = _apply_entity_aware_replacements(content, violations_with_spans, encrypt_with_entity_type)
        return (modified, True) if modified != content else (None, True)


class FilterAction(BaseAction):
    """Remove or replace problematic content.

    By default removes detected violations (replacement is empty string).
    Pass replacement to substitute with a placeholder (e.g. "[FILTERED]") instead.
    """

    action_type = ActionType.FILTER

    def __init__(self, replacement: str = "") -> None:
        """Initialize the filter action.

        Args:
            replacement: String to replace filtered content with; default "" removes it.
        """
        self.replacement = replacement

    def execute(
        self,
        result: ScanResult,
        content: str,
        **kwargs: Any,
    ) -> tuple[str | None, bool]:
        """Execute the filter action.

        Args:
            result: The scan result with violations
            content: The original content
            **kwargs: Additional arguments (may include 'replacement')

        Returns:
            Tuple of (filtered_content, should_continue=True)
        """
        replacement = kwargs.get("replacement", self.replacement)

        # Always apply our own replacement using violation spans/matched_text
        # instead of returning suggested_fix (which may use a different replacement
        # strategy, e.g. PII scanner's suggested_fix replaces with [ENTITY_TYPE] markers).
        violations_with_spans = [v for v in result.violations if v.span or v.metadata.get("matched_text")]
        if not violations_with_spans:
            if result.suggested_fix:
                return result.suggested_fix, True
            return None, True

        # Use span-based replacement when all violations have spans (avoids wrong occurrence with duplicates/overlaps)
        if all(v.span is not None for v in violations_with_spans):

            def replacer(v: Violation) -> str:
                return replacement

            filtered_content = _apply_replacements_by_spans(content, violations_with_spans, replacer)
        else:
            filtered_content = content
            for violation in violations_with_spans:
                matched_text = _get_matched_text(content, violation)
                if not matched_text:
                    continue
                filtered_content = filtered_content.replace(matched_text, replacement, 1)

        if filtered_content != content:
            return filtered_content, True

        # If no spans/matched_text for any violation, do not replace entire content;
        # return None so modified_content/final_content show unmodified content
        return None, True


class FailAction(BaseAction):
    """Raise an exception on violations.

    Stops processing and raises a ScanFailedError.
    """

    action_type = ActionType.FAIL

    def execute(
        self,
        result: ScanResult,
        content: str,
        **kwargs: Any,
    ) -> tuple[str | None, bool]:
        """Execute the fail action.

        Args:
            result: The scan result with violations
            content: The original content
            **kwargs: Additional arguments

        Returns:
            Never returns normally

        Raises:
            ScanFailedError: Always raised with violation details
        """
        raise ScanFailedError(
            scanner_name=result.scanner_name,
            violations=result.violations,
            message=f"Scan failed with {len(result.violations)} violation(s)",
        )


class IgnoreAction(BaseAction):
    """Log but allow content through unchanged.

    Does not modify the content; violations are recorded but ignored.
    """

    action_type = ActionType.IGNORE

    def execute(
        self,
        result: ScanResult,
        content: str,
        **kwargs: Any,
    ) -> tuple[str | None, bool]:
        """Execute the ignore action.

        Args:
            result: The scan result (logged but ignored)
            content: The original content
            **kwargs: Additional arguments

        Returns:
            Tuple of (None, should_continue=True)
        """
        # Content passes through unchanged
        return None, True


# Action type to action class mapping
_ACTION_MAP: dict[ActionType, type[BaseAction]] = {
    ActionType.FIX: FixAction,
    ActionType.MASK: MaskAction,
    ActionType.ENCRYPT: EncryptAction,
    ActionType.FILTER: FilterAction,
    ActionType.FAIL: FailAction,
    ActionType.IGNORE: IgnoreAction,
}


def get_action_for_type(action_type: ActionType) -> BaseAction:
    """Get an action instance for the given action type.

    Args:
        action_type: The action type

    Returns:
        An instance of the appropriate action class

    Raises:
        ValueError: If action type is not supported
    """
    if action_type == ActionType.CUSTOM:
        raise ValueError("Custom actions must be provided directly, not via get_action_for_type")

    if action_type not in _ACTION_MAP:
        raise ValueError(f"Unsupported action type: {action_type}")

    return _ACTION_MAP[action_type]()
