"""
Unit tests for flask_middleware() and starlette_middleware().

Flask and Starlette are mocked at the sys.modules level so these tests
run without either framework installed.  The key behaviours tested:

Flask:
  - create_trace + complete_trace called on successful request
  - status='failed' for 4xx responses
  - status='failed' on unhandled exception (teardown_request path)
  - double-call guard (_visibe_done) prevents two complete_trace calls
  - retry on complete_trace failure (background thread)
  - no tracing when _enabled=False
  - first_framework/model/prompt forwarded in payload

Starlette:
  - create_trace + complete_trace called on successful request
  - status='failed' for 4xx responses
  - status='failed' on unhandled exception
  - complete_trace waits for create_trace (threading.Event ordering fix)
  - retry on complete_trace failure
  - no tracing when _enabled=False
  - first_framework/model/prompt forwarded in payload
"""
import asyncio
import sys
import threading
import time
import types
import pytest

from visibe.client import Visibe
from visibe.context import _SessionContext, _active_session


# ------------------------------------------------------------------
# Shared helpers
# ------------------------------------------------------------------

class _MockAPI:
    """Minimal stand-in for APIClient — records calls, supports failure injection."""

    def __init__(self):
        self._enabled = True
        self.create_calls = []
        self.complete_calls = []
        self._complete_raise_n = 0   # raise on the first N calls to complete_trace

    def create_trace(self, data):
        self.create_calls.append(data)
        return True

    def complete_trace(self, trace_id, data):
        self.complete_calls.append((trace_id, data))
        if self._complete_raise_n > 0:
            self._complete_raise_n -= 1
            raise RuntimeError("simulated network error")
        return True


def _make_visibe():
    obs = Visibe(api_key="sk_live_abc", api_url="http://localhost:3000")
    obs.api_client = _MockAPI()
    obs._batcher = types.SimpleNamespace(flush=lambda: None, add=lambda *a: None)
    return obs


# ------------------------------------------------------------------
# Flask mock infrastructure
# ------------------------------------------------------------------

class _MockFlaskApp:
    """Records @before_request / @after_request / @teardown_request hooks."""

    def __init__(self):
        self._before = []
        self._after = []
        self._teardown = []

    def before_request(self, fn):
        self._before.append(fn)
        return fn

    def after_request(self, fn):
        self._after.append(fn)
        return fn

    def teardown_request(self, fn):
        self._teardown.append(fn)
        return fn

    # -- simulation helpers --

    def trigger_before(self):
        for fn in self._before:
            fn()

    def trigger_after(self, response):
        result = response
        for fn in self._after:
            result = fn(result)
        return result

    def trigger_teardown(self, exc=None):
        for fn in self._teardown:
            fn(exc)


class _MockResponse:
    def __init__(self, status_code=200):
        self.status_code = status_code


def _flask_setup(monkeypatch, obs):
    """
    Install a fake 'flask' module and return (app, fake_g).
    flask_middleware() imports 'flask' lazily, so this must be called
    before flask_middleware() is invoked.
    """
    fake_g = types.SimpleNamespace()
    fake_request = types.SimpleNamespace(method='GET', path='/api/chat')

    fake_flask = types.ModuleType('flask')
    fake_flask.g = fake_g
    fake_flask.request = fake_request
    monkeypatch.setitem(sys.modules, 'flask', fake_flask)

    app = _MockFlaskApp()
    obs.flask_middleware(app)
    return app, fake_g


# ------------------------------------------------------------------
# Flask tests
# ------------------------------------------------------------------

def test_flask_sends_create_and_complete_on_success(monkeypatch):
    obs = _make_visibe()
    app, _ = _flask_setup(monkeypatch, obs)

    app.trigger_before()
    app.trigger_after(_MockResponse(200))
    time.sleep(0.05)  # allow retry thread if any

    assert len(obs.api_client.create_calls) == 1
    assert len(obs.api_client.complete_calls) == 1
    assert obs.api_client.complete_calls[0][1]['status'] == 'completed'


def test_flask_failed_status_on_4xx(monkeypatch):
    obs = _make_visibe()
    app, _ = _flask_setup(monkeypatch, obs)

    app.trigger_before()
    app.trigger_after(_MockResponse(404))

    assert obs.api_client.complete_calls[0][1]['status'] == 'failed'


