"""
Utilities for normalizing and unwrapping structured JSON input/context fields.

The API stores input and context as JSON objects:
- input: {"user_message": "..."} for simple strings, or structured objects with additional keys
- context: {"value": "..."} for simple strings, or structured objects with additional keys

These utilities handle bidirectional conversion between the string format
used in the SDK's public API and the JSON format used by the API.
"""

from typing import Any, Optional, Union


def normalize_input_to_json(value: Any) -> Optional[dict[str, Any]]:
    """Wraps a string input as {"user_message": "..."}, passes dicts through.

    Args:
        value: The input value (string, dict, or None).

    Returns:
        A JSON-compatible dict or None.
    """
    if value is None:
        return None
    if isinstance(value, str):
        return {"user_message": value}
    if isinstance(value, dict):
        return value
    return None


def normalize_context_to_json(value: Any) -> Optional[dict[str, Any]]:
    """Wraps a string context as {"value": "..."}, passes dicts through.

    Args:
        value: The context value (string, dict, or None).

    Returns:
        A JSON-compatible dict or None.
    """
    if value is None:
        return None
    if isinstance(value, str):
        return {"value": value}
    if isinstance(value, dict):
        return value
    return None


def unwrap_input(input_value: Optional[dict[str, Any]]) -> Union[str, dict[str, Any], None]:
    """Unwraps a JSON input for backward compatibility.

    - If the input only contains "user_message", returns the raw string.
    - If the input has additional keys (structured input), returns the full object.
    - If None, returns None.

    Args:
        input_value: The JSON input object from the API.

    Returns:
        The unwrapped string, the full dict, or None.
    """
    if input_value is None:
        return None
    if isinstance(input_value, str):
        return input_value
    if isinstance(input_value, dict):
        if "user_message" in input_value and len(input_value) == 1:
            user_message = input_value["user_message"]
            if isinstance(user_message, str):
                return user_message
        return input_value
    return None


def unwrap_context(context_value: Optional[dict[str, Any]]) -> Union[str, dict[str, Any], None]:
    """Unwraps a JSON context for backward compatibility.

    - If the context only contains "value", returns the raw string.
    - If the context has additional keys (structured context), returns the full object.
    - If None, returns None.

    Args:
        context_value: The JSON context object from the API.

    Returns:
        The unwrapped string, the full dict, or None.
    """
    if context_value is None:
        return None
    if isinstance(context_value, str):
        return context_value
    if isinstance(context_value, dict):
        if "value" in context_value and len(context_value) == 1:
            val = context_value["value"]
            if isinstance(val, str):
                return val
        return context_value
    return None


def coerce_to_string(value: Union[str, dict[str, Any], None]) -> Optional[str]:
    """Coerces an unwrapped value to a string.

    - If string, returns as-is.
    - If dict, JSON-serializes it.
    - If None, returns None.

    Args:
        value: The unwrapped value.

    Returns:
        A string representation or None.
    """
    if value is None:
        return None
    if isinstance(value, str):
        return value
    if isinstance(value, dict):
        import json

        return json.dumps(value)
    return None
