"""
SnakeBridge Type Encoding System

Provides encoding and decoding functions for converting between Python and Elixir types.
Uses tagged JSON format with __type__ markers for special types that don't map directly to JSON.

Supported special types:
- atom: Elixir atoms (encoded with __type__ tag)
- tuple: Python tuples (encoded as lists with __type__ tag)
- set: Python sets (encoded as lists with __type__ tag)
- frozenset: Python frozensets (encoded as lists with __type__ tag)
- bytes: Python bytes (encoded as base64 strings with __type__ tag)
- complex: Python complex numbers (encoded as [real, imag] with __type__ tag)
- datetime: datetime objects (encoded as ISO 8601 strings with __type__ tag)
- date: date objects (encoded as ISO 8601 strings with __type__ tag)
- time: time objects (encoded as ISO 8601 strings with __type__ tag)
- special_float: Positive infinity, negative infinity, and NaN (encoded with __type__ tag)
"""

import base64
import math
from datetime import datetime, date, time
from typing import Any, Dict, List, Union

SCHEMA_VERSION = 1


class Atom:
    """Represents an Elixir atom value on the Python side."""

    def __init__(self, value: str):
        self.value = str(value)

    def __repr__(self) -> str:
        return f"Atom({self.value!r})"

    def __str__(self) -> str:
        return self.value


def _tag(type_tag: str, payload: Dict[str, Any]) -> Dict[str, Any]:
    tagged = dict(payload)
    tagged["__type__"] = type_tag
    tagged["__schema__"] = SCHEMA_VERSION
    return tagged


def encode(value: Any) -> Any:
    """
    Encode a Python value to a JSON-safe value with __type__ tags for special types.

    Args:
        value: The Python value to encode

    Returns:
        A JSON-safe value (possibly with __type__ tags)

    Examples:
        >>> encode((1, 2, 3))
        {'__type__': 'tuple', '__schema__': 1, 'elements': [1, 2, 3]}

        >>> encode({1, 2, 3})
        {'__type__': 'set', '__schema__': 1, 'elements': [1, 2, 3]}

        >>> encode(b'hello')
        {'__type__': 'bytes', '__schema__': 1, 'data': 'aGVsbG8='}

        >>> encode(1+2j)
        {'__type__': 'complex', '__schema__': 1, 'real': 1.0, 'imag': 2.0}
    """
    # Handle None
    if value is None:
        return None

    # Handle basic JSON-safe types
    if isinstance(value, (bool, int, str)):
        return value

    # Handle float (check for special values)
    if isinstance(value, float):
        if math.isinf(value):
            if value > 0:
                return _tag("special_float", {"value": "infinity"})
            else:
                return _tag("special_float", {"value": "neg_infinity"})
        elif math.isnan(value):
            return _tag("special_float", {"value": "nan"})
        else:
            return value

    # Handle tagged atoms
    if isinstance(value, Atom):
        return _tag("atom", {"value": value.value})

    # Handle tuple
    if isinstance(value, tuple):
        return _tag("tuple", {"elements": [encode(item) for item in value]})

    # Handle set
    if isinstance(value, set):
        return _tag(
            "set",
            {
                "elements": [
                    encode(item) for item in sorted(value, key=lambda x: (type(x).__name__, str(x)))
                ]
            },
        )

    # Handle frozenset
    if isinstance(value, frozenset):
        return _tag(
            "frozenset",
            {
                "elements": [
                    encode(item) for item in sorted(value, key=lambda x: (type(x).__name__, str(x)))
                ]
            },
        )

    # Handle bytes
    if isinstance(value, bytes):
        return _tag("bytes", {"data": base64.b64encode(value).decode("ascii")})

    # Handle bytearray
    if isinstance(value, bytearray):
        return _tag("bytes", {"data": base64.b64encode(bytes(value)).decode("ascii")})

    # Handle complex
    if isinstance(value, complex):
        return _tag("complex", {"real": value.real, "imag": value.imag})

    # Handle datetime
    if isinstance(value, datetime):
        return _tag("datetime", {"value": value.isoformat()})

    # Handle date
    if isinstance(value, date):
        return _tag("date", {"value": value.isoformat()})

    # Handle time
    if isinstance(value, time):
        return _tag("time", {"value": value.isoformat()})

    # Handle list
    if isinstance(value, list):
        return [encode(item) for item in value]

    # Handle dict
    if isinstance(value, dict):
        return {encode_dict_key(k): encode(v) for k, v in value.items()}

    # For any other type, try to convert to string
    try:
        return str(value)
    except Exception:
        return f"<non-serializable: {type(value).__name__}>"


