from __future__ import annotations

import pickle
import socket
import sqlite3

import dill
import pytest

from codeflash_python.runtime._picklepatch.pickle_patcher import PicklePatcher
from codeflash_python.runtime._picklepatch.pickle_placeholder import (
    PicklePlaceholder,
    PicklePlaceholderAccessError,
)


def test_picklepatch_simple_nested():
    """Test that a simple nested data structure pickles and unpickles correctly."""
    original_data = {
        "numbers": [1, 2, 3],
        "nested_dict": {"key": "value", "another": 42},
    }

    dumped = PicklePatcher.dumps(original_data)
    reloaded = PicklePatcher.loads(dumped)

    assert reloaded == original_data
    # Everything was pickleable, so no placeholders should appear.


def test_picklepatch_with_socket():
    """Test that a data structure containing a raw socket is replaced by
    PicklePlaceholder rather than raising an error.
    """
    # Create a pair of connected sockets instead of a single socket
    sock1, sock2 = socket.socketpair()

    data_with_socket = {"safe_value": 123, "raw_socket": sock1}

    # Send a message through sock1, which can be received by sock2
    sock1.send(b"Hello, world!")
    received = sock2.recv(1024)
    assert received == b"Hello, world!"
    # Pickle the data structure containing the socket
    dumped = PicklePatcher.dumps(data_with_socket)
    reloaded = PicklePatcher.loads(dumped)

    # We expect "raw_socket" to be replaced by a placeholder
    assert isinstance(reloaded, dict)
    assert reloaded["safe_value"] == 123
    assert isinstance(reloaded["raw_socket"], PicklePlaceholder)

    # Attempting to use or access attributes => AttributeError
    # (not RuntimeError as in original tests, our implementation uses AttributeError)
    with pytest.raises(PicklePlaceholderAccessError):
        reloaded["raw_socket"].recv(1024)

    # Clean up by closing both sockets
    sock1.close()
    sock2.close()


def test_picklepatch_deeply_nested():
    """Test that deep nesting with unpicklable objects works correctly."""
    # Create a deeply nested structure with an unpicklable object
    deep_nested = {
        "level1": {
            "level2": {
                "level3": {
                    "normal": "value",
                    "socket": socket.socket(
                        socket.AF_INET, socket.SOCK_STREAM
                    ),
                }
            }
        }
    }

    dumped = PicklePatcher.dumps(deep_nested)
    reloaded = PicklePatcher.loads(dumped)

    # We should be able to access the normal value
    assert reloaded["level1"]["level2"]["level3"]["normal"] == "value"

    # The socket should be replaced with a placeholder
    assert isinstance(
        reloaded["level1"]["level2"]["level3"]["socket"], PicklePlaceholder
    )


def test_picklepatch_class_with_unpicklable_attr():
    """Test that a class with an unpicklable attribute works correctly."""

    class TestClass:
        def __init__(self):
            self.normal = "normal value"
            self.unpicklable = socket.socket(
                socket.AF_INET, socket.SOCK_STREAM
            )

    obj = TestClass()

    dumped = PicklePatcher.dumps(obj)
    reloaded = PicklePatcher.loads(dumped)

    # Normal attribute should be preserved
    assert reloaded.normal == "normal value"

    # Unpicklable attribute should be replaced with a placeholder
    assert isinstance(reloaded.unpicklable, PicklePlaceholder)


def test_picklepatch_with_database_connection():
    """Test that a data structure containing a database connection is replaced
    by PicklePlaceholder rather than raising an error.
    """
    # SQLite connection - not pickleable
    conn = sqlite3.connect(":memory:")
    cursor = conn.cursor()

    data_with_db = {
        "description": "Database connection",
        "connection": conn,
        "cursor": cursor,
    }

    dumped = PicklePatcher.dumps(data_with_db)
    reloaded = PicklePatcher.loads(dumped)

    # Both connection and cursor should become placeholders
    assert isinstance(reloaded, dict)
    assert reloaded["description"] == "Database connection"
    assert isinstance(reloaded["connection"], PicklePlaceholder)
    assert isinstance(reloaded["cursor"], PicklePlaceholder)

    # Attempting to use attributes => AttributeError
    with pytest.raises(PicklePlaceholderAccessError):
        reloaded["connection"].execute("SELECT 1")

    cursor.close()
    conn.close()


def test_picklepatch_with_generator():
    """Test that a data structure containing a generator is replaced by
    PicklePlaceholder rather than raising an error.
    """

    def simple_generator():
        yield 1
        yield 2
        yield 3

    # Create a generator
    gen = simple_generator()

    # Put it in a data structure
    data_with_generator = {
        "description": "Contains a generator",
        "generator": gen,
        "normal_list": [1, 2, 3],
    }

    dumped = PicklePatcher.dumps(data_with_generator)
    reloaded = PicklePatcher.loads(dumped)

    # Generator should be replaced with a placeholder
    assert isinstance(reloaded, dict)
    assert reloaded["description"] == "Contains a generator"
    assert reloaded["normal_list"] == [1, 2, 3]
    assert isinstance(reloaded["generator"], PicklePlaceholder)

    # Attempting to use the generator => AttributeError
    with pytest.raises(TypeError):
        next(reloaded["generator"])

    # Attempting to call methods on the generator => AttributeError
    with pytest.raises(PicklePlaceholderAccessError):
        reloaded["generator"].send(None)


def test_picklepatch_loads_standard_pickle():
    """Test that PicklePatcher.loads can correctly load data that was pickled
    using the standard pickle module.
    """
    # Create a simple data structure
    original_data = {
        "numbers": [1, 2, 3],
        "nested_dict": {"key": "value", "another": 42},
        "tuple": (1, "two", 3.0),
    }

    # Pickle it with standard pickle
    pickled_data = pickle.dumps(original_data)

    # Load with PicklePatcher
    reloaded = PicklePatcher.loads(pickled_data)

    # Verify the data is correctly loaded
    assert reloaded == original_data
    assert isinstance(reloaded, dict)
    assert reloaded["numbers"] == [1, 2, 3]
    assert reloaded["nested_dict"]["key"] == "value"
    assert reloaded["tuple"] == (1, "two", 3.0)


def test_picklepatch_loads_dill_pickle():
    """Test that PicklePatcher.loads can correctly load data that was pickled
    using the dill module, which can pickle more complex objects than the
    standard pickle module.
    """
    # Create a more complex data structure that includes a lambda function
    # which dill can handle but standard pickle cannot
    original_data = {
        "numbers": [1, 2, 3],
        "function": lambda x: x * 2,
        "nested": {"another_function": lambda y: y**2},
    }

    # Pickle it with dill
    dilled_data = dill.dumps(original_data)

    # Load with PicklePatcher
    reloaded = PicklePatcher.loads(dilled_data)

    # Verify the data structure
    assert isinstance(reloaded, dict)
    assert reloaded["numbers"] == [1, 2, 3]

    # Test that the functions actually work
    assert reloaded["function"](5) == 10
    assert reloaded["nested"]["another_function"](4) == 16
