"""Contract and load tests for the specialist inference service.

Tests entity extraction, classification, relation extraction, structured
extraction, health/ready endpoints, batch size enforcement, dynamic batching,
bounded queue rejection, and model version.
"""

from __future__ import annotations

import asyncio
import os

import pytest
from fastapi.testclient import TestClient

# Force test mode before importing the app
os.environ["SPECIALIST_TEST_MODE"] = "1"

from services.specialist.app import app  # noqa: E402
from services.specialist.batching import DynamicBatcher, QueueFullError  # noqa: E402


@pytest.fixture
def client():
    """Create a test client with the specialist app."""
    with TestClient(app) as c:
        yield c


# ---------------------------------------------------------------------------
# Health / Ready endpoints
# ---------------------------------------------------------------------------


class TestHealthEndpoints:
    """Test health and readiness probes."""

    def test_health_returns_ok(self, client: TestClient):
        resp = client.get("/health")
        assert resp.status_code == 200
        data = resp.json()
        assert data["status"] == "ok"

    def test_ready_returns_ready_after_startup(self, client: TestClient):
        resp = client.get("/ready")
        assert resp.status_code == 200
        data = resp.json()
        assert data["status"] == "ready"
        assert "model" in data
        assert "uptime_seconds" in data

    def test_metrics_endpoint(self, client: TestClient):
        resp = client.get("/metrics")
        assert resp.status_code == 200
        data = resp.json()
        assert data["specialist_model_loaded"] == 1
        assert "specialist_uptime_seconds" in data
        assert "specialist_max_batch_size" in data
        assert "specialist_max_queue_size" in data
        assert "specialist_total_batches" in data
        assert "specialist_total_items" in data
        assert "specialist_total_rejections" in data
        assert "specialist_queue_depth" in data


# ---------------------------------------------------------------------------
# Entity extraction
# ---------------------------------------------------------------------------


class TestEntityExtraction:
    """Test POST /api/specialist/entities."""

    def test_entity_extraction_returns_spans_with_offsets(self, client: TestClient):
        payload = {
            "texts": ["Apple Inc reported Q3 revenue of $81.4 billion."],
            "schema_labels": ["company", "financial_metric", "date"],
        }
        resp = client.post("/api/specialist/entities", json=payload)
        assert resp.status_code == 200
        data = resp.json()

        assert "results" in data
        assert "model_version" in data
        assert "schema_version" in data
        assert "processing_time_ms" in data
        assert data["processing_time_ms"] >= 0

        # Results should be a list of lists (one per input text)
        assert len(data["results"]) == 1
        entities = data["results"][0]

        # Should find at least one entity
        assert len(entities) > 0

        # Each entity should have required fields
        for entity in entities:
            assert "text" in entity
            assert "entity_type" in entity
            assert "start_char" in entity
            assert "end_char" in entity
            assert "score" in entity
            assert "model_version" in entity
            assert "schema_version" in entity
            assert entity["start_char"] >= 0
            assert entity["end_char"] > entity["start_char"]
            assert 0.0 <= entity["score"] <= 1.0

    def test_entity_extraction_character_offsets_match_source(self, client: TestClient):
        text = "Apple Inc reported Q3 revenue of $81.4 billion."
        payload = {
            "texts": [text],
            "schema_labels": ["company", "date"],
        }
        resp = client.post("/api/specialist/entities", json=payload)
        data = resp.json()
        entities = data["results"][0]

        for entity in entities:
            # The extracted text should match the source at the given offsets
            extracted_from_source = text[entity["start_char"]:entity["end_char"]]
            assert extracted_from_source == entity["text"]

    def test_entity_extraction_batch_multiple_texts(self, client: TestClient):
        payload = {
            "texts": [
                "Apple reported strong earnings.",
                "Tesla announced new factory plans.",
                "Microsoft acquired a small startup.",
            ],
            "schema_labels": ["company"],
        }
        resp = client.post("/api/specialist/entities", json=payload)
        assert resp.status_code == 200
        data = resp.json()
        # Should have one result list per input text
        assert len(data["results"]) == 3

    def test_entity_extraction_with_batch_id(self, client: TestClient):
        payload = {
            "texts": ["Apple Q3 results beat expectations."],
            "schema_labels": ["company"],
            "batch_id": "test-batch-001",
        }
        resp = client.post("/api/specialist/entities", json=payload)
        assert resp.status_code == 200
        data = resp.json()
        assert data["batch_id"] == "test-batch-001"


# ---------------------------------------------------------------------------
# Classification
# ---------------------------------------------------------------------------