def test_flask_failed_status_on_5xx(monkeypatch):
    obs = _make_visibe()
    app, _ = _flask_setup(monkeypatch, obs)

    app.trigger_before()
    app.trigger_after(_MockResponse(500))

    assert obs.api_client.complete_calls[0][1]['status'] == 'failed'


def test_flask_failed_status_on_exception(monkeypatch):
    obs = _make_visibe()
    app, _ = _flask_setup(monkeypatch, obs)

    app.trigger_before()
    # after_request is not called when an unhandled exception occurs in Flask
    app.trigger_teardown(exc=RuntimeError("handler crashed"))

    assert len(obs.api_client.complete_calls) == 1
    assert obs.api_client.complete_calls[0][1]['status'] == 'failed'


def test_flask_double_call_guard(monkeypatch):
    """after_request + teardown_request must not send two complete_trace calls."""
    obs = _make_visibe()
    app, _ = _flask_setup(monkeypatch, obs)

    app.trigger_before()
    app.trigger_after(_MockResponse(200))   # sets _visibe_done=True
    app.trigger_teardown(exc=None)           # should be a no-op
    time.sleep(0.05)

    assert len(obs.api_client.complete_calls) == 1


def test_flask_no_tracing_when_disabled(monkeypatch):
    obs = _make_visibe()
    obs.api_client._enabled = False
    app, _ = _flask_setup(monkeypatch, obs)

    app.trigger_before()
    app.trigger_after(_MockResponse(200))

    assert len(obs.api_client.create_calls) == 0
    assert len(obs.api_client.complete_calls) == 0


def test_flask_complete_trace_payload_fields(monkeypatch):
    """Payload must include all required summary fields."""
    obs = _make_visibe()
    app, _ = _flask_setup(monkeypatch, obs)

    app.trigger_before()
    app.trigger_after(_MockResponse(200))

    payload = obs.api_client.complete_calls[0][1]
    for key in ('status', 'ended_at', 'duration_ms', 'llm_call_count',
                'total_cost', 'total_tokens', 'total_input_tokens',
                'total_output_tokens'):
        assert key in payload, f"missing key: {key}"


def test_flask_forwards_session_context_fields(monkeypatch):
    """first_framework/model/prompt from _SessionContext forwarded to payload."""
    obs = _make_visibe()
    app, fake_g = _flask_setup(monkeypatch, obs)

    app.trigger_before()

    # Simulate an integration setting session context fields
    sc = fake_g._visibe_session
    sc.first_framework = 'openai'
    sc.first_model = 'gpt-4o'
    sc.first_prompt = 'Hello!'

    app.trigger_after(_MockResponse(200))

    payload = obs.api_client.complete_calls[0][1]
    assert payload.get('framework') == 'openai'
    assert payload.get('model') == 'gpt-4o'
    assert payload.get('prompt') == 'Hello!'


def test_flask_omits_session_fields_when_no_llm_calls(monkeypatch):
    """When no LLM calls fire, first_* fields must NOT appear in payload."""
    obs = _make_visibe()
    app, _ = _flask_setup(monkeypatch, obs)

    app.trigger_before()
    app.trigger_after(_MockResponse(200))

    payload = obs.api_client.complete_calls[0][1]
    assert 'framework' not in payload
    assert 'model' not in payload
    assert 'prompt' not in payload


def test_flask_retry_on_complete_trace_failure(monkeypatch):
    """complete_trace failure triggers one background retry."""
    obs = _make_visibe()
    obs.api_client._complete_raise_n = 1  # fail first call, succeed second
    app, _ = _flask_setup(monkeypatch, obs)

    app.trigger_before()
    app.trigger_after(_MockResponse(200))
    time.sleep(1.5)  # retry fires after 1s

    assert len(obs.api_client.complete_calls) == 2


def test_flask_active_session_reset_after_request(monkeypatch):
    """_active_session ContextVar must be reset to None after the request."""
    obs = _make_visibe()
    app, fake_g = _flask_setup(monkeypatch, obs)

    app.trigger_before()
    active_during = _active_session.get()
    app.trigger_after(_MockResponse(200))
    active_after = _active_session.get()

    assert active_during is not None
    assert active_after is None


