"""
Tests for GateWay Job State Machine Module.
"""

import time

import pytest
import pytest_asyncio

from verlex_server.gateway.core.state_machine import (
    ACTIVE_STATES,
    TERMINAL_STATES,
    VALID_TRANSITIONS,
    JobContext,
    JobState,
    JobStateError,
    JobStateMachine,
    JobStateManager,
    MemoryStatePersistence,
    StateTransition,
    StateTransitionError,
)


class TestJobState:
    """Tests for JobState enum."""

    def test_all_states_have_transitions(self):
        """All states should be in the transitions dict."""
        for state in JobState:
            assert state in VALID_TRANSITIONS

    def test_terminal_states_have_no_transitions(self):
        """Terminal states should have no valid transitions."""
        for state in TERMINAL_STATES:
            assert len(VALID_TRANSITIONS[state]) == 0

    def test_active_states_defined(self):
        """Active states should be properly defined."""
        assert JobState.RUNNING in ACTIVE_STATES
        assert JobState.QUEUED in ACTIVE_STATES
        assert JobState.COMPLETED not in ACTIVE_STATES
        assert JobState.FAILED not in ACTIVE_STATES


class TestStateTransition:
    """Tests for StateTransition dataclass."""

    def test_transition_creation(self):
        """Test creating a state transition."""
        transition = StateTransition(
            from_state=JobState.PENDING,
            to_state=JobState.QUEUED,
            reason="Job submitted",
        )

        assert transition.from_state == JobState.PENDING
        assert transition.to_state == JobState.QUEUED
        assert transition.reason == "Job submitted"
        assert transition.timestamp > 0

    def test_transition_to_dict(self):
        """Test transition serialization."""
        transition = StateTransition(
            from_state=JobState.RUNNING,
            to_state=JobState.COMPLETED,
        )

        data = transition.to_dict()
        assert data["from_state"] == "running"
        assert data["to_state"] == "completed"

    def test_transition_from_dict(self):
        """Test transition deserialization."""
        data = {
            "from_state": "pending",
            "to_state": "queued",
            "timestamp": 12345.0,
            "reason": "test",
            "metadata": {"key": "value"},
        }

        transition = StateTransition.from_dict(data)
        assert transition.from_state == JobState.PENDING
        assert transition.to_state == JobState.QUEUED
        assert transition.reason == "test"


class TestJobContext:
    """Tests for JobContext dataclass."""

    def test_context_creation(self):
        """Test creating a job context."""
        context = JobContext(
            job_id="job_123",
            user_id="user_456",
        )

        assert context.job_id == "job_123"
        assert context.user_id == "user_456"
        assert context.state == JobState.PENDING
        assert context.created_at > 0

    def test_context_to_dict(self):
        """Test context serialization."""
        context = JobContext(
            job_id="job_1",
            user_id="user_1",
            state=JobState.RUNNING,
        )

        data = context.to_dict()
        assert data["job_id"] == "job_1"
        assert data["state"] == "running"

    def test_context_from_dict(self):
        """Test context deserialization."""
        data = {
            "job_id": "job_2",
            "user_id": "user_2",
            "state": "completed",
            "checkpoint_count": 5,
        }

        context = JobContext.from_dict(data)
        assert context.job_id == "job_2"
        assert context.state == JobState.COMPLETED
        assert context.checkpoint_count == 5


