"""
Privacy and content sanitization for Amplitude AI SDK.

Handles content redaction, hashing, and PII protection when privacy_mode is enabled.
"""

import hashlib
import re
from typing import Any
from urllib.parse import urlparse

# Constants for redacted content
REDACTED_IMAGE_PLACEHOLDER = "[base64 image redacted]"
REDACTED_CONTENT_PLACEHOLDER = "[content redacted]"

# Legacy chunking constants — kept only so get_text_from_llm_message() can
# still read old chunked events.  New events always use {"text": content}.
MAX_CHUNK_SIZE = 1024
MAX_CHUNKS = 8


def _normalize_tool_definitions(
    tool_definitions: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    """Normalize tool definitions from various provider formats into a canonical shape.

    Produces ``[{"name": ..., "description": ..., "parameters": ...}]`` regardless
    of whether the input is OpenAI, Anthropic, Bedrock, or other formats.
    """
    normalized: list[dict[str, Any]] = []
    for tool in tool_definitions:
        if not isinstance(tool, dict):
            continue
        # OpenAI Chat format: {"type": "function", "function": {"name", "description", "parameters"}}
        fn = tool.get("function")
        if isinstance(fn, dict):
            normalized.append({
                "name": fn.get("name", ""),
                "description": fn.get("description", ""),
                "parameters": fn.get("parameters"),
            })
            continue
        # Anthropic format: {"name", "description", "input_schema"}
        if "input_schema" in tool:
            normalized.append({
                "name": tool.get("name", ""),
                "description": tool.get("description", ""),
                "parameters": tool.get("input_schema"),
            })
            continue
        # Bedrock format: {"toolSpec": {"name", "description", "inputSchema"}}
        tool_spec = tool.get("toolSpec")
        if isinstance(tool_spec, dict):
            normalized.append({
                "name": tool_spec.get("name", ""),
                "description": tool_spec.get("description", ""),
                "parameters": tool_spec.get("inputSchema"),
            })
            continue
        # Gemini format: {"function_declarations": [{"name", "description", "parameters"}]}
        fn_decls = tool.get("function_declarations")
        if isinstance(fn_decls, list):
            for decl in fn_decls:
                if isinstance(decl, dict):
                    normalized.append({
                        "name": decl.get("name", ""),
                        "description": decl.get("description", ""),
                        "parameters": decl.get("parameters"),
                    })
            continue
        # OpenAI Responses / generic format: {"name", "description", "parameters"}
        if "name" in tool:
            normalized.append({
                "name": tool.get("name", ""),
                "description": tool.get("description", ""),
                "parameters": tool.get("parameters"),
            })
    return normalized


def is_base64_data_url(text: str) -> bool:
    """Check if text is a base64 data URL."""
    return bool(re.match(r"^data:([^;]+);base64,", text))


def is_valid_url(text: str) -> bool:
    """Check if text is a valid URL or relative path."""
    try:
        result = urlparse(text)
        return bool(result.scheme and result.netloc)
    except Exception:
        pass

    return text.startswith(("/", "./", "../"))


def is_raw_base64(text: str) -> bool:
    """Check if text appears to be raw base64 encoded data."""
    if is_valid_url(text):
        return False

    # Must be reasonably long and match base64 pattern
    return len(text) > 20 and bool(re.match(r"^[A-Za-z0-9+/]+=*$", text))


def create_content_hash(content: Any) -> str:
    """Create a SHA-256 hash of content for privacy mode."""
    if content is None:
        return ""

    # Convert to string representation
    content_str = str(content) if not isinstance(content, str) else content

    # Create hash
    return hashlib.sha256(content_str.encode("utf-8")).hexdigest()


def redact_base64_content(value: Any) -> Any:
    """Redact base64 encoded images and data URLs."""
    if not isinstance(value, str):
        return value

    if is_base64_data_url(value):
        return REDACTED_IMAGE_PLACEHOLDER

    if is_raw_base64(value):
        return REDACTED_IMAGE_PLACEHOLDER

    return value


def redact_pii_patterns(text: str) -> str:
    """Redact common PII patterns from text."""
    # Email addresses
    text = re.sub(r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b", "[email]", text)

    # Phone numbers (US format)
    text = re.sub(r"\b\(?([0-9]{3})\)?[-. ]?([0-9]{3})[-. ]?([0-9]{4})\b", "[phone]", text)

    # Credit card numbers (basic pattern)
    text = re.sub(r"\b(?:\d{4}[-\s]?){3}\d{4}\b", "[credit_card]", text)

    # SSN pattern
    text = re.sub(r"\b\d{3}-\d{2}-\d{4}\b", "[ssn]", text)

    return text


def extract_text_from_openai_message(message: dict[str, Any]) -> str:
    """Extract text content from OpenAI message format."""
    if not isinstance(message, dict) or "content" not in message:
        return ""

    content = message["content"]

    if isinstance(content, str):
        return content
    elif isinstance(content, list):
        # Multi-modal content (text + images)
        text_parts = []
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                text_parts.append(item.get("text", ""))
        return "".join(text_parts)

    return str(content)


def extract_text_from_anthropic_message(message: dict[str, Any]) -> str:
    """Extract text content from Anthropic message format."""
    if not isinstance(message, dict) or "content" not in message:
        return ""

    content = message["content"]

    if isinstance(content, str):
        return content
    elif isinstance(content, list):
        # Anthropic's content blocks format
        text_parts = []
        for block in content:
            if isinstance(block, dict) and block.get("type") == "text":
                text_parts.append(block.get("text", ""))
        return "".join(text_parts)

    return str(content)


def extract_text_from_gemini_contents(contents: Any) -> str:
    """Extract text content from Gemini contents format."""
    if isinstance(contents, str):
        return contents
    elif isinstance(contents, list):
        text_parts = []
        for content in contents:
            if isinstance(content, dict):
                parts = content.get("parts", [])
                if isinstance(parts, list) and parts:
                    text = parts[0] if isinstance(parts[0], str) else str(parts[0])
                    text_parts.append(text)
                elif isinstance(parts, str):
                    text_parts.append(parts)
            elif isinstance(content, str):
                text_parts.append(content)
        return "".join(text_parts)

    return str(contents)


def _chunk_content(text: str, **_kwargs: Any) -> dict[str, Any]:
    """Return ``$llm_message`` payload for the given text.

    Content is stored as ``{"text": content}`` at full length — the
    :class:`AgentEvent` subclass prevents the base SDK from truncating it,
    and Nova already whitelists ``$llm_message`` server-side.

    Previous versions split long content into ``c0..c7`` chunks.  That
    format is still readable via :func:`get_text_from_llm_message` for
    backward compatibility, but is no longer produced.
    """
    return {"text": text}


def get_text_from_llm_message(llm_message: dict[str, Any]) -> str:
    """Extract the full text from a $llm_message dict, whether chunked or not.

    Callers that need the plain-text content (e.g. legacy provider wrappers)
    should use this instead of accessing ['text'] directly, since chunked
    content uses c0..cN keys rather than 'text'.
    """
    if "text" in llm_message:
        return str(llm_message["text"])

    n = llm_message.get("n", 0)
    if isinstance(n, int) and n > 0:
        return "".join(llm_message.get(f"c{i}", "") for i in range(n))

    return ""


def sanitize_any_content(content: Any, privacy_mode: bool = False, redact_pii: bool = True) -> dict[str, Any]:
    """
    Universal content sanitizer that works with any format.

    Args:
        content: The content to sanitize (string, dict, list, etc.)
        privacy_mode: If True, replace content with hash
        redact_pii: If True, redact PII patterns

    Returns:
        Dictionary with sanitized content and/or hash
    """
    result: dict[str, Any] = {}

    if content is None:
        return result

    # Convert any content type to text
    if isinstance(content, str):
        text_content = content
    elif isinstance(content, (dict, list)):
        text_content = _extract_text_from_structured_content(content)
    else:
        text_content = str(content)

    # Tool-call-only LLM responses have content=None (caught above) or
    # empty string.  Emitting $llm_message with blank text causes the
    # response to appear as "missing text" in the Amplitude thread view.
    if len(text_content) == 0:
        return result

    if redact_pii:
        text_content = redact_pii_patterns(text_content)
        text_content = redact_base64_content(text_content)

    if privacy_mode:
        result["content_hash"] = create_content_hash(text_content)
    else:
        result["$llm_message"] = _chunk_content(text_content)

    return result


def _extract_text_from_structured_content(content: Any) -> str:
    """Extract text from any structured content format.

    Handles dict, list, str, None, and other scalar types that may appear
    when recursing into message fields (e.g. OpenAI tool-call responses
    with ``{"content": None}``).
    """
    if content is None:
        return ""

    if isinstance(content, str):
        return content

    if isinstance(content, dict):
        for field in ["content", "text", "message"]:
            if field in content:
                return _extract_text_from_structured_content(content[field])

        return str(content)

    if isinstance(content, list):
        text_parts: list[str] = []
        for item in content:
            if isinstance(item, str):
                text_parts.append(item)
            elif isinstance(item, dict):
                text_parts.append(_extract_text_from_structured_content(item))
            else:
                text_parts.append(str(item))
        return "".join(text_parts)

    return str(content)


# Keep this for backward compatibility
def sanitize_message_content(content: Any, privacy_mode: bool = False, redact_pii: bool = True) -> dict[str, Any]:
    """Backward compatibility wrapper for sanitize_any_content."""
    return sanitize_any_content(content, privacy_mode, redact_pii)


def _sanitize_structured_content(
    content: Any, redact_pii: bool
) -> Any:
    """Recursively sanitize structured content (strings, dicts, and lists).

    - **str**: PII-redacted and base64-scrubbed.
    - **dict**: each value sanitized recursively.
    - **list**: each element sanitized recursively.
    - Other types pass through unchanged.
    """
    if isinstance(content, str):
        if redact_pii:
            content = redact_pii_patterns(content)
        return redact_base64_content(content)

    if isinstance(content, dict):
        sanitized_dict: dict[str, Any] = {}
        for key, value in content.items():
            sanitized_dict[key] = _sanitize_structured_content(value, redact_pii)
        return sanitized_dict

    if isinstance(content, list):
        return [_sanitize_structured_content(item, redact_pii) for item in content]

    return content


# Legacy functions kept for backward compatibility
def sanitize_openai_message(
    message: dict[str, Any], privacy_mode: bool = False, redact_pii: bool = True
) -> dict[str, Any]:
    """Legacy function - use sanitize_any_content instead."""
    text_content = extract_text_from_openai_message(message)
    sanitized_result = sanitize_any_content(text_content, privacy_mode, redact_pii)

    result = message.copy() if isinstance(message, dict) else {}
    if privacy_mode:
        result["content_hash"] = sanitized_result.get("content_hash")
        result.pop("content", None)
    elif "$llm_message" in sanitized_result:
        result["content"] = get_text_from_llm_message(sanitized_result["$llm_message"])

    return result


def sanitize_anthropic_message(
    message: dict[str, Any], privacy_mode: bool = False, redact_pii: bool = True
) -> dict[str, Any]:
    """Legacy function - use sanitize_any_content instead."""
    text_content = extract_text_from_anthropic_message(message)
    sanitized_result = sanitize_any_content(text_content, privacy_mode, redact_pii)

    result = message.copy() if isinstance(message, dict) else {}
    if privacy_mode:
        result["content_hash"] = sanitized_result.get("content_hash")
        result.pop("content", None)
    elif "$llm_message" in sanitized_result:
        result["content"] = get_text_from_llm_message(sanitized_result["$llm_message"])

    return result


def sanitize_tool_input_output(
    tool_input: Any = None, tool_output: Any = None, privacy_mode: bool = False
) -> dict[str, Any]:
    """
    Sanitize tool input/output for tracking.

    Returns dictionary with either raw values or hashes based on privacy_mode.
    """
    result: dict[str, Any] = {}

    if tool_input is not None:
        if privacy_mode:
            result["tool_input_hash"] = create_content_hash(tool_input)
        # Sanitize but keep raw input
        elif isinstance(tool_input, str):
            result["tool_input"] = redact_pii_patterns(tool_input)
        else:
            result["tool_input"] = tool_input

    if tool_output is not None:
        if privacy_mode:
            result["tool_output_hash"] = create_content_hash(tool_output)
            result["tool_output_size"] = len(str(tool_output))
        # Sanitize but keep raw output
        elif isinstance(tool_output, str):
            result["tool_output"] = redact_pii_patterns(tool_output)
        else:
            result["tool_output"] = tool_output

    return result


class PrivacyConfig:
    """Configuration for privacy and sanitization settings.

    Parameters
    ----------
    privacy_mode:
        Legacy boolean flag.  When *True*, content is stripped entirely
        (equivalent to ``content_mode="metadata_only"``).  Prefer using
        ``content_mode`` instead.
    redact_pii:
        Scrub PII patterns from content when ``content_mode="full"``.
    custom_redaction_patterns:
        Additional regex patterns to redact.
    content_mode:
        The :class:`~amplitude_ai.config.ContentMode` enum value (or
        its string equivalent).  When provided, takes precedence over
        ``privacy_mode``.  If *None*, behavior falls back to the boolean
        ``privacy_mode`` flag for backward compatibility.
    """

    _VALID_CONTENT_MODES = {"full", "metadata_only", "customer_enriched"}

    def __init__(
        self,
        privacy_mode: bool = False,
        redact_pii: bool = False,
        custom_redaction_patterns: list[str] | None = None,
        content_mode: str | None = None,
        validate: bool = False,
        debug: bool = False,
    ) -> None:
        self.privacy_mode = privacy_mode
        self.redact_pii = redact_pii
        self.validate = validate
        self.debug = debug
        self.custom_patterns = custom_redaction_patterns or []
        # Store as a plain string so we can compare with StrEnum values
        # without importing ContentMode (avoids circular import).
        # None means "use legacy privacy_mode boolean".
        mode_str = getattr(content_mode, "value", content_mode) if content_mode is not None else None
        if mode_str is not None:
            mode_str = str(mode_str)
        if mode_str is not None and mode_str not in self._VALID_CONTENT_MODES:
            raise ValueError(
                f"Invalid content_mode {content_mode!r}. "
                f"Must be one of: {', '.join(sorted(self._VALID_CONTENT_MODES))}"
            )
        self._content_mode: str | None = mode_str

    @property
    def content_mode(self) -> str | None:
        """The active content mode, or *None* if using legacy ``privacy_mode``."""
        return self._content_mode

    def _apply_custom_patterns(self, text: str) -> str:
        """Apply customer-provided redaction patterns to text."""
        if not self.custom_patterns or not isinstance(text, str):
            return text
        for pattern in self.custom_patterns:
            try:
                text = re.sub(pattern, "[REDACTED]", text)
            except re.error:
                pass  # Skip invalid patterns silently
        return text

    def _apply_custom_patterns_to_llm_message(self, llm_message: dict[str, Any]) -> None:
        """Apply custom redaction patterns to an ``$llm_message`` dict in-place.

        New events always use the ``"text"`` key.  Legacy chunked events
        (``"c0"``..``"cN"``) are also handled for backward compatibility.
        """
        if "text" in llm_message:
            llm_message["text"] = self._apply_custom_patterns(llm_message["text"])
        else:
            n = llm_message.get("n", 0)
            if isinstance(n, int):
                for i in range(n):
                    key = f"c{i}"
                    if key in llm_message:
                        llm_message[key] = self._apply_custom_patterns(llm_message[key])

    def sanitize_content(self, content: Any) -> dict[str, Any]:
        """Sanitize content based on the active privacy tier.

        Returns a dict of event properties to merge into the event.

        - ``"full"``: PII-scrubbed text under ``$llm_message.text``
        - ``"metadata_only"`` / ``"customer_enriched"``: empty dict (no content)
        - Legacy (no content_mode): delegates to ``sanitize_any_content``
        """
        if self._content_mode is None:
            # Legacy path — privacy_mode=True maps to metadata_only behavior
            if self.privacy_mode:
                return {}
            result = sanitize_any_content(
                content,
                privacy_mode=False,
                redact_pii=self.redact_pii,
            )
            if self.custom_patterns and "$llm_message" in result:
                self._apply_custom_patterns_to_llm_message(result["$llm_message"])
            return result

        if self._content_mode == "full":
            result = sanitize_any_content(
                content, privacy_mode=False, redact_pii=self.redact_pii
            )
            if self.custom_patterns and "$llm_message" in result:
                self._apply_custom_patterns_to_llm_message(result["$llm_message"])
            return result
        else:
            # metadata_only, customer_enriched → no content
            return {}

    def sanitize_system_prompt(self, system_prompt: str | None) -> dict[str, Any]:
        """Sanitize a system prompt for AI response events.

        Returns event properties for the system prompt, respecting privacy tiers:

        - ``PROP_SYSTEM_PROMPT_LENGTH`` is **always** included (all tiers)
          so users can track prompt size trends without content.
        - ``"full"``: includes PII-redacted + custom-pattern-redacted text,
          truncated at 10 000 characters.
        - ``"metadata_only"`` / ``"customer_enriched"``: length only, no content.
        """
        from .tracking import PROP_SYSTEM_PROMPT, PROP_SYSTEM_PROMPT_LENGTH

        if not system_prompt:
            return {}

        result: dict[str, Any] = {
            PROP_SYSTEM_PROMPT_LENGTH: len(system_prompt),
        }

        # Determine effective mode
        mode = self._content_mode
        if mode is None:
            mode = "metadata_only" if self.privacy_mode else "full"

        if mode == "full":
            sanitized = system_prompt
            if self.redact_pii:
                sanitized = redact_pii_patterns(sanitized)
            sanitized = self._apply_custom_patterns(sanitized)
            truncated = sanitized[:10000] if len(sanitized) > 10000 else sanitized
            result[PROP_SYSTEM_PROMPT] = truncated

        return result

    def sanitize_tool_definitions(
        self, tool_definitions: list[dict[str, Any]] | None
    ) -> dict[str, Any]:
        """Sanitize request-side tool definitions for AI response events.

        Returns event properties for tool definitions, respecting privacy tiers:

        - ``PROP_TOOL_DEFINITIONS_COUNT`` is **always** included (all tiers)
        - ``PROP_TOOL_DEFINITIONS_HASH`` is **always** included (all tiers)
          — a SHA-256 hash of the canonical tool definitions for grouping/dedup
        - ``"full"``: includes PII-redacted + custom-pattern-redacted JSON,
          truncated at 10 000 characters.
        - ``"metadata_only"`` / ``"customer_enriched"``: count + hash only.
        """
        from .tracking import (
            PROP_TOOL_DEFINITIONS,
            PROP_TOOL_DEFINITIONS_COUNT,
            PROP_TOOL_DEFINITIONS_HASH,
        )
        import json

        if not tool_definitions:
            return {}

        result: dict[str, Any] = {
            PROP_TOOL_DEFINITIONS_COUNT: len(tool_definitions),
        }

        normalized = _normalize_tool_definitions(tool_definitions)
        canonical_json = json.dumps(normalized, sort_keys=True, separators=(",", ":"))
        result[PROP_TOOL_DEFINITIONS_HASH] = hashlib.sha256(
            canonical_json.encode()
        ).hexdigest()[:16]

        mode = self._content_mode
        if mode is None:
            mode = "metadata_only" if self.privacy_mode else "full"

        if mode == "full":
            serialized = json.dumps(normalized, separators=(",", ":"))
            if self.redact_pii:
                serialized = redact_pii_patterns(serialized)
            serialized = self._apply_custom_patterns(serialized)
            truncated = serialized[:10000] if len(serialized) > 10000 else serialized
            result[PROP_TOOL_DEFINITIONS] = truncated

        return result

    def sanitize_reasoning_content(
        self, reasoning_content: str | None, reasoning_tokens: int | None = None
    ) -> dict[str, Any]:
        """Sanitize reasoning/thinking content for AI response events.

        Returns event properties for reasoning, respecting privacy tiers:

        - Always includes ``[Agent] Has Reasoning`` when content is present
          **or** when ``reasoning_tokens`` is positive (indicating the model
          used reasoning even if the content wasn't captured).
        - ``"full"``: includes raw text (truncated at 10K chars).
        - ``"metadata_only"`` / ``"customer_enriched"``: no text.
        - Reasoning tokens are always included when provided (all tiers).
        """
        from .tracking import (
            PROP_HAS_REASONING,
            PROP_REASONING_CONTENT,
            PROP_REASONING_TOKENS,
        )

        result: dict[str, Any] = {}

        has_reasoning = bool(reasoning_content) or (
            reasoning_tokens is not None and reasoning_tokens > 0
        )
        if not has_reasoning and reasoning_tokens is None:
            return result

        if has_reasoning:
            result[PROP_HAS_REASONING] = True

        if reasoning_tokens is not None:
            result[PROP_REASONING_TOKENS] = reasoning_tokens

        if not has_reasoning or reasoning_content is None:
            return result

        # Determine effective mode
        mode = self._content_mode
        if mode is None:
            mode = "metadata_only" if self.privacy_mode else "full"

        if mode == "full":
            # Apply PII + custom redaction, then truncate at 10K characters
            sanitized = reasoning_content
            if self.redact_pii:
                sanitized = redact_pii_patterns(sanitized)
            sanitized = self._apply_custom_patterns(sanitized)
            truncated = sanitized[:10000] if len(sanitized) > 10000 else sanitized
            result[PROP_REASONING_CONTENT] = truncated
        # metadata_only / customer_enriched → no content

        return result