class TestClassification:
    """Test POST /api/specialist/classify."""

    def test_classification_returns_labels_with_scores(self, client: TestClient):
        payload = {
            "texts": ["Apple reported quarterly earnings beating analyst expectations."],
            "schema_labels": ["earnings", "acquisition", "product_launch", "legal"],
        }
        resp = client.post("/api/specialist/classify", json=payload)
        assert resp.status_code == 200
        data = resp.json()

        assert len(data["results"]) == 1
        classifications = data["results"][0]
        assert len(classifications) > 0

        for cls in classifications:
            assert "text" in cls
            assert "label" in cls
            assert "score" in cls
            assert "model_version" in cls
            assert "schema_version" in cls
            assert 0.0 <= cls["score"] <= 1.0

    def test_classification_batch_processing(self, client: TestClient):
        payload = {
            "texts": [
                "Company announces merger.",
                "New product launched today.",
                "CEO resigned unexpectedly.",
            ],
            "schema_labels": ["acquisition", "product_launch", "management"],
        }
        resp = client.post("/api/specialist/classify", json=payload)
        assert resp.status_code == 200
        data = resp.json()
        assert len(data["results"]) == 3


# ---------------------------------------------------------------------------
# Relation extraction
# ---------------------------------------------------------------------------


class TestRelationExtraction:
    """Test POST /api/specialist/relations."""

    def test_relation_extraction_returns_triples(self, client: TestClient):
        payload = {
            "texts": ["Apple acquired Google subsidiary for $2 billion."],
            "schema_labels": ["acquired", "competes_with", "supplies"],
        }
        resp = client.post("/api/specialist/relations", json=payload)
        assert resp.status_code == 200
        data = resp.json()

        assert len(data["results"]) == 1
        relations = data["results"][0]

        # With Apple and Google in text, the mock should find a relation
        if relations:
            for rel in relations:
                assert "subject" in rel
                assert "subject_type" in rel
                assert "subject_start" in rel
                assert "subject_end" in rel
                assert "relation" in rel
                assert "object" in rel
                assert "object_type" in rel
                assert "object_start" in rel
                assert "object_end" in rel
                assert "score" in rel
                assert "model_version" in rel
                assert "schema_version" in rel
                assert 0.0 <= rel["score"] <= 1.0


# ---------------------------------------------------------------------------
# Structured extraction
# ---------------------------------------------------------------------------


class TestStructuredExtraction:
    """Test POST /api/specialist/extract."""

    def test_structured_extraction_returns_facts(self, client: TestClient):
        payload = {
            "texts": ["Revenue was $81.4 billion, up 8% year over year."],
            "schema_labels": ["revenue", "growth_rate"],
        }
        resp = client.post("/api/specialist/extract", json=payload)
        assert resp.status_code == 200
        data = resp.json()

        assert len(data["results"]) == 1
        structured = data["results"][0]

        if structured:
            for item in structured:
                assert "text" in item
                assert "field" in item
                assert "value" in item
                assert "start_char" in item
                assert "end_char" in item
                assert "score" in item
                assert "model_version" in item
                assert "schema_version" in item


# ---------------------------------------------------------------------------
# Model version in response
# ---------------------------------------------------------------------------


class TestModelVersion:
    """Test that model version and schema version are present in all responses."""

    def test_entity_response_contains_model_version(self, client: TestClient):
        payload = {
            "texts": ["Tesla reported record deliveries."],
            "schema_labels": ["company"],
        }
        resp = client.post("/api/specialist/entities", json=payload)
        data = resp.json()
        assert "model_version" in data
        assert "schema_version" in data
        assert data["schema_version"] == "specialist-v1"

    def test_classification_response_contains_model_version(self, client: TestClient):
        payload = {
            "texts": ["Earnings beat expectations."],
            "schema_labels": ["earnings"],
        }
        resp = client.post("/api/specialist/classify", json=payload)
        data = resp.json()
        assert "model_version" in data
        assert "schema_version" in data

    def test_relations_response_contains_model_version(self, client: TestClient):
        payload = {
            "texts": ["Apple and Google compete in AI."],
            "schema_labels": ["competes_with"],
        }
        resp = client.post("/api/specialist/relations", json=payload)
        data = resp.json()
        assert "model_version" in data
        assert "schema_version" in data

    def test_structured_response_contains_model_version(self, client: TestClient):
        payload = {
            "texts": ["Revenue was $50 billion."],
            "schema_labels": ["revenue"],
        }
        resp = client.post("/api/specialist/extract", json=payload)
        data = resp.json()
        assert "model_version" in data
        assert "schema_version" in data