class TestJobStateMachine:
    """Tests for JobStateMachine."""

    @pytest_asyncio.fixture
    async def machine(self):
        return JobStateMachine(
            job_id="test_job",
            user_id="test_user",
        )

    @pytest.mark.asyncio
    async def test_initial_state(self, machine):
        """Test initial state is PENDING."""
        assert machine.state == JobState.PENDING
        assert not machine.is_active
        assert not machine.is_terminal

    @pytest.mark.asyncio
    async def test_valid_transition(self, machine):
        """Test valid state transition."""
        transition = await machine.transition(
            JobState.QUEUED,
            reason="Job submitted",
        )

        assert machine.state == JobState.QUEUED
        assert transition.from_state == JobState.PENDING
        assert transition.to_state == JobState.QUEUED

    @pytest.mark.asyncio
    async def test_invalid_transition(self, machine):
        """Test invalid state transition raises error."""
        with pytest.raises(StateTransitionError):
            await machine.transition(JobState.COMPLETED)

    @pytest.mark.asyncio
    async def test_can_transition_to(self, machine):
        """Test transition validation."""
        assert machine.can_transition_to(JobState.QUEUED)
        assert machine.can_transition_to(JobState.CANCELLED)
        assert not machine.can_transition_to(JobState.COMPLETED)
        assert not machine.can_transition_to(JobState.RUNNING)

    @pytest.mark.asyncio
    async def test_get_valid_transitions(self, machine):
        """Test getting valid transitions."""
        valid = machine.get_valid_transitions()

        assert JobState.QUEUED in valid
        assert JobState.CANCELLED in valid
        assert JobState.COMPLETED not in valid

    @pytest.mark.asyncio
    async def test_full_lifecycle(self, machine):
        """Test complete job lifecycle."""
        # Submit
        await machine.transition(JobState.QUEUED)
        assert machine.is_active

        # Schedule
        await machine.transition(JobState.SCHEDULED)

        # Provision
        await machine.transition(JobState.PROVISIONING)

        # Deploy
        await machine.transition(JobState.DEPLOYING)

        # Run
        await machine.transition(JobState.RUNNING)
        assert machine.is_running

        # Complete
        await machine.transition(JobState.COMPLETING)
        await machine.transition(JobState.COMPLETED)

        assert machine.state == JobState.COMPLETED
        assert machine.is_terminal
        assert not machine.is_active

    @pytest.mark.asyncio
    async def test_history_tracking(self, machine):
        """Test state history is tracked."""
        await machine.transition(JobState.QUEUED)
        await machine.transition(JobState.SCHEDULED)
        await machine.transition(JobState.PROVISIONING)

        history = machine.get_history()
        assert len(history) == 3
        assert history[0].to_state == JobState.QUEUED
        assert history[1].to_state == JobState.SCHEDULED
        assert history[2].to_state == JobState.PROVISIONING

    @pytest.mark.asyncio
    async def test_timing_updates(self, machine):
        """Test timing fields are updated."""
        assert machine.context.started_at is None

        await machine.transition(JobState.QUEUED)
        await machine.transition(JobState.SCHEDULED)
        await machine.transition(JobState.PROVISIONING)
        await machine.transition(JobState.DEPLOYING)
        await machine.transition(JobState.RUNNING)

        assert machine.context.started_at is not None
        assert machine.context.completed_at is None

        await machine.transition(JobState.COMPLETING)
        await machine.transition(JobState.COMPLETED)

        assert machine.context.completed_at is not None

    @pytest.mark.asyncio
    async def test_get_duration(self, machine):
        """Test duration calculation."""
        assert machine.get_duration() is None

        await machine.transition(JobState.QUEUED)
        await machine.transition(JobState.SCHEDULED)
        await machine.transition(JobState.PROVISIONING)
        await machine.transition(JobState.DEPLOYING)
        await machine.transition(JobState.RUNNING)

        time.sleep(0.1)

        duration = machine.get_duration()
        assert duration is not None
        assert duration >= 0.1

    @pytest.mark.asyncio
    async def test_set_error(self, machine):
        """Test setting error."""
        await machine.transition(JobState.QUEUED)
        await machine.transition(JobState.SCHEDULED)
        await machine.transition(JobState.PROVISIONING)
        await machine.transition(JobState.DEPLOYING)
        await machine.transition(JobState.RUNNING)

        await machine.set_error("Something went wrong")

        assert machine.context.error == "Something went wrong"
        assert machine.context.error_count == 1
        assert machine.state == JobState.FAILED

    @pytest.mark.asyncio
    async def test_checkpoint(self, machine):
        """Test checkpoint recording."""
        await machine.checkpoint("checkpoint_123")

        assert machine.context.checkpoint_id == "checkpoint_123"
        assert machine.context.checkpoint_count == 1

        await machine.checkpoint("checkpoint_456")
        assert machine.context.checkpoint_count == 2

    @pytest.mark.asyncio
    async def test_set_instance(self, machine):
        """Test instance information."""
        await machine.set_instance(
            instance_id="i-12345",
            provider="aws",
            region="us-east-1",
        )

        assert machine.context.instance_id == "i-12345"
        assert machine.context.provider == "aws"
        assert machine.context.region == "us-east-1"

    @pytest.mark.asyncio
    async def test_disconnect_reconnect(self, machine):
        """Test disconnect and reconnect handling."""
        await machine.transition(JobState.QUEUED)
        await machine.transition(JobState.SCHEDULED)
        await machine.transition(JobState.PROVISIONING)
        await machine.transition(JobState.DEPLOYING)
        await machine.transition(JobState.RUNNING)

        # Disconnect
        await machine.handle_disconnect()
        assert machine.state == JobState.DISCONNECTED
        assert not machine.context.client_connected

        # Reconnect
        await machine.handle_reconnect()
        assert machine.state == JobState.RUNNING
        assert machine.context.client_connected

    @pytest.mark.asyncio
    async def test_callbacks(self, machine):
        """Test state change callbacks."""
        events = []

        def callback(job_id, from_state, to_state, metadata):
            events.append((job_id, from_state, to_state))

        machine.register_callback(callback)

        await machine.transition(JobState.QUEUED)
        await machine.transition(JobState.CANCELLED)

        assert len(events) == 2
        assert events[0] == ("test_job", JobState.PENDING, JobState.QUEUED)
        assert events[1] == ("test_job", JobState.QUEUED, JobState.CANCELLED)

    @pytest.mark.asyncio
    async def test_callbacks_with_filter(self, machine):
        """Test callbacks with state filter."""
        events = []

        def callback(job_id, from_state, to_state, metadata):
            events.append(to_state)

        # Only trigger on terminal states
        machine.register_callback(callback, states=TERMINAL_STATES)

        await machine.transition(JobState.QUEUED)
        assert len(events) == 0

        await machine.transition(JobState.CANCELLED)
        assert len(events) == 1
        assert events[0] == JobState.CANCELLED


