"""Integration tests for chatCompletionOverDataset GraphQL subscription."""

from __future__ import annotations

import asyncio
import json
from typing import Any, cast

import httpx
import pytest

from tests.integration._helpers import _AppInfo, _parse_apollo_multipart_response

from .conftest import CustomProviders, DatasetEvaluators, _gql

# =============================================================================
# GraphQL Operations
# =============================================================================

CHAT_COMPLETION_OVER_DATASET = """
subscription ChatCompletionOverDataset($input: ChatCompletionOverDatasetInput!) {
    chatCompletionOverDataset(input: $input) {
        __typename
        ... on TextChunk {
            content
            datasetExampleId
            repetitionNumber
        }
        ... on ChatCompletionSubscriptionExperiment {
            experiment {
                id
                name
            }
        }
        ... on ChatCompletionSubscriptionResult {
            datasetExampleId
            repetitionNumber
            span {
                id
            }
            experimentRun {
                id
            }
        }
        ... on ChatCompletionSubscriptionError {
            datasetExampleId
            repetitionNumber
            message
        }
        ... on EvaluationChunk {
            datasetExampleId
            repetitionNumber
            evaluatorName
            error
            experimentRunEvaluation {
                id
                name
                label
                score
            }
        }
    }
}
"""

RESUME_EXPERIMENT = """
mutation ResumeExperiment($experimentId: ID!) {
    resumeExperiment(experimentId: $experimentId) {
        job {
            id
            status
        }
    }
}
"""

GET_EXPERIMENT_JOB_STATUS = """
query GetExperimentJobStatus($id: ID!) {
    node(id: $id) {
        ... on ExperimentJob {
            id
            status
        }
    }
}
"""

GET_EXPERIMENT = """
query GetExperiment($id: ID!) {
    node(id: $id) {
        ... on Experiment {
            id
            name
            description
            repetitions
            runs(first: 100) {
                edges {
                    node {
                        id
                        output
                        error
                        traceId
                        trace {
                            id
                            traceId
                            numSpans
                            spans(first: 100) {
                                edges {
                                    node {
                                        id
                                        name
                                        latencyMs
                                        tokenCountPrompt
                                        tokenCountCompletion
                                        tokenCountTotal
                                    }
                                }
                            }
                        }
                        annotations {
                            edges {
                                node {
                                    id
                                    name
                                    label
                                    score
                                    explanation
                                    error
                                    traceId
                                    trace {
                                        id
                                        traceId
                                        numSpans
                                        spans(first: 100) {
                                            edges {
                                                node {
                                                    id
                                                    name
                                                    tokenCountTotal
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
}
"""


# =============================================================================
# Helper Functions
# =============================================================================


