"""
Unit tests covering fixes applied in this session:

- Issue #2:  Race condition in init() — _init_lock prevents concurrent corruption
- Issue #3:  SpanBatcher worker thread crash recovery — try-except + re-queue
- Issue #4:  Monkey-patch re-entry guard — _patched_modules prevents double-patching
- Issue #8:  LangChain tool stack cleanup on error — pop on exception in on_tool_start

No framework dependencies or network required.
"""
import threading
import time
import uuid
from unittest.mock import MagicMock, patch

import pytest

import visibe
import visibe as _mod
from visibe.api import APIClient, SpanBatcher
from visibe.integrations.langchain import (
    LangGraphCallback,
    _tool_stack_ctx,
)


# ---------------------------------------------------------------------------
# Fixture: reset visibe global state around every test
# ---------------------------------------------------------------------------

@pytest.fixture(autouse=True)
def reset_visibe_globals():
    """Ensure clean global state before and after each test."""
    _mod._global_client = None
    _mod._auto_patched_frameworks.clear()
    _mod._patched_modules.clear()
    yield
    _mod._global_client = None
    _mod._auto_patched_frameworks.clear()
    _mod._patched_modules.clear()


# ===================================================================
# Issue #2 — init() race condition (threading.Lock)
# ===================================================================

def test_init_lock_exists():
    """_init_lock should be a threading.Lock instance."""
    assert isinstance(_mod._init_lock, type(threading.Lock()))


def test_concurrent_init_only_creates_one_client():
    """Multiple threads calling init() simultaneously must produce a single client."""
    results = []
    barrier = threading.Barrier(4)

    def worker():
        barrier.wait()
        client = visibe.init(
            api_key="sk_live_test",
            api_url="http://localhost:3000",
            frameworks=[],
        )
        results.append(id(client))

    threads = [threading.Thread(target=worker) for _ in range(4)]
    for t in threads:
        t.start()
    for t in threads:
        t.join(timeout=10)

    # All threads should get the exact same client instance
    assert len(set(results)) == 1


def test_concurrent_init_patches_only_once():
    """Even with concurrent init() calls, _auto_patched_frameworks should not have duplicates."""
    barrier = threading.Barrier(3)

    def worker():
        barrier.wait()
        visibe.init(
            api_key="sk_live_test",
            api_url="http://localhost:3000",
            frameworks=[],
        )

    threads = [threading.Thread(target=worker) for _ in range(3)]
    for t in threads:
        t.start()
    for t in threads:
        t.join(timeout=10)

    # No framework duplication
    assert len(_mod._auto_patched_frameworks) == len(set(_mod._auto_patched_frameworks))


# ===================================================================
# Issue #3 — SpanBatcher worker crash recovery
# ===================================================================

def _make_batcher(api_client=None, **kwargs):
    """Create a SpanBatcher with a mock API client."""
    if api_client is None:
        api_client = MagicMock(spec=APIClient)
        api_client._enabled = True
    return SpanBatcher(api_client, **kwargs)


def test_batcher_worker_survives_send_exception():
    """If _send raises, the worker thread should stay alive."""
    client = MagicMock(spec=APIClient)
    client._enabled = True
    # First batch send raises, second succeeds
    client.send_spans_batch.side_effect = [ConnectionError("boom"), True]

    batcher = _make_batcher(client, flush_interval=0.1)
    batcher.add("t1", {"span_id": "s1"})

    # Wait for the worker to attempt the flush and hit the error
    time.sleep(0.4)

    # Worker thread should still be alive
    assert batcher._thread.is_alive()

    batcher.shutdown()


def test_batcher_requeues_failed_batch():
    """Failed batch should be re-added to the buffer for retry."""
    client = MagicMock(spec=APIClient)
    client._enabled = True
    client.send_spans_batch.side_effect = ConnectionError("network error")

    batcher = _make_batcher(client, flush_interval=60.0)  # long interval so worker doesn't auto-flush
    batcher.add("t1", {"span_id": "s1"})
    batcher.add("t1", {"span_id": "s2"})

    # Manually trigger a worker cycle by signaling the event
    batcher._flush_event.set()
    time.sleep(0.3)

    # The failed spans should be back in the buffer
    with batcher._lock:
        span_ids = [s[1]["span_id"] for s in batcher._buffer]
    assert "s1" in span_ids
    assert "s2" in span_ids

    batcher._stopped = True
    batcher._flush_event.set()