def encode_dict_key(key: Any) -> str:
    """
    Encode a dictionary key to a string.

    Args:
        key: The dictionary key to encode

    Returns:
        A string representation of the key
    """
    if isinstance(key, str):
        return key
    elif isinstance(key, (int, float, bool)):
        return str(key)
    elif isinstance(key, tuple):
        return str(key)
    else:
        return repr(key)


def decode(value: Any) -> Any:
    """
    Decode a JSON value with __type__ tags back to Python values.

    Args:
        value: The JSON value to decode (possibly with __type__ tags)

    Returns:
        The decoded Python value

    Examples:
        >>> decode({'__type__': 'tuple', '__schema__': 1, 'elements': [1, 2, 3]})
        (1, 2, 3)

        >>> decode({'__type__': 'set', '__schema__': 1, 'elements': [1, 2, 3]})
        {1, 2, 3}

        >>> decode({'__type__': 'bytes', '__schema__': 1, 'data': 'aGVsbG8='})
        b'hello'

        >>> decode({'__type__': 'complex', '__schema__': 1, 'real': 1.0, 'imag': 2.0})
        (1+2j)
    """
    # Handle None
    if value is None:
        return None

    # Handle basic types
    if isinstance(value, (bool, int, float, str)):
        return value

    # Handle list
    if isinstance(value, list):
        return [decode(item) for item in value]

    # Handle dict (check for __type__ tag)
    if isinstance(value, dict):
        if "__type__" in value:
            type_tag = value["__type__"]

            if type_tag == "atom":
                atom_value = value.get("value")
                if atom_value is None:
                    return value
                return Atom(atom_value)

            if type_tag == "tuple":
                elements = value.get("elements") or value.get("value") or []
                return tuple(decode(item) for item in elements)

            elif type_tag == "set":
                elements = value.get("elements") or value.get("value") or []
                return set(decode(item) for item in elements)

            elif type_tag == "frozenset":
                elements = value.get("elements") or value.get("value") or []
                return frozenset(decode(item) for item in elements)

            elif type_tag == "bytes":
                data = value.get("data") or value.get("value")
                if data is None:
                    return value
                return base64.b64decode(data)

            elif type_tag == "complex":
                return complex(value["real"], value["imag"])

            elif type_tag == "datetime":
                return datetime.fromisoformat(value["value"])

            elif type_tag == "date":
                return date.fromisoformat(value["value"])

            elif type_tag == "time":
                return time.fromisoformat(value["value"])

            elif type_tag == "special_float":
                special = value.get("value")
                if special == "infinity":
                    return float("inf")
                if special == "neg_infinity":
                    return float("-inf")
                if special == "nan":
                    return float("nan")
                return value

            elif type_tag == "infinity":
                return float("inf")

            elif type_tag == "neg_infinity":
                return float("-inf")

            elif type_tag == "nan":
                return float("nan")

            else:
                # Unknown type tag, return as-is
                return {k: decode(v) for k, v in value.items()}
        else:
            # Regular dict, decode recursively
            return {k: decode(v) for k, v in value.items()}

    # Return as-is for any other type
    return value


# Convenience functions for common operations
def encode_args(args: Dict[str, Any]) -> Dict[str, Any]:
    """
    Encode a dictionary of function arguments.

    Args:
        args: Dictionary of argument names to values

    Returns:
        Dictionary with encoded values
    """
    return {name: encode(value) for name, value in args.items()}


def decode_args(args: Dict[str, Any]) -> Dict[str, Any]:
    """
    Decode a dictionary of function arguments.

    Args:
        args: Dictionary of argument names to encoded values

    Returns:
        Dictionary with decoded values
    """
    return {name: decode(value) for name, value in args.items()}


def encode_result(result: Any) -> Dict[str, Any]:
    """
    Encode a function result in SnakeBridge result format.

    Args:
        result: The result value to encode

    Returns:
        Dictionary with success status and encoded result
    """
    return {
        "success": True,
        "result": encode(result)
    }


def encode_error(error: Exception) -> Dict[str, Any]:
    """
    Encode an error in SnakeBridge result format.

    Args:
        error: The exception to encode

    Returns:
        Dictionary with success status and error information
    """
    return {
        "success": False,
        "error": str(error),
        "error_type": type(error).__name__
    }