class TestMemoryStatePersistence:
    """Tests for MemoryStatePersistence."""

    @pytest.fixture
    def persistence(self):
        return MemoryStatePersistence()

    @pytest.mark.asyncio
    async def test_save_load_context(self, persistence):
        """Test saving and loading context."""
        context = JobContext(
            job_id="job_1",
            user_id="user_1",
            state=JobState.RUNNING,
        )

        await persistence.save_context(context)
        loaded = await persistence.load_context("job_1")

        assert loaded is not None
        assert loaded.job_id == "job_1"
        assert loaded.state == JobState.RUNNING

    @pytest.mark.asyncio
    async def test_save_load_transition(self, persistence):
        """Test saving and loading transitions."""
        transition = StateTransition(
            from_state=JobState.PENDING,
            to_state=JobState.QUEUED,
        )

        await persistence.save_transition("job_1", transition)
        history = await persistence.load_history("job_1")

        assert len(history) == 1
        assert history[0].from_state == JobState.PENDING

    @pytest.mark.asyncio
    async def test_delete(self, persistence):
        """Test deleting state data."""
        context = JobContext(job_id="job_2", user_id="user_2")
        await persistence.save_context(context)

        await persistence.delete("job_2")

        loaded = await persistence.load_context("job_2")
        assert loaded is None


class TestJobStateManager:
    """Tests for JobStateManager."""

    @pytest_asyncio.fixture
    async def manager(self):
        return JobStateManager()

    @pytest.mark.asyncio
    async def test_create_job(self, manager):
        """Test creating a new job."""
        machine = await manager.create_job("job_1", "user_1")

        assert machine is not None
        assert machine.job_id == "job_1"
        assert machine.state == JobState.PENDING

    @pytest.mark.asyncio
    async def test_create_duplicate_job(self, manager):
        """Test creating duplicate job raises error."""
        await manager.create_job("job_2", "user_2")

        with pytest.raises(JobStateError):
            await manager.create_job("job_2", "user_2")

    @pytest.mark.asyncio
    async def test_get_job(self, manager):
        """Test getting a job."""
        await manager.create_job("job_3", "user_3")

        machine = await manager.get_job("job_3")
        assert machine is not None
        assert machine.job_id == "job_3"

    @pytest.mark.asyncio
    async def test_get_nonexistent_job(self, manager):
        """Test getting nonexistent job returns None."""
        machine = await manager.get_job("nonexistent")
        assert machine is None

    @pytest.mark.asyncio
    async def test_get_jobs_by_state(self, manager):
        """Test filtering jobs by state."""
        job1 = await manager.create_job("job_a", "user_1")
        job2 = await manager.create_job("job_b", "user_1")
        await manager.create_job("job_c", "user_1")

        await job1.transition(JobState.QUEUED)
        await job2.transition(JobState.QUEUED)

        queued = await manager.get_jobs_by_state(JobState.QUEUED)
        assert len(queued) == 2

        pending = await manager.get_jobs_by_state(JobState.PENDING)
        assert len(pending) == 1

    @pytest.mark.asyncio
    async def test_get_jobs_by_user(self, manager):
        """Test filtering jobs by user."""
        await manager.create_job("job_1", "alice")
        await manager.create_job("job_2", "alice")
        await manager.create_job("job_3", "bob")

        alice_jobs = await manager.get_jobs_by_user("alice")
        assert len(alice_jobs) == 2

        bob_jobs = await manager.get_jobs_by_user("bob")
        assert len(bob_jobs) == 1

    @pytest.mark.asyncio
    async def test_get_active_jobs(self, manager):
        """Test getting active jobs."""
        job1 = await manager.create_job("job_1", "user_1")
        await manager.create_job("job_2", "user_1")

        await job1.transition(JobState.QUEUED)

        active = await manager.get_active_jobs()
        assert len(active) == 1
        assert active[0].job_id == "job_1"

    @pytest.mark.asyncio
    async def test_global_callback(self, manager):
        """Test global callbacks."""
        events = []

        def callback(job_id, from_state, to_state, metadata):
            events.append(job_id)

        manager.register_global_callback(callback)

        job1 = await manager.create_job("job_1", "user_1")
        await job1.transition(JobState.QUEUED)

        job2 = await manager.create_job("job_2", "user_2")
        await job2.transition(JobState.QUEUED)

        assert "job_1" in events
        assert "job_2" in events

    @pytest.mark.asyncio
    async def test_get_stats(self, manager):
        """Test getting statistics."""
        job1 = await manager.create_job("job_1", "user_1")
        await manager.create_job("job_2", "user_1")

        await job1.transition(JobState.QUEUED)

        stats = manager.get_stats()

        assert stats["total_jobs"] == 2
        assert stats["active_jobs"] == 1
        assert stats["by_state"]["pending"] == 1
        assert stats["by_state"]["queued"] == 1