# ---------------------------------------------------------------------------
# Batch size enforcement
# ---------------------------------------------------------------------------


class TestBatchSizeEnforcement:
    """Test that exceeding max_batch_size is rejected."""

    def test_exceeding_max_batch_size_returns_422(self, client: TestClient):
        # Default max_batch_size is 32, send 33 texts
        texts = [f"Text number {i}" for i in range(33)]
        payload = {
            "texts": texts,
            "schema_labels": ["company"],
        }
        resp = client.post("/api/specialist/entities", json=payload)
        assert resp.status_code == 422
        data = resp.json()
        assert "maximum" in data["detail"].lower() or "exceeds" in data["detail"].lower()

    def test_at_max_batch_size_succeeds(self, client: TestClient):
        # 32 texts should be fine
        texts = [f"Apple reported earnings for period {i}." for i in range(32)]
        payload = {
            "texts": texts,
            "schema_labels": ["company"],
        }
        resp = client.post("/api/specialist/entities", json=payload)
        assert resp.status_code == 200
        data = resp.json()
        assert len(data["results"]) == 32

    def test_empty_texts_rejected(self, client: TestClient):
        payload = {
            "texts": [],
            "schema_labels": ["company"],
        }
        resp = client.post("/api/specialist/entities", json=payload)
        # Pydantic min_length=1 should reject this
        assert resp.status_code == 422

    def test_empty_labels_rejected(self, client: TestClient):
        payload = {
            "texts": ["Some text"],
            "schema_labels": [],
        }
        resp = client.post("/api/specialist/entities", json=payload)
        assert resp.status_code == 422

    def test_classify_enforces_batch_limit(self, client: TestClient):
        texts = [f"Text {i}" for i in range(33)]
        payload = {"texts": texts, "schema_labels": ["earnings"]}
        resp = client.post("/api/specialist/classify", json=payload)
        assert resp.status_code == 422

    def test_relations_enforces_batch_limit(self, client: TestClient):
        texts = [f"Text {i}" for i in range(33)]
        payload = {"texts": texts, "schema_labels": ["competes_with"]}
        resp = client.post("/api/specialist/relations", json=payload)
        assert resp.status_code == 422

    def test_extract_enforces_batch_limit(self, client: TestClient):
        texts = [f"Text {i}" for i in range(33)]
        payload = {"texts": texts, "schema_labels": ["revenue"]}
        resp = client.post("/api/specialist/extract", json=payload)
        assert resp.status_code == 422


# ---------------------------------------------------------------------------
# Load test (lightweight simulation)
# ---------------------------------------------------------------------------


class TestLoadSimulation:
    """Basic load simulation — process many batches sequentially."""

    def test_sequential_batch_throughput(self, client: TestClient):
        """Process 10 batches of 10 texts and ensure consistent results."""
        total_ms = 0.0
        for i in range(10):
            payload = {
                "texts": [f"Apple reported Q{j % 4 + 1} results." for j in range(10)],
                "schema_labels": ["company", "date"],
            }
            resp = client.post("/api/specialist/entities", json=payload)
            assert resp.status_code == 200
            data = resp.json()
            assert len(data["results"]) == 10
            total_ms += data["processing_time_ms"]

        # All 100 documents processed — just confirm no errors
        assert total_ms >= 0

    def test_mixed_endpoint_load(self, client: TestClient):
        """Call all endpoints in sequence to simulate mixed load."""
        entity_payload = {
            "texts": ["Apple Q3 revenue beat."],
            "schema_labels": ["company"],
        }
        classify_payload = {
            "texts": ["Major acquisition announced."],
            "schema_labels": ["acquisition", "earnings"],
        }
        relation_payload = {
            "texts": ["Apple and Google compete in phones."],
            "schema_labels": ["competes_with"],
        }
        structured_payload = {
            "texts": ["Revenue was $50 billion."],
            "schema_labels": ["revenue"],
        }

        for _ in range(5):
            assert client.post("/api/specialist/entities", json=entity_payload).status_code == 200
            assert client.post("/api/specialist/classify", json=classify_payload).status_code == 200
            assert client.post("/api/specialist/relations", json=relation_payload).status_code == 200
            assert client.post("/api/specialist/extract", json=structured_payload).status_code == 200

    def test_concurrent_batch_load(self, client: TestClient):
        """Simulate rapid sequential calls to stress the service."""
        payload = {
            "texts": [f"Company {i} announced results." for i in range(16)],
            "schema_labels": ["company", "earnings", "date"],
        }
        # 20 rapid sequential requests
        for _ in range(20):
            resp = client.post("/api/specialist/entities", json=payload)
            assert resp.status_code == 200
            data = resp.json()
            assert len(data["results"]) == 16
            assert data["processing_time_ms"] >= 0