def test_batcher_successful_send_clears_buffer():
    """After a successful send, buffer should be empty."""
    client = MagicMock(spec=APIClient)
    client._enabled = True
    client.send_spans_batch.return_value = True

    batcher = _make_batcher(client, flush_interval=0.1)
    batcher.add("t1", {"span_id": "s1"})

    # Wait for auto-flush
    time.sleep(0.4)

    with batcher._lock:
        assert len(batcher._buffer) == 0

    batcher.shutdown()


# ===================================================================
# Issue #4 — Monkey-patch re-entry guard (_patched_modules)
# ===================================================================

def test_patched_modules_initially_empty():
    """_patched_modules should start empty."""
    assert len(_mod._patched_modules) == 0


def test_patch_openai_sets_patched_flag(monkeypatch):
    """_patch_openai should add 'openai' to _patched_modules."""
    import types
    fake_openai = types.ModuleType('openai')

    class FakeOpenAI:
        def __init__(self): pass

    class FakeAsyncOpenAI:
        def __init__(self): pass

    fake_openai.OpenAI = FakeOpenAI
    fake_openai.AsyncOpenAI = FakeAsyncOpenAI
    monkeypatch.setitem(__import__('sys').modules, 'openai', fake_openai)

    client = MagicMock()
    client._get_openai_integration.return_value = MagicMock()

    _mod._patch_openai(client)
    assert 'openai' in _mod._patched_modules


def test_patch_openai_idempotent(monkeypatch):
    """Calling _patch_openai twice should not overwrite _original_openai_init."""
    import types
    fake_openai = types.ModuleType('openai')

    class FakeOpenAI:
        def __init__(self): pass

    class FakeAsyncOpenAI:
        def __init__(self): pass

    fake_openai.OpenAI = FakeOpenAI
    fake_openai.AsyncOpenAI = FakeAsyncOpenAI
    monkeypatch.setitem(__import__('sys').modules, 'openai', fake_openai)

    client = MagicMock()
    client._get_openai_integration.return_value = MagicMock()

    _mod._patch_openai(client)
    saved_original = _mod._original_openai_init

    # Patch again — should be a no-op, not overwrite
    _mod._patch_openai(client)
    assert _mod._original_openai_init is saved_original


def test_unpatch_openai_clears_patched_flag(monkeypatch):
    """_unpatch_openai should remove 'openai' from _patched_modules."""
    import types
    fake_openai = types.ModuleType('openai')

    class FakeOpenAI:
        def __init__(self): pass

    class FakeAsyncOpenAI:
        def __init__(self): pass

    fake_openai.OpenAI = FakeOpenAI
    fake_openai.AsyncOpenAI = FakeAsyncOpenAI
    monkeypatch.setitem(__import__('sys').modules, 'openai', fake_openai)

    client = MagicMock()
    client._get_openai_integration.return_value = MagicMock()

    _mod._patch_openai(client)
    assert 'openai' in _mod._patched_modules

    _mod._unpatch_openai()
    assert 'openai' not in _mod._patched_modules


def test_shutdown_clears_patched_modules():
    """shutdown() should clear _patched_modules."""
    _mod._patched_modules.add('openai')
    _mod._patched_modules.add('bedrock')
    visibe.init(api_key="sk_live_test", api_url="http://localhost:3000", frameworks=[])
    visibe.shutdown()
    assert len(_mod._patched_modules) == 0


