"""Serialization utilities for telemetry data capture.

This module provides safe JSON serialization for telemetry data with fallback
handling for non-serializable types and truncation for large payloads.
"""

import json
import logging
import reprlib
import sys
import types
from typing import Any, Dict, Optional

logger = logging.getLogger(__name__)

# Issue #9 Fix: Hard limit for collection sizes to prevent OOM on deep nesting
MAX_COLLECTION_LENGTH = 100000

# Issue #3 Complete Fix: Configure reprlib for safe stringification of nested structures
# Prevents OOM on deeply nested collections like [[0] * 100_000_000]
_SAFE_REPR = reprlib.Repr()
_SAFE_REPR.maxlevel = 5  # Maximum nesting depth
_SAFE_REPR.maxdict = 100  # Max dict items to show
_SAFE_REPR.maxlist = 100  # Max list items to show
_SAFE_REPR.maxtuple = 100  # Max tuple items to show
_SAFE_REPR.maxset = 100  # Max set items to show
_SAFE_REPR.maxstring = 1000  # Max string length in repr output


def serialize_for_telemetry(
    data: Any,
    max_length: int = 1048576,  # 1MB limit for large payloads
    truncate_marker: str = "...[truncated]",
) -> Any:
    """Serialize data to a telemetry-safe value.

    Returns basic Python types (None, bool, int, float, str, list, dict) unchanged
    since they are already OTel-compatible. Converts complex types to JSON strings
    with graceful fallback to str() for non-serializable objects. Truncates large
    string payloads to prevent excessive telemetry data.

    Args:
        data: Data to serialize (any type)
        max_length: Maximum length of serialized string (default: 1MB/1048576 bytes)
        truncate_marker: Marker to append when truncating (default: "...[truncated]")

    Returns:
        Original value for basic OTel-compatible types; JSON/str string for complex types

    Examples:
        >>> serialize_for_telemetry({"key": "value"})
        {'key': 'value'}

        >>> serialize_for_telemetry([1, 2, 3])
        [1, 2, 3]

        >>> serialize_for_telemetry("a" * 2000, max_length=100)
        'aaa...aaa...[truncated]'
    """
    # None, bool, int, float are returned unchanged (OTel-compatible)
    # Note: OTel span attributes DO NOT accept dicts, and only accept homogeneous
    # lists of primitives. To be safe for both logs and spans, we stringify dicts/lists.
    if data is None or isinstance(data, (bool, int, float)):
        return data

    # Strings are returned as-is but with truncation applied if needed
    if isinstance(data, str):
        if len(data) <= max_length:
            return data
        keep_length = max_length - len(truncate_marker)
        if keep_length > 0:
            return data[:keep_length] + truncate_marker
        return truncate_marker

    try:
        # Custom JSON encoder to handle special types (NumPy, Pandas, etc.)
        serialized = json.dumps(data, default=_json_default, ensure_ascii=False)
    except (TypeError, ValueError, OverflowError):
        # Fallback to string representation for non-serializable objects
        serialized = str(data)

    # Truncate if exceeds max length
    if len(serialized) > max_length:
        # Keep some content from the beginning and add truncation marker
        keep_length = max_length - len(truncate_marker)
        if keep_length > 0:
            serialized = serialized[:keep_length] + truncate_marker
        else:
            serialized = truncate_marker

    return serialized


def _json_default(obj: Any) -> Any:
    """Custom JSON encoder default handler for special types.

    Args:
        obj: Object that couldn't be serialized by default JSON encoder

    Returns:
        JSON-serializable representation of the object

    Raises:
        TypeError: If object cannot be converted to JSON-serializable form
    """
    obj_type = type(obj).__name__

    # Handle NumPy arrays
    if hasattr(obj, "tolist"):  # NumPy arrays have tolist() method
        try:
            result = obj.tolist()
            logger.debug(f"Serialized {obj_type} using tolist() method")
            return result
        except (AttributeError, TypeError, ValueError) as e:
            logger.warning(
                f"Failed to serialize {obj_type} using tolist(): {type(e).__name__}. "
                f"Trying next fallback method."
            )

    # Handle Pandas DataFrames/Series
    if hasattr(obj, "to_dict"):  # Pandas DataFrames/Series have to_dict() method
        try:
            result = obj.to_dict()
            logger.debug(f"Serialized {obj_type} using to_dict() method")
            return result
        except (AttributeError, TypeError, ValueError) as e:
            logger.warning(
                f"Failed to serialize {obj_type} using to_dict(): {type(e).__name__}. "
                f"Trying next fallback method."
            )

    # Handle dataclasses
    if hasattr(obj, "__dataclass_fields__"):
        try:
            result = {field: getattr(obj, field) for field in obj.__dataclass_fields__}
            logger.debug(f"Serialized dataclass {obj_type} using __dataclass_fields__")
            return result
        except (AttributeError, TypeError, ValueError) as e:
            logger.warning(
                f"Failed to serialize dataclass {obj_type}: {type(e).__name__}. "
                f"Trying next fallback method."
            )

    # Handle objects with __dict__ (custom classes)
    if hasattr(obj, "__dict__"):
        try:
            result = obj.__dict__
            logger.debug(f"Serialized {obj_type} using __dict__ attribute")
            return result
        except (AttributeError, TypeError, ValueError) as e:
            logger.warning(
                f"Failed to serialize {obj_type} using __dict__: {type(e).__name__}. "
                f"Using str() as final fallback."
            )

    # Final fallback: convert to string
    logger.debug(f"Serialized {obj_type} using str() fallback")
    return str(obj)


