"""Tests for the v3 pipeline orchestrator — state machine, queues, leases, flags."""

from __future__ import annotations

from datetime import timedelta
from uuid import uuid4

import pytest

from services.intelligence_pipeline_v3.orchestrator.feature_flags import (
    FeatureFlags,
    PipelineVersion,
)
from services.intelligence_pipeline_v3.orchestrator.leases import (
    LeaseExpiredError,
    LeaseManager,
)
from services.intelligence_pipeline_v3.orchestrator.queues import (
    QueueMessage,
    QueueName,
    QueueRouter,
)
from services.intelligence_pipeline_v3.orchestrator.state import (
    PipelineState,
    PipelineStateMachine,
    StageState,
)

# ---------------------------------------------------------------------------
# State Machine Tests
# ---------------------------------------------------------------------------


class TestPipelineStateMachine:
    """Task 41.1: Explicit stage state transitions and idempotency keys."""

    def test_initial_state_is_pending(self):
        sm = PipelineStateMachine(document_id="doc-001")
        assert sm.state == PipelineState.PENDING

    def test_valid_transition_pending_to_segmenting(self):
        sm = PipelineStateMachine(document_id="doc-001")
        t = sm.transition_pipeline(PipelineState.SEGMENTING, "start processing")
        assert sm.state == PipelineState.SEGMENTING
        assert t.from_state == PipelineState.PENDING
        assert t.to_state == PipelineState.SEGMENTING
        assert t.idempotency_key != ""

    def test_invalid_transition_raises(self):
        sm = PipelineStateMachine(document_id="doc-001")
        with pytest.raises(ValueError, match="Invalid pipeline transition"):
            sm.transition_pipeline(PipelineState.COMPLETED)

    def test_full_happy_path_transitions(self):
        sm = PipelineStateMachine(document_id="doc-001")
        states = [
            PipelineState.SEGMENTING,
            PipelineState.EXTRACTING,
            PipelineState.RESOLVING,
            PipelineState.VERIFYING,
            PipelineState.ROUTING,
            PipelineState.IMPACT,
            PipelineState.PERSISTING,
            PipelineState.COMPLETED,
        ]
        for state in states:
            sm.transition_pipeline(state)
        assert sm.state == PipelineState.COMPLETED
        assert len(sm.history) == len(states)

    def test_routing_can_go_to_adjudication(self):
        sm = PipelineStateMachine(document_id="doc-001")
        for s in [
            PipelineState.SEGMENTING,
            PipelineState.EXTRACTING,
            PipelineState.RESOLVING,
            PipelineState.VERIFYING,
            PipelineState.ROUTING,
        ]:
            sm.transition_pipeline(s)
        sm.transition_pipeline(PipelineState.ADJUDICATING)
        assert sm.state == PipelineState.ADJUDICATING

    def test_stage_state_transitions(self):
        sm = PipelineStateMachine(document_id="doc-001")
        sm.transition_stage("extraction", StageState.LEASED)
        assert sm.stage_states["extraction"] == StageState.LEASED
        sm.transition_stage("extraction", StageState.RUNNING)
        assert sm.stage_states["extraction"] == StageState.RUNNING
        sm.transition_stage("extraction", StageState.SUCCEEDED)
        assert sm.stage_states["extraction"] == StageState.SUCCEEDED

    def test_stage_invalid_transition_raises(self):
        sm = PipelineStateMachine(document_id="doc-001")
        sm.transition_stage("extraction", StageState.LEASED)
        with pytest.raises(ValueError, match="Invalid stage transition"):
            sm.transition_stage("extraction", StageState.SUCCEEDED)

    def test_idempotency_key_is_deterministic(self):
        sm1 = PipelineStateMachine(document_id="doc-001")
        sm2 = PipelineStateMachine(document_id="doc-001")
        t1 = sm1.transition_pipeline(PipelineState.SEGMENTING)
        t2 = sm2.transition_pipeline(PipelineState.SEGMENTING)
        assert t1.idempotency_key == t2.idempotency_key

    def test_can_retry_tracks_attempts(self):
        sm = PipelineStateMachine(document_id="doc-001", max_retries=2)
        sm.transition_stage("extraction", StageState.LEASED)
        sm.transition_stage("extraction", StageState.RUNNING)
        sm.transition_stage("extraction", StageState.RETRYING)
        assert sm.can_retry("extraction")
        sm.transition_stage("extraction", StageState.QUEUED)
        sm.transition_stage("extraction", StageState.LEASED)
        sm.transition_stage("extraction", StageState.RUNNING)
        sm.transition_stage("extraction", StageState.RETRYING)
        assert not sm.can_retry("extraction")

    def test_dead_letter_after_max_retries(self):
        sm = PipelineStateMachine(document_id="doc-001", max_retries=1)
        sm.transition_stage("extraction", StageState.LEASED)
        sm.transition_stage("extraction", StageState.RUNNING)
        sm.transition_stage("extraction", StageState.RETRYING)
        sm.transition_pipeline(PipelineState.SEGMENTING)
        sm.transition_pipeline(PipelineState.FAILED)
        assert sm.should_dead_letter()