# ------------------------------------------------------------------
# Starlette mock infrastructure
# ------------------------------------------------------------------

def _starlette_setup(monkeypatch):
    """Install fake starlette module so BaseHTTPMiddleware is importable."""
    class _FakeBaseHTTPMiddleware:
        def __init__(self, app=None):
            self.app = app

    fake_base = types.ModuleType('starlette.middleware.base')
    fake_base.BaseHTTPMiddleware = _FakeBaseHTTPMiddleware

    fake_mid = types.ModuleType('starlette.middleware')
    fake_star = types.ModuleType('starlette')

    monkeypatch.setitem(sys.modules, 'starlette', fake_star)
    monkeypatch.setitem(sys.modules, 'starlette.middleware', fake_mid)
    monkeypatch.setitem(sys.modules, 'starlette.middleware.base', fake_base)


def _make_asgi_request(method='GET', path='/api/chat'):
    return types.SimpleNamespace(
        method=method,
        url=types.SimpleNamespace(path=path),
    )


def _make_call_next(status_code=200, raises=None):
    """Return an async call_next function."""
    async def call_next(request):
        if raises:
            raise raises
        return types.SimpleNamespace(status_code=status_code)
    return call_next


async def _run_dispatch(obs, request, call_next_fn, monkeypatch):
    """Get the middleware class and call dispatch()."""
    MiddlewareClass = obs.starlette_middleware()
    middleware = MiddlewareClass.__new__(MiddlewareClass)
    try:
        return await middleware.dispatch(request, call_next_fn)
    except Exception:
        pass  # let tests that expect exceptions not crash here


# ------------------------------------------------------------------
# Starlette tests
# ------------------------------------------------------------------

@pytest.mark.asyncio
async def test_starlette_sends_create_and_complete_on_success(monkeypatch):
    _starlette_setup(monkeypatch)
    obs = _make_visibe()

    await _run_dispatch(obs, _make_asgi_request(), _make_call_next(200), monkeypatch)
    await asyncio.sleep(0.2)  # wait for _finalize thread

    assert len(obs.api_client.create_calls) == 1
    assert len(obs.api_client.complete_calls) == 1
    assert obs.api_client.complete_calls[0][1]['status'] == 'completed'


@pytest.mark.asyncio
async def test_starlette_failed_status_on_4xx(monkeypatch):
    _starlette_setup(monkeypatch)
    obs = _make_visibe()

    await _run_dispatch(obs, _make_asgi_request(), _make_call_next(400), monkeypatch)
    await asyncio.sleep(0.2)

    assert obs.api_client.complete_calls[0][1]['status'] == 'failed'


@pytest.mark.asyncio
async def test_starlette_failed_status_on_exception(monkeypatch):
    _starlette_setup(monkeypatch)
    obs = _make_visibe()

    await _run_dispatch(
        obs, _make_asgi_request(),
        _make_call_next(raises=RuntimeError("handler crashed")),
        monkeypatch,
    )
    await asyncio.sleep(0.2)

    assert len(obs.api_client.complete_calls) == 1
    assert obs.api_client.complete_calls[0][1]['status'] == 'failed'


@pytest.mark.asyncio
async def test_starlette_no_tracing_when_disabled(monkeypatch):
    _starlette_setup(monkeypatch)
    obs = _make_visibe()
    obs.api_client._enabled = False

    await _run_dispatch(obs, _make_asgi_request(), _make_call_next(200), monkeypatch)
    await asyncio.sleep(0.1)

    assert len(obs.api_client.create_calls) == 0
    assert len(obs.api_client.complete_calls) == 0


@pytest.mark.asyncio
async def test_starlette_complete_trace_after_create_trace(monkeypatch):
    """
    The ordering fix: complete_trace must always be called AFTER create_trace
    returns, even when create_trace is slow.  This verifies the threading.Event
    synchronisation added to _finalize().
    """
    _starlette_setup(monkeypatch)
    obs = _make_visibe()

    call_order = []
    create_done_event = threading.Event()

    original_create = obs.api_client.create_trace
    def slow_create(data):
        time.sleep(0.08)          # simulate a slow backend POST
        call_order.append('create')
        original_create(data)
    obs.api_client.create_trace = slow_create

    original_complete = obs.api_client.complete_trace
    def recording_complete(trace_id, data):
        call_order.append('complete')
        original_complete(trace_id, data)
        create_done_event.set()
    obs.api_client.complete_trace = recording_complete

    await _run_dispatch(obs, _make_asgi_request(), _make_call_next(200), monkeypatch)
    create_done_event.wait(timeout=3)

    assert call_order == ['create', 'complete'], (
        f"Ordering violated: {call_order}. "
        "complete_trace must not arrive before create_trace."
    )