def test_reinit_after_shutdown_can_repatch(monkeypatch):
    """After shutdown clears _patched_modules, init() should be able to patch again."""
    import types
    fake_openai = types.ModuleType('openai')

    class FakeOpenAI:
        def __init__(self): pass

    class FakeAsyncOpenAI:
        def __init__(self): pass

    fake_openai.OpenAI = FakeOpenAI
    fake_openai.AsyncOpenAI = FakeAsyncOpenAI
    monkeypatch.setitem(__import__('sys').modules, 'openai', fake_openai)

    client = MagicMock()
    client._get_openai_integration.return_value = MagicMock()

    _mod._patch_openai(client)
    assert 'openai' in _mod._patched_modules

    _mod._unpatch_openai()
    assert 'openai' not in _mod._patched_modules

    # Should be able to patch again
    _mod._patch_openai(client)
    assert 'openai' in _mod._patched_modules


# ===================================================================
# Issue #8 — LangChain tool stack cleanup on error
# ===================================================================

@pytest.fixture()
def clean_tool_stack():
    """Reset _tool_stack_ctx before and after each test."""
    token = _tool_stack_ctx.set(())
    yield
    _tool_stack_ctx.set(())


def test_tool_stack_pushed_on_tool_start(clean_tool_stack):
    """on_tool_start should push tool name onto _tool_stack_ctx."""
    cb = LangGraphCallback()
    run_id = uuid.uuid4()

    cb.on_tool_start(
        {"name": "search"},
        "query",
        run_id=run_id,
    )

    assert _tool_stack_ctx.get() == ("search",)


def test_tool_stack_popped_on_tool_end(clean_tool_stack):
    """on_tool_end should pop tool name from _tool_stack_ctx."""
    cb = LangGraphCallback()
    run_id = uuid.uuid4()

    cb.on_tool_start({"name": "search"}, "query", run_id=run_id)
    cb.on_tool_end("result", run_id=run_id)

    assert _tool_stack_ctx.get() == ()


def test_tool_stack_popped_on_tool_error(clean_tool_stack):
    """on_tool_error should pop tool name from _tool_stack_ctx."""
    cb = LangGraphCallback()
    run_id = uuid.uuid4()

    cb.on_tool_start({"name": "search"}, "query", run_id=run_id)
    cb.on_tool_error(RuntimeError("fail"), run_id=run_id)

    assert _tool_stack_ctx.get() == ()


def test_tool_stack_cleaned_on_tool_start_exception(clean_tool_stack):
    """If super().on_tool_start() raises, the stack push should be reverted."""
    cb = LangGraphCallback()
    run_id = uuid.uuid4()

    # Make the parent on_tool_start raise
    with patch.object(
        cb.__class__.__bases__[0],  # LangChainCallback
        'on_tool_start',
        side_effect=RuntimeError("callback error"),
    ):
        with pytest.raises(RuntimeError, match="callback error"):
            cb.on_tool_start({"name": "broken_tool"}, "input", run_id=run_id)

    # Stack should be clean — the push was reverted
    assert _tool_stack_ctx.get() == ()


def test_tool_stack_nested_cleanup_on_error(clean_tool_stack):
    """Nested tool stack: error in inner tool should only pop the inner entry."""
    cb = LangGraphCallback()
    run_id_outer = uuid.uuid4()
    run_id_inner = uuid.uuid4()

    # Push outer tool successfully
    cb.on_tool_start({"name": "outer"}, "input", run_id=run_id_outer)
    assert _tool_stack_ctx.get() == ("outer",)

    # Inner tool errors via on_tool_error callback
    cb.on_tool_start({"name": "inner"}, "input", run_id=run_id_inner)
    assert _tool_stack_ctx.get() == ("outer", "inner")

    cb.on_tool_error(RuntimeError("inner failed"), run_id=run_id_inner)
    assert _tool_stack_ctx.get() == ("outer",)

    # Outer tool completes normally
    cb.on_tool_end("result", run_id=run_id_outer)
    assert _tool_stack_ctx.get() == ()


def test_tool_stack_hierarchical_name(clean_tool_stack):
    """Nested tools should produce hierarchical agent names."""
    cb = LangGraphCallback()
    run_id1 = uuid.uuid4()
    run_id2 = uuid.uuid4()

    cb.on_tool_start({"name": "plan"}, "input", run_id=run_id1)
    cb.on_tool_start({"name": "research"}, "input", run_id=run_id2)

    assert _tool_stack_ctx.get() == ("plan", "research")
    assert cb.current_agent == "plan/research"