# ---------------------------------------------------------------------------
# Queue Tests
# ---------------------------------------------------------------------------


class TestQueueRouter:
    """Task 41.2: Fast-path, adjudication, persistence, and review queues."""

    def test_all_queue_names_defined(self):
        assert QueueName.INCOMING
        assert QueueName.FAST_PATH
        assert QueueName.ADJUDICATION
        assert QueueName.PERSISTENCE
        assert QueueName.REVIEW
        assert QueueName.DEAD_LETTER

    def test_enqueue_and_dequeue(self):
        router = QueueRouter()
        msg = QueueMessage.create(
            queue=QueueName.FAST_PATH,
            run_id=uuid4(),
            document_id="doc-001",
        )
        assert router.enqueue(msg)
        assert router.depth(QueueName.FAST_PATH) == 1
        dequeued = router.dequeue(QueueName.FAST_PATH)
        assert dequeued is not None
        assert dequeued.document_id == "doc-001"

    def test_backpressure_rejects_at_max_depth(self):
        router = QueueRouter(max_depth=2)
        run_id = uuid4()
        for i in range(2):
            msg = QueueMessage.create(
                queue=QueueName.FAST_PATH, run_id=run_id, document_id=f"doc-{i}"
            )
            assert router.enqueue(msg)
        # Third should be rejected
        msg = QueueMessage.create(
            queue=QueueName.FAST_PATH, run_id=run_id, document_id="doc-3"
        )
        assert not router.enqueue(msg)

    def test_idempotency_rejects_duplicate_keys(self):
        router = QueueRouter()
        run_id = uuid4()
        msg = QueueMessage.create(
            queue=QueueName.FAST_PATH,
            run_id=run_id,
            document_id="doc-001",
            idempotency_key="key-123",
        )
        assert router.enqueue(msg)
        router.dequeue(QueueName.FAST_PATH)
        # Second enqueue with same key should be rejected
        msg2 = QueueMessage.create(
            queue=QueueName.FAST_PATH,
            run_id=run_id,
            document_id="doc-001",
            idempotency_key="key-123",
        )
        assert not router.enqueue(msg2)

    def test_move_to_dead_letter(self):
        router = QueueRouter()
        msg = QueueMessage.create(
            queue=QueueName.FAST_PATH, run_id=uuid4(), document_id="doc-001"
        )
        router.enqueue(msg)
        original = router.dequeue(QueueName.FAST_PATH)
        assert original is not None
        dlq_msg = router.move_to_dead_letter(original)
        assert dlq_msg.queue == QueueName.DEAD_LETTER
        assert router.depth(QueueName.DEAD_LETTER) == 1

    def test_dequeue_empty_returns_none(self):
        router = QueueRouter()
        assert router.dequeue(QueueName.REVIEW) is None

    def test_is_saturated(self):
        router = QueueRouter(max_depth=5)
        run_id = uuid4()
        for i in range(5):
            msg = QueueMessage.create(
                queue=QueueName.ADJUDICATION, run_id=run_id, document_id=f"doc-{i}"
            )
            router.enqueue(msg)
        assert router.is_saturated(QueueName.ADJUDICATION)


# ---------------------------------------------------------------------------
# Lease Tests
# ---------------------------------------------------------------------------