@pytest.mark.asyncio
async def test_starlette_complete_trace_even_if_create_fails(monkeypatch):
    """
    If create_trace raises (auth error etc.), _finalize must still attempt
    complete_trace — matches Node.js behaviour of swallowing the create error.
    """
    _starlette_setup(monkeypatch)
    obs = _make_visibe()

    def failing_create(data):
        raise RuntimeError("auth error")
    obs.api_client.create_trace = failing_create

    await _run_dispatch(obs, _make_asgi_request(), _make_call_next(200), monkeypatch)
    await asyncio.sleep(0.3)

    # complete_trace is attempted even though create failed
    assert len(obs.api_client.complete_calls) == 1


@pytest.mark.asyncio
async def test_starlette_forwards_session_context_fields(monkeypatch):
    """first_framework/model/prompt forwarded when integrations fire."""
    _starlette_setup(monkeypatch)
    obs = _make_visibe()

    async def call_next_with_llm(request):
        # Simulate an integration updating the active session
        sc = _active_session.get()
        if sc:
            sc.first_framework = 'anthropic'
            sc.first_model = 'claude-3-5-sonnet-20241022'
            sc.first_prompt = 'What is the weather?'
        return types.SimpleNamespace(status_code=200)

    await _run_dispatch(obs, _make_asgi_request(), call_next_with_llm, monkeypatch)
    await asyncio.sleep(0.2)

    payload = obs.api_client.complete_calls[0][1]
    assert payload.get('framework') == 'anthropic'
    assert payload.get('model') == 'claude-3-5-sonnet-20241022'
    assert payload.get('prompt') == 'What is the weather?'


@pytest.mark.asyncio
async def test_starlette_omits_session_fields_when_no_llm_calls(monkeypatch):
    _starlette_setup(monkeypatch)
    obs = _make_visibe()

    await _run_dispatch(obs, _make_asgi_request(), _make_call_next(200), monkeypatch)
    await asyncio.sleep(0.2)

    payload = obs.api_client.complete_calls[0][1]
    assert 'framework' not in payload
    assert 'model' not in payload
    assert 'prompt' not in payload


@pytest.mark.asyncio
async def test_starlette_retry_on_complete_trace_failure(monkeypatch):
    """complete_trace failure triggers one retry after 1s."""
    _starlette_setup(monkeypatch)
    obs = _make_visibe()
    obs.api_client._complete_raise_n = 1  # fail first, succeed second

    await _run_dispatch(obs, _make_asgi_request(), _make_call_next(200), monkeypatch)
    await asyncio.sleep(1.5)  # retry fires after 1s

    assert len(obs.api_client.complete_calls) == 2


@pytest.mark.asyncio
async def test_starlette_active_session_reset_after_request(monkeypatch):
    """_active_session ContextVar must be None after dispatch returns."""
    _starlette_setup(monkeypatch)
    obs = _make_visibe()

    active_during_request = []

    async def call_next_capture(request):
        active_during_request.append(_active_session.get())
        return types.SimpleNamespace(status_code=200)

    await _run_dispatch(obs, _make_asgi_request(), call_next_capture, monkeypatch)

    assert active_during_request[0] is not None   # set during handler
    assert _active_session.get() is None           # reset after dispatch


@pytest.mark.asyncio
async def test_starlette_complete_trace_payload_fields(monkeypatch):
    _starlette_setup(monkeypatch)
    obs = _make_visibe()

    await _run_dispatch(obs, _make_asgi_request(), _make_call_next(200), monkeypatch)
    await asyncio.sleep(0.2)

    payload = obs.api_client.complete_calls[0][1]
    for key in ('status', 'ended_at', 'duration_ms', 'llm_call_count',
                'total_cost', 'total_tokens', 'total_input_tokens',
                'total_output_tokens'):
        assert key in payload, f"missing key: {key}"