async def _run_subscription_with_evaluators(
    client: httpx.AsyncClient,
    *,
    custom_provider_id: str,
    dataset_id: str,
    evaluators: DatasetEvaluators,
    stream_model_output: bool = True,
) -> tuple[str, list[dict[str, Any]]]:
    """Run chatCompletionOverDataset subscription with evaluators.

    Uses the custom OpenAI provider pointing to mock server.
    Includes all 4 evaluators (one per provider type).

    Returns a tuple of (experiment_id, subscription_results) where
    subscription_results is a list of ChatCompletionSubscriptionResult payloads.
    """
    variables = {
        "input": {
            "promptVersion": {
                "templateFormat": "MUSTACHE",
                "template": {
                    "messages": [
                        {
                            "role": "SYSTEM",
                            "content": [{"text": {"text": "You are a helpful assistant."}}],
                        },
                        {
                            "role": "USER",
                            "content": [{"text": {"text": "{{question}}"}}],
                        },
                    ]
                },
                "invocationParameters": {"openai": {}},
                "modelProvider": "OPENAI",
                "modelName": "gpt-4o-mini",
                "customProviderId": custom_provider_id,
            },
            "repetitions": 2,
            "datasetId": dataset_id,
            "evaluators": [
                {
                    "id": evaluators.openai,
                    "name": "openai_eval",
                    "inputMapping": {
                        "pathMapping": {
                            "input": "$.input",
                            "output": "$.output",
                            "reference": "$.reference",
                        },
                        "literalMapping": {},
                    },
                },
                {
                    "id": evaluators.openai_responses,
                    "name": "openai_responses_eval",
                    "inputMapping": {
                        "pathMapping": {
                            "input": "$.input",
                            "output": "$.output",
                            "reference": "$.reference",
                        },
                        "literalMapping": {},
                    },
                },
                {
                    "id": evaluators.anthropic,
                    "name": "anthropic_eval",
                    "inputMapping": {
                        "pathMapping": {
                            "input": "$.input",
                            "output": "$.output",
                            "reference": "$.reference",
                        },
                        "literalMapping": {},
                    },
                },
                {
                    "id": evaluators.google_genai,
                    "name": "google_eval",
                    "inputMapping": {
                        "pathMapping": {
                            "input": "$.input",
                            "output": "$.output",
                            "reference": "$.reference",
                        },
                        "literalMapping": {},
                    },
                },
                {
                    "id": evaluators.bedrock,
                    "name": "bedrock_eval",
                    "inputMapping": {
                        "pathMapping": {
                            "input": "$.input",
                            "output": "$.output",
                            "reference": "$.reference",
                        },
                        "literalMapping": {},
                    },
                },
            ],
            "streamModelOutput": stream_model_output,
        }
    }

    async def _execute_subscription() -> tuple[str, list[dict[str, Any]]]:
        """Execute the subscription and return the experiment ID and results."""
        experiment_id: str | None = None
        errors_seen: list[dict[str, Any]] = []
        completion_errors_seen: list[dict[str, Any]] = []
        eval_errors_seen: list[dict[str, Any]] = []
        results_seen: list[dict[str, Any]] = []

        # Apollo multipart subscription protocol
        async with client.stream(
            "POST",
            "/graphql",
            json={"query": CHAT_COMPLETION_OVER_DATASET, "variables": variables},
            headers={
                "Accept": "multipart/mixed;boundary=graphql;subscriptionSpec=1.0,application/json",
                "Content-Type": "application/json",
            },
        ) as response:
            response.raise_for_status()

            # Parse multipart response
            async for payload in _parse_apollo_multipart_response(response):
                # Apollo multipart wraps data in a "payload" field
                inner = payload.get("payload", payload)
                # Capture any errors for debugging
                if errs := inner.get("errors"):
                    errors_seen.extend(errs)
                if data := inner.get("data"):
                    subscription_data = data.get("chatCompletionOverDataset")
                    if subscription_data:
                        typename = subscription_data.get("__typename")
                        if typename == "ChatCompletionSubscriptionExperiment":
                            experiment_id = str(subscription_data["experiment"]["id"])
                        elif typename == "ChatCompletionSubscriptionResult":
                            results_seen.append(subscription_data)
                        elif typename == "ChatCompletionSubscriptionError":
                            completion_errors_seen.append(subscription_data)
                        elif typename == "EvaluationChunk":
                            if subscription_data.get("error"):
                                eval_errors_seen.append(subscription_data)

            assert not errors_seen, f"GraphQL errors during subscription: {errors_seen}"
            assert not completion_errors_seen, (
                f"Chat completion / playground errors during subscription: {completion_errors_seen}"
            )
            assert not eval_errors_seen, (
                f"Evaluation errors received during subscription: {eval_errors_seen}"
            )

        if experiment_id is None:
            error_msg = "Did not receive experiment ID from subscription"
            if errors_seen:
                error_msg += f"\nGraphQL errors: {errors_seen}"
            raise AssertionError(error_msg)

        return experiment_id, results_seen

    # Run subscription with a 120 second timeout
    try:
        return await asyncio.wait_for(_execute_subscription(), timeout=120.0)
    except asyncio.TimeoutError:
        raise AssertionError("Subscription timed out after 120 seconds")