class TestLeaseManager:
    """Task 41.3: Leases, retry policies, dead-letter handling."""

    def test_acquire_lease(self):
        mgr = LeaseManager()
        run_id = uuid4()
        lease = mgr.acquire(run_id, "extraction", "worker-1")
        assert lease is not None
        assert lease.is_active
        assert not lease.is_expired

    def test_cannot_double_acquire(self):
        mgr = LeaseManager()
        run_id = uuid4()
        lease1 = mgr.acquire(run_id, "extraction", "worker-1")
        lease2 = mgr.acquire(run_id, "extraction", "worker-2")
        assert lease1 is not None
        assert lease2 is None

    def test_release_allows_reacquisition(self):
        mgr = LeaseManager()
        run_id = uuid4()
        lease = mgr.acquire(run_id, "extraction", "worker-1")
        assert lease is not None
        mgr.release(lease)
        lease2 = mgr.acquire(run_id, "extraction", "worker-2")
        assert lease2 is not None

    def test_expired_lease_allows_reacquisition(self):
        mgr = LeaseManager(default_ttl=timedelta(seconds=-1))
        run_id = uuid4()
        lease = mgr.acquire(run_id, "extraction", "worker-1")
        assert lease is not None
        assert lease.is_expired
        # Another worker can acquire
        lease2 = mgr.acquire(run_id, "extraction", "worker-2")
        assert lease2 is not None

    def test_renew_extends_lease(self):
        mgr = LeaseManager(default_ttl=timedelta(seconds=60))
        run_id = uuid4()
        lease = mgr.acquire(run_id, "extraction", "worker-1")
        assert lease is not None
        original_expiry = lease.expires_at
        mgr.renew(lease, timedelta(seconds=120))
        assert lease.expires_at > original_expiry
        assert lease.renewed_count == 1

    def test_renew_expired_raises(self):
        mgr = LeaseManager(default_ttl=timedelta(seconds=-1))
        run_id = uuid4()
        lease = mgr.acquire(run_id, "extraction", "worker-1")
        assert lease is not None
        with pytest.raises(LeaseExpiredError):
            mgr.renew(lease)

    def test_active_count(self):
        mgr = LeaseManager()
        run_id = uuid4()
        mgr.acquire(run_id, "extraction", "worker-1")
        mgr.acquire(run_id, "sentiment", "worker-2")
        assert mgr.active_count() == 2


# ---------------------------------------------------------------------------
# Feature Flag Tests
# ---------------------------------------------------------------------------


class TestFeatureFlags:
    """Task 41.4: Independent v2/v3 routing behind feature flags."""

    def test_default_routes_to_v2(self):
        flags = FeatureFlags()
        assert flags.resolve("doc-001") == PipelineVersion.V2

    def test_v3_enabled_routes_to_v3(self):
        flags = FeatureFlags(v3_enabled=True, default_version=PipelineVersion.V3)
        assert flags.resolve("doc-001") == PipelineVersion.V3

    def test_shadow_mode_returns_shadow(self):
        flags = FeatureFlags(shadow_enabled=True)
        assert flags.resolve("doc-001") == PipelineVersion.SHADOW

    def test_percentage_routing_is_deterministic(self):
        flags = FeatureFlags(v3_enabled=True, v3_percentage=50)
        result1 = flags.resolve("doc-001")
        result2 = flags.resolve("doc-001")
        assert result1 == result2

    def test_agent_override_takes_precedence(self):
        flags = FeatureFlags(v3_enabled=True, v3_percentage=0)
        flags.set_agent_override("agent-1", PipelineVersion.V3)
        assert (
            flags.resolve("doc-001", agent_id="agent-1") == PipelineVersion.V3
        )
        # Different agent uses default
        result = flags.resolve("doc-001", agent_id="agent-2")
        # Not v3 since percentage is 0 and no override for agent-2
        assert result in (PipelineVersion.V2, PipelineVersion.V3)

    def test_document_type_override(self):
        flags = FeatureFlags(v3_enabled=True)
        flags.document_type_overrides["filing"] = PipelineVersion.V3
        assert (
            flags.resolve("doc-001", document_type="filing")
            == PipelineVersion.V3
        )

    def test_excluded_document_type(self):
        flags = FeatureFlags(v3_enabled=True, v3_percentage=100)
        flags.excluded_document_types.add("transcript")
        assert (
            flags.resolve("doc-001", document_type="transcript")
            == PipelineVersion.V2
        )

    def test_is_v3_active(self):
        flags = FeatureFlags()
        assert not flags.is_v3_active()
        flags.v3_enabled = True
        assert flags.is_v3_active()

    def test_to_dict_serialization(self):
        flags = FeatureFlags(v3_enabled=True, v3_percentage=25)
        d = flags.to_dict()
        assert d["v3_enabled"] is True
        assert d["v3_percentage"] == 25
