"""DataShield bridge — mask evidence samples before storing."""
from __future__ import annotations

import hashlib
import random
import re
import string
from typing import Any, Dict, List, Optional

from .contracts import EvidenceSample


# ---------------------------------------------------------------------------
# Try to import real DataShield — fall back to built-in stub
# ---------------------------------------------------------------------------

_DATASHIELD_AVAILABLE = False

try:
    from private_plugins.databridge_pro.datashield.masker import DataShieldMasker  # type: ignore
    _DATASHIELD_AVAILABLE = True
except ImportError:
    pass


class ShieldBridge:
    """Bridge between BLCE evidence samples and DataShield masking.

    If DataShield is available (Pro/Enterprise), delegates masking to its
    engine.  Otherwise uses a built-in stub that scrambles values while
    preserving data shape (column count, row count, types).
    """

    def __init__(self, use_datashield: bool = True):
        self._use_real = _DATASHIELD_AVAILABLE and use_datashield
        if self._use_real:
            self._masker = DataShieldMasker()  # type: ignore
        else:
            self._masker = None

    @property
    def backend(self) -> str:
        """Return which masking backend is active."""
        return "datashield" if self._use_real else "builtin_stub"

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

    def mask_sample(
        self,
        sample: EvidenceSample,
        rows: List[Dict[str, Any]],
    ) -> Dict[str, Any]:
        """Mask evidence-sample rows and return a masked copy.

        Returns a dict with ``masked_rows``, ``masked_sample`` (updated
        EvidenceSample with ``masked=True``), and ``backend`` indicator.
        """
        if not rows:
            return {
                "masked_rows": [],
                "masked_sample": sample.model_dump(),
                "backend": self.backend,
                "note": "No rows to mask.",
            }

        if self._use_real:
            masked_rows = self._mask_with_datashield(rows)
        else:
            masked_rows = self._mask_with_stub(rows, sample.columns)

        # Build a new EvidenceSample marked as masked
        masked_hash = self._hash_rows(masked_rows)
        masked_sample = EvidenceSample(
            sample_id=sample.sample_id,
            artifact_id=sample.artifact_id,
            table_name=sample.table_name,
            filter_desc=sample.filter_desc,
            row_count=len(masked_rows),
            columns=sample.columns,
            sample_hash=masked_hash,
            masked=True,
        )

        return {
            "masked_rows": masked_rows,
            "masked_sample": masked_sample.model_dump(),
            "backend": self.backend,
            "original_row_count": len(rows),
            "masked_row_count": len(masked_rows),
        }

    def mask_value(self, value: Any, column_name: str = "") -> Any:
        """Mask a single value based on its type."""
        if self._use_real:
            return self._masker.mask_value(value, column_name)  # type: ignore
        return self._stub_mask_value(value, column_name)

    def is_available(self) -> bool:
        """Return whether DataShield Pro is available."""
        return self._use_real

    # ------------------------------------------------------------------
    # DataShield Pro masking (delegates to real engine)
    # ------------------------------------------------------------------

    def _mask_with_datashield(self, rows: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Delegate to real DataShield masker."""
        return [
            {k: self._masker.mask_value(v, k) for k, v in row.items()}  # type: ignore
            for row in rows
        ]

    # ------------------------------------------------------------------
    # Built-in stub masking
    # ------------------------------------------------------------------

    def _mask_with_stub(
        self,
        rows: List[Dict[str, Any]],
        columns: List[str],
    ) -> List[Dict[str, Any]]:
        """Stub masking — scramble values while preserving shape and types."""
        masked: List[Dict[str, Any]] = []
        for row in rows:
            masked_row = {}
            for k, v in row.items():
                masked_row[k] = self._stub_mask_value(v, k)
            masked.append(masked_row)
        return masked

    @staticmethod
    def _stub_mask_value(value: Any, column_name: str = "") -> Any:
        """Mask a single value with type-preserving scrambling."""
        if value is None:
            return None

        if isinstance(value, bool):
            return random.choice([True, False])

        if isinstance(value, int):
            magnitude = max(1, abs(value))
            return random.randint(1, magnitude * 2)

        if isinstance(value, float):
            magnitude = max(1.0, abs(value))
            return round(random.uniform(0.1, magnitude * 2), 2)

        if isinstance(value, str):
            # Preserve date-like strings
            if re.match(r"\d{4}-\d{2}-\d{2}", value):
                return value  # dates are kept for temporal context

            # Mask string content
            length = max(len(value), 4)
            return "".join(random.choices(string.ascii_uppercase + string.digits, k=length))

        # Fallback — stringify and mask
        return "MASKED"

    @staticmethod
    def _hash_rows(rows: List[Dict[str, Any]]) -> str:
        """Hash masked rows for integrity tracking."""
        content = str(sorted(str(r) for r in rows))
        return hashlib.sha256(content.encode()).hexdigest()[:16]