async def _run_subscription_without_evaluators(
    client: httpx.AsyncClient,
    *,
    custom_provider_id: str,
    model_name: str,
    model_provider: str,
    dataset_id: str,
    invocation_parameters: dict[str, Any],
    stream_model_output: bool = True,
) -> tuple[str, list[dict[str, Any]]]:
    """Run chatCompletionOverDataset subscription without evaluators.

    Tests the primary prompt with a specific custom provider.

    Returns ``(experiment_id, subscription_results)`` where ``subscription_results``
    lists ``ChatCompletionSubscriptionResult`` payloads from the stream.
    """
    variables = {
        "input": {
            "promptVersion": {
                "templateFormat": "MUSTACHE",
                "template": {
                    "messages": [
                        {
                            "role": "SYSTEM",
                            "content": [{"text": {"text": "You are a helpful assistant."}}],
                        },
                        {
                            "role": "USER",
                            "content": [{"text": {"text": "{{question}}"}}],
                        },
                    ]
                },
                "invocationParameters": invocation_parameters,
                "modelProvider": model_provider,
                "modelName": model_name,
                "customProviderId": custom_provider_id,
            },
            "repetitions": 1,
            "datasetId": dataset_id,
            "evaluators": [],
            "streamModelOutput": stream_model_output,
        }
    }

    async def _execute_subscription() -> tuple[str, list[dict[str, Any]]]:
        """Execute the subscription and return the experiment ID and results."""
        experiment_id: str | None = None
        errors_seen: list[dict[str, Any]] = []
        completion_errors_seen: list[dict[str, Any]] = []
        results_seen: list[dict[str, Any]] = []

        # Apollo multipart subscription protocol
        async with client.stream(
            "POST",
            "/graphql",
            json={"query": CHAT_COMPLETION_OVER_DATASET, "variables": variables},
            headers={
                "Accept": "multipart/mixed;boundary=graphql;subscriptionSpec=1.0,application/json",
                "Content-Type": "application/json",
            },
        ) as response:
            response.raise_for_status()

            # Parse multipart response
            async for payload in _parse_apollo_multipart_response(response):
                # Apollo multipart wraps data in a "payload" field
                inner = payload.get("payload", payload)
                # Capture any errors for debugging
                if errs := inner.get("errors"):
                    errors_seen.extend(errs)
                if data := inner.get("data"):
                    subscription_data = data.get("chatCompletionOverDataset")
                    if subscription_data:
                        typename = subscription_data.get("__typename")
                        if typename == "ChatCompletionSubscriptionExperiment":
                            experiment_id = str(subscription_data["experiment"]["id"])
                        elif typename == "ChatCompletionSubscriptionResult":
                            results_seen.append(subscription_data)
                        elif typename == "ChatCompletionSubscriptionError":
                            completion_errors_seen.append(subscription_data)

            assert not errors_seen, f"GraphQL errors during subscription: {errors_seen}"
            assert not completion_errors_seen, (
                f"Chat completion errors during subscription: {completion_errors_seen}"
            )

        if experiment_id is None:
            error_msg = "Did not receive experiment ID from subscription"
            if errors_seen:
                error_msg += f"\nGraphQL errors: {errors_seen}"
            raise AssertionError(error_msg)

        return experiment_id, results_seen

    # Run subscription with a 60 second timeout (no evaluators = faster)
    try:
        return await asyncio.wait_for(_execute_subscription(), timeout=60.0)
    except asyncio.TimeoutError:
        raise AssertionError("Subscription timed out after 60 seconds")


async def _wait_for_job_status(
    client: httpx.AsyncClient,
    job_id: str,
    expected_status: str = "COMPLETED",
    timeout: float = 30.0,
) -> None:
    """Poll an ExperimentJob until it reaches the expected status."""

    async def _poll() -> None:
        for _ in range(int(timeout / 0.5)):
            data = await _gql(client, GET_EXPERIMENT_JOB_STATUS, {"id": job_id})
            if data["node"]["status"] == expected_status:
                return
            await asyncio.sleep(0.5)
        raise AssertionError(
            f"Experiment job {job_id} did not reach {expected_status} within {timeout}s"
        )

    await asyncio.wait_for(_poll(), timeout=timeout)


async def _get_experiment(
    client: httpx.AsyncClient,
    experiment_id: str,
) -> dict[str, Any] | None:
    """Query experiment by ID."""
    data = await _gql(client, GET_EXPERIMENT, {"id": experiment_id})
    node = data.get("node")
    return cast(dict[str, Any], node) if node is not None else None


# =============================================================================
# Tests
# =============================================================================


