import unittest
from unittest.mock import patch

from posthog.contexts import (
    get_tags,
    new_context,
    scoped,
    tag,
    identify_context,
    set_context_session,
    get_context_session_id,
    get_context_distinct_id,
)


class TestContexts(unittest.TestCase):
    def test_tag_and_get_tags(self):
        with new_context(fresh=True):
            tag("key1", "value1")
            tag("key2", 2)

            tags = get_tags()
            assert tags["key1"] == "value1"
            assert tags["key2"] == 2

    def test_new_context_isolation(self):
        with new_context(fresh=True):
            # Set tag in outer context
            tag("outer", "value")

            with new_context(fresh=True):
                # Inner context should start empty
                assert get_tags() == {}

                # Set tag in inner context
                tag("inner", "value")
                assert get_tags()["inner"] == "value"

                # Outer tag should not be visible
                self.assertNotIn("outer", get_tags())

            with new_context(fresh=False):
                # Inner context should inherit outer tag
                assert get_tags() == {"outer": "value"}

            # After exiting context, inner tag should be gone
            self.assertNotIn("inner", get_tags())

            # Outer tag should still be there
            assert get_tags()["outer"] == "value"

    def test_nested_contexts(self):
        with new_context(fresh=True):
            tag("level1", "value1")

            with new_context(fresh=True):
                tag("level2", "value2")

                with new_context(fresh=True):
                    tag("level3", "value3")
                    assert get_tags() == {"level3": "value3"}

                # Back to level 2
                assert get_tags() == {"level2": "value2"}

            # Back to level 1
            assert get_tags() == {"level1": "value1"}

    @patch("posthog.capture_exception")
    def test_scoped_decorator_success(self, mock_capture):
        @scoped()
        def successful_function(x, y):
            tag("x", x)
            tag("y", y)
            return x + y

        result = successful_function(1, 2)

        # Function should execute normally
        assert result == 3

        # No exception should be captured
        mock_capture.assert_not_called()

        # Context should be cleared after function execution
        assert get_tags() == {}

    @patch("posthog.capture_exception")
    def test_scoped_decorator_exception(self, mock_capture):
        test_exception = ValueError("Test exception")

        def check_context_on_capture(exception, **kwargs):
            # Assert tags are available when capture_exception is called
            current_tags = get_tags()
            assert current_tags.get("important_context") == "value"

        mock_capture.side_effect = check_context_on_capture

        @scoped()
        def failing_function():
            tag("important_context", "value")
            raise test_exception

        # Function should raise the exception
        with self.assertRaises(ValueError):
            failing_function()

        # Verify capture_exception was called
        mock_capture.assert_called_once_with(test_exception)

        # Context should be cleared after function execution
        assert get_tags() == {}

    @patch("posthog.capture_exception")
    def test_new_context_exception_handling(self, mock_capture):
        test_exception = RuntimeError("Context exception")

        def check_context_on_capture(exception, **kwargs):
            # Assert inner context tags are available when capture_exception is called
            current_tags = get_tags()
            assert current_tags.get("inner_context") == "inner_value"

        mock_capture.side_effect = check_context_on_capture

        # Set up outer context
        with new_context():
            tag("outer_context", "outer_value")

            try:
                with new_context():
                    tag("inner_context", "inner_value")
                    raise test_exception
            except RuntimeError:
                pass  # Expected exception

            # Outer context should still be intact
            assert get_tags()["outer_context"] == "outer_value"

        # Verify capture_exception was called
        mock_capture.assert_called_once_with(test_exception)

    def test_identify_context(self):
        with new_context(fresh=True):
            # Initially no distinct ID
            assert get_context_distinct_id() is None

            # Set distinct ID
            identify_context("user123")
            assert get_context_distinct_id() == "user123"

    def test_set_context_session(self):
        with new_context(fresh=True):
            # Initially no session ID
            assert get_context_session_id() is None

            # Set session ID
            set_context_session("session456")
            assert get_context_session_id() == "session456"

    def test_context_inheritance_fresh_context(self):
        with new_context(fresh=True):
            identify_context("user123")
            set_context_session("session456")

            with new_context(fresh=True):
                # Fresh context should not inherit
                assert get_context_distinct_id() is None
                assert get_context_session_id() is None

            # Original context should still have values
            assert get_context_distinct_id() == "user123"
            assert get_context_session_id() == "session456"

    def test_context_inheritance_non_fresh_context(self):
        with new_context(fresh=True):
            identify_context("user123")
            set_context_session("session456")

            with new_context(fresh=False):
                # Non-fresh context should inherit
                assert get_context_distinct_id() == "user123"
                assert get_context_session_id() == "session456"

                # Override in child context
                identify_context("user789")
                set_context_session("session999")
                assert get_context_distinct_id() == "user789"
                assert get_context_session_id() == "session999"

            # Original context should still have original values
            assert get_context_distinct_id() == "user123"
            assert get_context_session_id() == "session456"

    def test_scoped_decorator_with_context_ids(self):
        @scoped()
        def function_with_context():
            identify_context("user456")
            set_context_session("session789")
            return get_context_distinct_id(), get_context_session_id()

        distinct_id, session_id = function_with_context()
        assert distinct_id == "user456"
        assert session_id == "session789"

        # Context should be cleared after function execution
        assert get_context_distinct_id() is None
        assert get_context_session_id() is None