import time

import pytest

from higgsmeter import AppState
from higgsmeter import enable_tracker
from higgsmeter import state_tracker
from higgsmeter import track_state


def test_basic_state_changes():
    """Test basic state transitions"""
    enable_tracker(rank=0)

    state_tracker.set_state(AppState.IDLE)
    assert state_tracker._tracker._current_state == AppState.IDLE

    state_tracker.set_state(AppState.CPU)
    assert state_tracker._tracker._current_state == AppState.CPU


def test_context_manager_scope():
    """Test context manager automatically reverts state"""
    enable_tracker(rank=0)

    state_tracker.set_state(AppState.IDLE)
    assert state_tracker._tracker._current_state == AppState.IDLE

    with state_tracker.scope(AppState.IO):
        assert state_tracker._tracker._current_state == AppState.IO

    assert state_tracker._tracker._current_state == AppState.IDLE


def test_nested_scopes():
    """Test nested context managers work correctly"""
    enable_tracker(rank=0)

    state_tracker.set_state(AppState.IDLE)

    with state_tracker.scope(AppState.CPU):
        assert state_tracker._tracker._current_state == AppState.CPU

        with state_tracker.scope(AppState.IO):
            assert state_tracker._tracker._current_state == AppState.IO

        assert state_tracker._tracker._current_state == AppState.CPU

    assert state_tracker._tracker._current_state == AppState.IDLE


def test_decorator_with_manual_scope():
    """Test decorator works with manual scopes inside"""
    enable_tracker(rank=0)

    @track_state(AppState.CPU)
    def outer_function():
        assert state_tracker._tracker._current_state == AppState.CPU

        with state_tracker.scope(AppState.IO):
            assert state_tracker._tracker._current_state == AppState.IO

        assert state_tracker._tracker._current_state == AppState.CPU

    state_tracker.set_state(AppState.IDLE)
    outer_function()
    assert state_tracker._tracker._current_state == AppState.IDLE


def test_decorator_calling_decorated_method():
    """Test nested decorators work correctly"""
    enable_tracker(rank=0)

    @track_state(AppState.IO)
    def io_function():
        assert state_tracker._tracker._current_state == AppState.IO

    @track_state(AppState.CPU)
    def cpu_function():
        assert state_tracker._tracker._current_state == AppState.CPU
        io_function()
        assert state_tracker._tracker._current_state == AppState.CPU

    state_tracker.set_state(AppState.IDLE)
    cpu_function()
    assert state_tracker._tracker._current_state == AppState.IDLE


def test_operations_tracking():
    """Test operation labels are tracked correctly"""
    enable_tracker(rank=0)

    state_tracker.set_state(AppState.IDLE, operation='waiting')
    assert state_tracker._tracker._current_operation == 'waiting'

    with state_tracker.scope(AppState.IO, 'fetch_data'):
        assert state_tracker._tracker._current_operation == 'fetch_data'

    assert state_tracker._tracker._current_operation == 'waiting'


def test_rank_zero_only_tracking():
    """Test that only rank 0 enables tracking"""
    enable_tracker(rank=1)
    assert not state_tracker._tracker._enabled

    enable_tracker(rank=0)
    assert state_tracker._tracker._enabled


@pytest.mark.asyncio
async def test_async_decorator():
    """Test decorator works with async functions"""
    enable_tracker(rank=0)

    @track_state(AppState.GPU)
    async def async_gpu_work():
        assert state_tracker._tracker._current_state == AppState.GPU
        await asyncio.sleep(0.01)

    import asyncio

    state_tracker.set_state(AppState.IDLE)
    await async_gpu_work()
    assert state_tracker._tracker._current_state == AppState.IDLE


def test_metrics_collection():
    """Test that Prometheus metrics are recorded"""
    from prometheus_client import REGISTRY
    from prometheus_client import generate_latest

    enable_tracker(rank=0)

    state_tracker.set_state(AppState.IDLE)
    time.sleep(0.1)
    state_tracker.set_state(AppState.GPU, operation='test')
    time.sleep(0.1)
    state_tracker.finalize()

    metrics = generate_latest(REGISTRY).decode('utf-8')

    assert 'worker_activity_time_seconds_total' in metrics
    assert 'worker_current_state' in metrics
    assert 'state="idle"' in metrics
    assert 'state="gpu"' in metrics