class TestChatCompletionOverDataset:
    """Tests for chatCompletionOverDataset subscription with evaluators."""

    @pytest.mark.parametrize(
        "stream_model_output",
        [
            pytest.param(True, id="streaming"),
            pytest.param(False, id="non_streaming"),
        ],
    )
    async def test_experiment_with_evaluators(
        self,
        _app: _AppInfo,
        _custom_providers: CustomProviders,
        _dataset_id: str,
        _dataset_evaluators: DatasetEvaluators,
        stream_model_output: bool,
    ) -> None:
        """Test chatCompletionOverDataset with all 4 provider evaluators.

        Uses OpenAI custom provider (pointing to mock server) for the experiment,
        and attaches evaluators using all 4 custom providers (OpenAI, Anthropic,
        Google GenAI, Bedrock).
        """
        # Long timeout so stream() read doesn't hit default 5s during slow CI
        async with httpx.AsyncClient(base_url=_app.base_url, timeout=120.0) as client:
            # Run subscription with all evaluators
            experiment_id, subscription_results = await _run_subscription_with_evaluators(
                client,
                custom_provider_id=_custom_providers.openai,
                dataset_id=_dataset_id,
                evaluators=_dataset_evaluators,
                stream_model_output=stream_model_output,
            )

            # Query experiment
            experiment = await _get_experiment(client, experiment_id)
            assert experiment is not None

            # Verify we have 4 runs (2 examples × 2 repetitions)
            runs = experiment["runs"]["edges"]
            assert len(runs) == 4, (
                f"Expected 4 runs, got {len(runs)}. "
                f"Subscription results: {json.dumps(subscription_results)}"
            )

            # Verify each run has spans and evaluations
            for run in runs:
                run_node = run["node"]
                assert run_node["error"] is None, f"Run had error: {run_node['error']}"

                # Verify trace and spans exist
                assert run_node["traceId"] is not None, "Run should have a traceId"
                trace = run_node["trace"]
                assert trace is not None, "Run should have a trace"
                assert trace["numSpans"] >= 1, "Trace should have at least 1 span"

                spans = trace["spans"]["edges"]
                assert len(spans) >= 1, "Trace should have at least 1 span"

                # Verify span has expected fields
                for span in spans:
                    span_node = span["node"]
                    assert span_node["id"] is not None
                    assert span_node["name"] is not None
                    assert span_node["latencyMs"] is not None

                # Verify token counts sum to > 0 for spans in the trace
                total_tokens = sum((span["node"]["tokenCountTotal"] or 0) for span in spans)
                assert total_tokens > 0, (
                    f"Total token count for trace should be > 0, got {total_tokens}"
                )

                # Verify we have 5 annotations (one per evaluator)
                annotations = run_node["annotations"]["edges"]
                assert len(annotations) == 5, (
                    f"Expected 5 annotations per run, got {len(annotations)}"
                )

                # Verify we have annotations from all 5 evaluators
                annotation_names = {ann["node"]["name"] for ann in annotations}
                expected_names = {
                    "openai_eval",
                    "openai_responses_eval",
                    "anthropic_eval",
                    "google_eval",
                    "bedrock_eval",
                }
                assert annotation_names == expected_names, (
                    f"Expected annotations {expected_names}, got {annotation_names}"
                )

                # Verify each evaluator annotation has its own trace with token counts
                for ann in annotations:
                    ann_node = ann["node"]
                    ann_name = ann_node["name"]

                    # Verify annotation has label, score, and explanation
                    assert ann_node["label"] is not None, (
                        f"Evaluator {ann_name} should have a label, error={ann_node.get('error')}"
                    )
                    assert ann_node["score"] is not None, (
                        f"Evaluator {ann_name} should have a score"
                    )
                    assert ann_node["explanation"] is not None, (
                        f"Evaluator {ann_name} should have an explanation"
                    )

                    # Each evaluator should have a trace
                    assert ann_node["traceId"] is not None, (
                        f"Evaluator {ann_name} should have a traceId"
                    )
                    ann_trace = ann_node["trace"]
                    assert ann_trace is not None, f"Evaluator {ann_name} should have a trace"
                    assert ann_trace["numSpans"] >= 1, (
                        f"Evaluator {ann_name} trace should have at least 1 span"
                    )

                    # Sum token counts for all spans in the evaluator trace
                    ann_spans = ann_trace["spans"]["edges"]
                    ann_total_tokens = sum(
                        (span["node"]["tokenCountTotal"] or 0) for span in ann_spans
                    )
                    assert ann_total_tokens > 0, (
                        f"Evaluator {ann_name} trace should have token counts > 0, "
                        f"got {ann_total_tokens}"
                    )

            # Resume the completed experiment and verify it completes quickly.
            resume_data = await _gql(client, RESUME_EXPERIMENT, {"experimentId": experiment_id})
            job = resume_data["resumeExperiment"]["job"]
            assert job["status"] == "RUNNING"
            await _wait_for_job_status(client, job["id"])

    @pytest.mark.parametrize(
        "stream_model_output",
        [
            pytest.param(True, id="streaming"),
            pytest.param(False, id="non_streaming"),
        ],
    )
    @pytest.mark.parametrize(
        "provider_key,model_name,model_provider,invocation_parameters",
        [
            pytest.param("openai", "gpt-4o-mini", "OPENAI", {"openai": {}}, id="openai"),
            pytest.param(
                "openai_responses",
                "gpt-4o-mini",
                "OPENAI",
                {"openai": {}},
                id="openai_responses",
            ),
            pytest.param(
                "anthropic",
                "claude-3-5-sonnet-latest",
                "ANTHROPIC",
                {"anthropic": {"maxTokens": 1024}},
                id="anthropic",
            ),
            pytest.param(
                "google_genai", "gemini-2.0-flash", "GOOGLE", {"google": {}}, id="google_genai"
            ),
            pytest.param(
                "bedrock",
                "anthropic.claude-3-haiku-20240307-v1:0",
                "AWS",
                {"aws": {}},
                id="bedrock",
            ),
        ],
    )
    async def test_provider_without_evaluators(
        self,
        _app: _AppInfo,
        _custom_providers: CustomProviders,
        _dataset_id: str,
        provider_key: str,
        model_name: str,
        model_provider: str,
        invocation_parameters: dict[str, Any],
        stream_model_output: bool,
    ) -> None:
        """Test chatCompletionOverDataset with a provider on primary prompt.

        Tests the custom provider for the primary prompt without any evaluators.
        This exercises the text generation path for the provider.
        """
        # Look up provider ID from the fixture
        provider_id = getattr(_custom_providers, provider_key)

        # Long timeout so stream() read doesn't hit default 5s during slow CI
        async with httpx.AsyncClient(base_url=_app.base_url, timeout=60.0) as client:
            # Run subscription without evaluators
            experiment_id, subscription_results = await _run_subscription_without_evaluators(
                client,
                custom_provider_id=provider_id,
                model_name=model_name,
                model_provider=model_provider,
                dataset_id=_dataset_id,
                invocation_parameters=invocation_parameters,
                stream_model_output=stream_model_output,
            )

            # Query experiment
            experiment = await _get_experiment(client, experiment_id)
            assert experiment is not None, "Experiment should exist"

            # Verify we have 2 runs (2 examples × 1 repetition)
            runs = experiment["runs"]["edges"]
            assert len(runs) == 2, (
                f"Expected 2 runs, got {len(runs)}. "
                f"Subscription results: {json.dumps(subscription_results)}"
            )

            # Verify each run completed successfully
            for run in runs:
                run_node = run["node"]
                assert run_node["error"] is None, f"Run had error: {run_node['error']}"

                # Verify output is not empty (echo response)
                assert run_node["output"] is not None, "Run should have output"

                # Verify trace and spans exist
                assert run_node["traceId"] is not None, "Run should have a traceId"
                trace = run_node["trace"]
                assert trace is not None, "Run should have a trace"
                assert trace["numSpans"] >= 1, "Trace should have at least 1 span"

                spans = trace["spans"]["edges"]
                assert len(spans) >= 1, "Trace should have at least 1 span"

                # Verify span has expected fields
                for span in spans:
                    span_node = span["node"]
                    assert span_node["id"] is not None, "Span should have id"
                    assert span_node["name"] is not None, "Span should have name"
                    assert span_node["latencyMs"] is not None, "Span should have latencyMs"

                # Verify token counts sum to > 0 for spans in the trace
                total_tokens = sum((span["node"]["tokenCountTotal"] or 0) for span in spans)
                assert total_tokens > 0, f"Total token count should be > 0, got {total_tokens}"

                # Verify no annotations (since no evaluators)
                annotations = run_node["annotations"]["edges"]
                assert len(annotations) == 0, f"Expected 0 annotations, got {len(annotations)}"

            # Resume the completed experiment and verify it completes quickly.
            resume_data = await _gql(client, RESUME_EXPERIMENT, {"experimentId": experiment_id})
            job = resume_data["resumeExperiment"]["job"]
            assert job["status"] == "RUNNING"
            await _wait_for_job_status(client, job["id"])
