"""Test for goal tracking context token fix (v0.6.1).

This test specifically validates the bug fix for context token handling
in record_goal_started() and record_goal_completed().

Bug: record_goal_started was storing a context manager instead of a Token,
causing TypeError when record_goal_completed tried to detach it.

Fix: Changed to use context.attach(trace.set_span_in_context(span)) for
proper Token handling.
"""

import pytest
import time
from mca_sdk import MCAClient
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry import context as otel_context


def test_goal_tracking_context_token_detach():
    """Test that context token is properly detached without TypeError.

    Regression test for v0.6.1 bug fix.

    Previous error:
        TypeError: expected an instance of Token, got _AgnosticContextManager
    """
    # Create client with in-memory span exporter
    span_exporter = InMemorySpanExporter()
    tracer_provider = TracerProvider()
    tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter))

    client = MCAClient(
        service_name="test-goal-context",
        model_id="test-001",
        team_name="test-team",
        collector_endpoint="http://localhost:4318",
        allow_insecure_collector=True,
    )

    # Inject test provider
    client._tracer_provider = tracer_provider
    client._tracer = tracer_provider.get_tracer(__name__)

    try:
        # Start a goal - this should create and attach a proper Token
        goal_id = client.record_goal_started(
            "Test goal for context token validation", goal_type="test"
        )

        assert goal_id is not None
        assert isinstance(goal_id, str)

        # Check that context was properly stored
        assert goal_id in client._active_goals
        span, start_time, context_token = client._active_goals[goal_id]

        # Verify the token is actually a Token object, not a context manager
        # Token objects have a var attribute
        assert hasattr(context_token, "var"), f"Expected Token object, got {type(context_token)}"

        # Simulate some work
        time.sleep(0.01)

        # Complete the goal - this should detach the token WITHOUT error
        # This is where the bug would occur with TypeError
        client.record_goal_completed(goal_id=goal_id, status="success")

        # Verify goal was removed from active goals
        assert goal_id not in client._active_goals

        # Verify spans were created
        spans = span_exporter.get_finished_spans()
        assert len(spans) > 0

        goal_spans = [s for s in spans if s.name == "agent.goal"]
        assert len(goal_spans) == 1

        goal_span = goal_spans[0]
        assert goal_span.attributes.get("goal_id") == goal_id
        assert goal_span.attributes.get("goal_type") == "test"
        assert goal_span.attributes.get("goal_status") == "success"

    finally:
        client.shutdown()


def test_multiple_concurrent_goals_context_isolation():
    """Test that multiple concurrent goals have isolated contexts.

    Ensures that context tokens don't interfere with each other.
    """
    client = MCAClient(
        service_name="test-concurrent-goals",
        model_id="test-002",
        team_name="test-team",
        collector_endpoint="http://localhost:4318",
        allow_insecure_collector=True,
    )

    try:
        # Start multiple goals
        goal_ids = []
        for i in range(3):
            goal_id = client.record_goal_started(
                f"Concurrent goal {i}", goal_type="concurrent_test"
            )
            goal_ids.append(goal_id)

        # Verify all have unique tokens
        tokens = set()
        for goal_id in goal_ids:
            _, _, token = client._active_goals[goal_id]
            tokens.add(id(token))  # Use id() to check object identity

        assert len(tokens) == 3, "Each goal should have a unique context token"

        # Complete in different order
        client.record_goal_completed(goal_ids[1], status="success")
        client.record_goal_completed(goal_ids[0], status="success")
        client.record_goal_completed(goal_ids[2], status="success")

        # Verify all completed without errors
        assert len(client._active_goals) == 0

    finally:
        client.shutdown()


if __name__ == "__main__":
    # Run tests directly
    print("Testing goal context token fix...")
    test_goal_tracking_context_token_detach()
    print("✓ Single goal test passed")

    test_multiple_concurrent_goals_context_isolation()
    print("✓ Concurrent goals test passed")

    print("\n✅ All context token tests passed!")