def _get_shape(data: Any) -> Optional[str]:
    """Extract shape information from data.

    Issue #4 Fix: Deep shape extraction for nested lists/tuples.
    Issue #5 Fix: Unwrap {"args": [...], "kwargs": {...}} structures.

    Args:
        data: Input or output data (numpy array, pandas DataFrame, etc.).

    Returns:
        String representation of shape, or None if shape cannot be determined.
    """
    try:
        # Issue #5 Fix: Unwrap decorator-wrapped multi-argument structures
        if isinstance(data, dict):
            # Handle {"args": [...], "kwargs": {...}} from decorator
            if set(data.keys()) == {"args", "kwargs"}:
                args = data.get("args", ())
                kwargs = data.get("kwargs", {})

                # Try args first
                if args and len(args) > 0:
                    return _get_shape(args[0])

                # Issue #5 Complete Fix: Check kwargs if args is empty
                # Common pattern: model.predict(X=data)
                if kwargs:
                    # Try common parameter names for feature matrices
                    for param_name in ["X", "data", "input", "features", "x"]:
                        if param_name in kwargs:
                            return _get_shape(kwargs[param_name])
                    # Fallback: use first kwarg value
                    first_kwarg = next(iter(kwargs.values()), None)
                    if first_kwarg is not None:
                        return _get_shape(first_kwarg)

                return None

            # Handle {"kwargs": {...}} from autolog when only kwargs passed
            elif set(data.keys()) == {"kwargs"}:
                kwargs = data.get("kwargs", {})
                if kwargs:
                    # Try common parameter names
                    for param_name in ["X", "data", "input", "features", "x"]:
                        if param_name in kwargs:
                            return _get_shape(kwargs[param_name])
                    # Fallback: use first kwarg value
                    first_kwarg = next(iter(kwargs.values()), None)
                    if first_kwarg is not None:
                        return _get_shape(first_kwarg)
                return None

        # Handle numpy arrays, pandas DataFrames, scipy sparse matrices
        if hasattr(data, "shape"):
            return str(data.shape)

        # Issue #4 Fix: Deep shape extraction for nested lists/tuples
        elif isinstance(data, (list, tuple)):
            if len(data) == 0:
                return "(0,)"

            # Recursively determine shape for nested structures
            def _get_nested_shape(obj, depth=0, max_depth=10):
                """Recursively calculate shape of nested list/tuple."""
                if depth > max_depth:
                    return [len(obj)]  # Prevent infinite recursion

                if not isinstance(obj, (list, tuple)) or len(obj) == 0:
                    return []

                shape = [len(obj)]
                # Check if all elements are also lists/tuples with consistent length
                if all(isinstance(item, (list, tuple)) for item in obj):
                    first_item_shape = _get_nested_shape(obj[0], depth + 1, max_depth)
                    # Verify all items have the same shape
                    if all(
                        isinstance(item, (list, tuple)) and len(item) == len(obj[0]) for item in obj
                    ):
                        shape.extend(first_item_shape)

                return shape

            shape_list = _get_nested_shape(data)
            return str(tuple(shape_list))

        else:
            return None
    except Exception:
        return None