# ---------------------------------------------------------------------------
# Dynamic batching
# ---------------------------------------------------------------------------


class TestDynamicBatching:
    """Test that the DynamicBatcher correctly collects and processes requests."""

    @pytest.mark.asyncio
    async def test_batcher_processes_single_item(self):
        """Single item submitted should be processed as a batch of one."""
        processed_batches: list[list] = []

        def process_fn(payloads):
            processed_batches.append(payloads)
            return [p * 2 for p in payloads]

        batcher = DynamicBatcher(
            max_batch_size=4, max_wait_ms=50.0, max_queue_size=16
        )
        batcher.start(process_fn)

        result = await batcher.submit(5)
        assert result == 10
        assert len(processed_batches) >= 1

        await batcher.stop()

    @pytest.mark.asyncio
    async def test_batcher_collects_concurrent_items(self):
        """Multiple concurrent submissions should be batched together."""
        processed_batches: list[list] = []

        def process_fn(payloads):
            processed_batches.append(list(payloads))
            return [p + 100 for p in payloads]

        batcher = DynamicBatcher(
            max_batch_size=8, max_wait_ms=200.0, max_queue_size=64
        )
        batcher.start(process_fn)

        results = await asyncio.gather(
            batcher.submit(1),
            batcher.submit(2),
            batcher.submit(3),
            batcher.submit(4),
        )

        assert sorted(results) == [101, 102, 103, 104]
        total_items = sum(len(b) for b in processed_batches)
        assert total_items == 4

        await batcher.stop()

    @pytest.mark.asyncio
    async def test_batcher_respects_max_batch_size(self):
        """Batcher should not exceed max_batch_size per batch."""
        batch_sizes: list[int] = []

        def process_fn(payloads):
            batch_sizes.append(len(payloads))
            return list(range(len(payloads)))

        batcher = DynamicBatcher(
            max_batch_size=3, max_wait_ms=500.0, max_queue_size=64
        )
        batcher.start(process_fn)

        await asyncio.gather(
            batcher.submit("a"),
            batcher.submit("b"),
            batcher.submit("c"),
            batcher.submit("d"),
            batcher.submit("e"),
        )

        for size in batch_sizes:
            assert size <= 3

        await batcher.stop()

    @pytest.mark.asyncio
    async def test_batcher_metrics_tracked(self):
        """Batcher should track processed items and batches."""

        def process_fn(payloads):
            return [None] * len(payloads)

        batcher = DynamicBatcher(
            max_batch_size=4, max_wait_ms=50.0, max_queue_size=16
        )
        batcher.start(process_fn)

        await asyncio.gather(
            batcher.submit("x"),
            batcher.submit("y"),
        )
        await asyncio.sleep(0.1)

        assert batcher.total_items_processed >= 2
        assert batcher.total_batches_processed >= 1
        assert batcher.total_rejections == 0

        await batcher.stop()


# ---------------------------------------------------------------------------
# Bounded queue rejection
# ---------------------------------------------------------------------------


class TestBoundedQueue:
    """Test that the bounded queue rejects overflow."""

    @pytest.mark.asyncio
    async def test_queue_full_raises_error(self):
        """When the queue is full, new submissions raise QueueFullError."""
        batcher = DynamicBatcher(
            max_batch_size=32, max_wait_ms=50.0, max_queue_size=3
        )
        # Intentionally NOT calling batcher.start() — no background loop
        # means items remain in the queue.
        batcher._running = True  # Allow submit to not raise other issues

        # Fill the queue to capacity
        loop = asyncio.get_running_loop()
        for i in range(3):
            from services.specialist.batching import _PendingRequest
            pending = _PendingRequest(
                payload=i,
                future=loop.create_future(),
            )
            batcher._queue.put_nowait(pending)

        # Queue is full — next submit should raise QueueFullError
        with pytest.raises(QueueFullError):
            await batcher.submit(999)

        assert batcher.total_rejections >= 1
        assert batcher.queue_size == 3

    @pytest.mark.asyncio
    async def test_queue_size_property(self):
        """queue_size should reflect current pending items."""

        def process_fn(payloads):
            return [None] * len(payloads)

        batcher = DynamicBatcher(
            max_batch_size=32, max_wait_ms=500.0, max_queue_size=100
        )
        batcher.start(process_fn)

        assert batcher.queue_size == 0
        await batcher.submit("test")
        await asyncio.sleep(0.15)
        assert batcher.queue_size == 0

        await batcher.stop()