def extract_model_payload(
    data: Any,
    max_size_bytes: int = 10000,
    max_object_size_multiplier: int = 10,
) -> Dict[str, Any]:
    """Extract serialized data and shape with safe size limits.

    This function provides unified data extraction for both manual and auto
    instrumentation. It always extracts shape information (lightweight) and
    safely serializes data with byte-level size checking to prevent OOM issues.

    Issue #3 Fix: Robust OOM protection with tiered size estimation.
    Issue #7 Fix: Truncation respects exact byte limits including message.
    Issue #9 Fix: Detect and handle generators/iterators without consuming them.
    Issue #11 Fix: Serialization errors logged at warning level, not debug.

    Args:
        data: Data to extract (numpy array, DataFrame, dict, list, etc.).
        max_size_bytes: Maximum size in bytes for serialized data (default: 10000).
        max_object_size_multiplier: Safety multiplier for object size check (default: 10).

    Returns:
        Dictionary with keys:
            - "data": Serialized string representation of data
            - "shape": Shape string (e.g., "(100, 5)") or None
            - "truncated": Boolean indicating if data was truncated

    Examples:
        >>> import numpy as np
        >>> payload = extract_model_payload(np.array([[1, 2], [3, 4]]))
        >>> payload["shape"]
        '(2, 2)'
        >>> payload["truncated"]
        False

        >>> large_data = "x" * 20000
        >>> payload = extract_model_payload(large_data, max_size_bytes=1000)
        >>> payload["truncated"]
        True
    """
    result: Dict[str, Any] = {
        "data": None,
        "shape": None,
        "truncated": False,
    }

    try:
        # Always extract shape (lightweight operation)
        shape_str = _get_shape(data)
        if shape_str:
            result["shape"] = shape_str

        # Issue #9 Fix: Detect generators and iterators early
        if isinstance(data, types.GeneratorType) or (
            hasattr(data, "__iter__") and not hasattr(data, "__len__")
        ):
            result["data"] = "[Unconsumable Iterator/Generator]"
            result["truncated"] = True
            logger.debug("Skipped serialization of unconsumable iterator/generator")
            return result

        # Issue #3 Fix: Robust OOM protection with tiered size estimation
        estimated_size = 0

        # For NumPy arrays: check actual memory usage
        if hasattr(data, "nbytes"):
            estimated_size = data.nbytes

        # For Pandas DataFrames: check memory usage
        elif hasattr(data, "memory_usage") and callable(data.memory_usage):
            try:
                estimated_size = data.memory_usage(deep=True).sum()
            except Exception:
                estimated_size = sys.getsizeof(data)

        # For collections: check length to prevent deep nested structures
        elif isinstance(data, (list, tuple, dict, set)):
            if len(data) > MAX_COLLECTION_LENGTH:
                result["data"] = (
                    f"[Collection too large to serialize: {len(data)} items exceeds "
                    f"limit of {MAX_COLLECTION_LENGTH}]"
                )
                result["truncated"] = True
                logger.debug(f"Skipped serialization of large collection: {len(data)} items")
                return result
            estimated_size = sys.getsizeof(data)

        # Fallback: use sys.getsizeof
        else:
            estimated_size = sys.getsizeof(data)

        # If estimated size is huge, skip string conversion entirely
        if estimated_size > max_size_bytes * max_object_size_multiplier:
            result["data"] = f"[Object too large to serialize: {estimated_size} bytes]"
            result["truncated"] = True
            logger.debug(
                f"Skipped data serialization: estimated size {estimated_size} bytes "
                f"exceeds safety threshold of {max_size_bytes * max_object_size_multiplier} bytes"
            )
            return result

        # Issue #3 Complete Fix: Use reprlib for safe stringification of collections
        # Prevents OOM on deeply nested structures like [[0] * 100_000_000]
        # reprlib safely truncates during serialization without evaluating entire structure
        try:
            if isinstance(data, (list, tuple, dict, set)):
                # Use reprlib for safe truncation - drops exec time from >5s to <1ms
                data_str = _SAFE_REPR.repr(data)
            else:
                # For primitives and NumPy/Pandas objects, use str()
                data_str = str(data)
        except Exception:
            # Fallback to str() if reprlib fails
            data_str = str(data)

        # Measure actual string byte size for consistent truncation
        # Use UTF-8 encoding length, not character count
        data_bytes = data_str.encode("utf-8", errors="replace")
        data_size_bytes = len(data_bytes)

        if data_size_bytes > max_size_bytes:
            # Issue #7 Fix: Truncate to fit within exact byte limit including message
            # Calculate truncation message size first
            trunc_msg = f"...[truncated, {data_size_bytes} bytes total]"
            trunc_msg_bytes = trunc_msg.encode("utf-8")
            allowed_data_bytes = max_size_bytes - len(trunc_msg_bytes)

            if allowed_data_bytes <= 0:
                # Not enough space for any data, just return the message
                result["data"] = trunc_msg
                result["truncated"] = True
                logger.debug(
                    f"Truncated data: {data_size_bytes} bytes exceeds limit, returning message only"
                )
            else:
                # Slice strictly by bytes, decode ignoring partial chars, then append message
                truncated_bytes = data_bytes[:allowed_data_bytes]
                safe_string = truncated_bytes.decode("utf-8", errors="ignore")
                result["data"] = safe_string + trunc_msg
                result["truncated"] = True
                logger.debug(
                    f"Truncated data: {data_size_bytes} bytes exceeds limit of {max_size_bytes} bytes"
                )
        else:
            result["data"] = data_str
            result["truncated"] = False

        return result

    except Exception as e:
        # Issue #11 Fix: Elevate serialization errors to warning level
        logger.warning(f"Failed to serialize model payload ({type(e).__name__}): {e}")
        # Return safe fallback
        result["data"] = "[Failed to serialize data]"
        result["truncated"] = True
        return result
