"""E2E tests for multi-model-type workflows.

Tests verify that all SDK model types (Internal, Generative/LLM, Agentic, Vendor)
successfully complete the full telemetry workflow with proper attribute validation.
"""

import time
import pytest
from mca_sdk import MCAClient


@pytest.mark.integration
@pytest.mark.error_scenario
class TestMultiModelE2E:
    """End-to-end tests for different model categories and types."""

    def test_internal_model_full_workflow(
        self,
        docker_compose_environment,
        gcp_monitoring_client,
        gcp_logging_client,
        gcp_trace_client,
        gcp_project_id,
        wait_for_telemetry,
        cleanup_collector_exports,
    ):
        """Test internal/regression model completes full E2E workflow.

        Given: Internal regression model configured
        When: Model records prediction with metrics, logs, traces
        Then: All telemetry appears in respective GCP services
        """
        # Arrange
        client = MCAClient(
            service_name="e2e-internal-model",
            model_id="e2e-int-mdl-001",
            team_name="clinical-ai",
            model_category="internal",
            model_type="regression",
            otel_endpoint=docker_compose_environment["collector_otlp"],
            otel_protocol="http",
        )

        start_time = time.time()

        # Act: Execute prediction workflow
        @client.predict()
        def predict_readmission_risk(patient_data: dict) -> dict:
            """Predict 30-day readmission risk."""
            time.sleep(0.05)
            return {"risk_score": 0.72, "confidence": 0.85}

        result = predict_readmission_risk({"age": 65, "comorbidities": 3})

        client.shutdown()
        wait_for_telemetry(10, "internal model telemetry")

        # Assert: Verify metric exists and has correct type
        end_time = time.time()
        metrics = gcp_monitoring_client.query_metrics(
            project_id=gcp_project_id,
            metric_type="mca.internal.regression",
            start_time=start_time,
            end_time=end_time,
        )

        assert len(metrics) > 0, "Internal model metrics not found"

        # Verify at least one metric has the expected attributes
        found_internal_metric = False
        for metric in metrics:
            attrs = metric.get("attributes", {})
            model_category = attrs.get("model.category") or attrs.get("model_category")
            model_type = attrs.get("model.type") or attrs.get("model_type")

            if model_category == "internal" and model_type == "regression":
                found_internal_metric = True
                break

        assert (
            found_internal_metric
        ), "No metrics found with model_category=internal and model_type=regression"

    def test_generative_llm_model_full_workflow(
        self,
        docker_compose_environment,
        gcp_monitoring_client,
        gcp_logging_client,
        gcp_trace_client,
        gcp_project_id,
        wait_for_telemetry,
        cleanup_collector_exports,
    ):
        """Test generative/LLM model completes full E2E workflow.

        Given: Generative LLM model configured
        When: Model generates completion with token tracking
        Then: Token counts are accurately captured in metrics
        """
        # Arrange
        client = MCAClient(
            service_name="e2e-genai-model",
            model_id="e2e-llm-mdl-001",
            team_name="clinical-ai",
            model_category="generative",
            model_type="llm",
            otel_endpoint=docker_compose_environment["collector_otlp"],
            otel_protocol="http",
        )

        start_time = time.time()

        # Act: Record LLM completion with specific token counts
        expected_input_tokens = 50
        expected_output_tokens = 100
        expected_total_tokens = 150

        client.record_llm_completion(
            prompt="Summarize patient history",
            completion="Patient has history of hypertension...",
            model_name="gpt-4",
            input_tokens=expected_input_tokens,
            output_tokens=expected_output_tokens,
            total_tokens=expected_total_tokens,
            latency=1.2,
            cost=0.005,
        )

        client.shutdown()
        wait_for_telemetry(10, "LLM telemetry with token tracking")

        # Assert: Verify token metrics with exact counts
        end_time = time.time()

        # Query for token metrics
        token_metrics = gcp_monitoring_client.query_metrics(
            project_id=gcp_project_id,
            metric_type="tokens",  # Match any token-related metric
            start_time=start_time,
            end_time=end_time,
        )

        assert len(token_metrics) > 0, "No token metrics found for LLM completion"

        # Verify specific token counts are captured
        found_input = False
        found_output = False

        for metric in token_metrics:
            attrs = metric.get("attributes", {})
            value = metric.get("value", 0)

            # Check for input tokens
            token_type = attrs.get("token.type") or attrs.get("token_type")
            if token_type == "input" or "input" in metric.get("metric_type", ""):
                if value == expected_input_tokens:
                    found_input = True

            # Check for output tokens
            if token_type == "output" or "output" in metric.get("metric_type", ""):
                if value == expected_output_tokens:
                    found_output = True

        # At minimum, verify total tokens or that input/output were captured
        assert (
            found_input
            or found_output
            or any(m.get("value") == expected_total_tokens for m in token_metrics)
        ), (
            f"Token counts not accurately captured. Expected input={expected_input_tokens}, "
            f"output={expected_output_tokens}. Found metrics: {token_metrics}"
        )

    def test_agentic_ai_model_full_workflow(
        self,
        docker_compose_environment,
        gcp_monitoring_client,
        gcp_trace_client,
        gcp_project_id,
        wait_for_telemetry,
        cleanup_collector_exports,
    ):
        """Test agentic AI model completes full E2E workflow.

        Given: Agentic AI model with goals and tools
        When: Agent executes multi-step reasoning workflow
        Then: Goal outcome and tool call states are correctly captured
        """
        # Arrange
        client = MCAClient(
            service_name="e2e-agentic-model",
            model_id="e2e-agent-mdl-001",
            team_name="clinical-ai",
            model_category="agentic",
            model_type="agent",
            otel_endpoint=docker_compose_environment["collector_otlp"],
            otel_protocol="http",
        )

        start_time = time.time()

        # Act: Execute agentic workflow with specific outcome
        expected_outcome = "diagnosis_complete"
        goal_id = client.start_goal("diagnose_patient", agent_id="agent-001")

        # Simulate tool calls
        client.record_tool_call(
            goal_id=goal_id,
            tool_name="query_patient_history",
            status="success",
            latency=0.05,
        )

        client.record_tool_call(
            goal_id=goal_id,
            tool_name="analyze_lab_results",
            status="success",
            latency=0.08,
        )

        client.complete_goal(goal_id, outcome=expected_outcome)

        client.shutdown()
        wait_for_telemetry(10, "agentic AI telemetry")

        # Assert: Verify goal metrics with outcome
        end_time = time.time()
        goal_metrics = gcp_monitoring_client.query_metrics(
            project_id=gcp_project_id,
            metric_type="goal",
            start_time=start_time,
            end_time=end_time,
        )

        assert len(goal_metrics) > 0, "Agentic goal metrics not found"

        # Verify outcome is captured
        found_outcome = False
        for metric in goal_metrics:
            attrs = metric.get("attributes", {})
            outcome = attrs.get("goal.outcome") or attrs.get("outcome")

            if outcome == expected_outcome:
                found_outcome = True
                break

        assert found_outcome, (
            f"Goal outcome '{expected_outcome}' not found in metrics. "
            f"Available outcomes: {[m.get('attributes', {}).get('outcome') for m in goal_metrics]}"
        )

        # Verify tool metrics
        tool_metrics = gcp_monitoring_client.query_metrics(
            project_id=gcp_project_id,
            metric_type="tool",
            start_time=start_time,
            end_time=end_time,
        )

        assert (
            len(tool_metrics) >= 2
        ), f"Expected at least 2 tool call metrics, found {len(tool_metrics)}"

        # Verify tool names are captured
        tool_names = []
        for metric in tool_metrics:
            attrs = metric.get("attributes", {})
            tool_name = attrs.get("tool.name") or attrs.get("tool_name")
            if tool_name:
                tool_names.append(tool_name)

        assert (
            "query_patient_history" in tool_names
        ), f"Tool 'query_patient_history' not found. Available: {tool_names}"
        assert (
            "analyze_lab_results" in tool_names
        ), f"Tool 'analyze_lab_results' not found. Available: {tool_names}"

    def test_vendor_model_bridge_full_workflow(
        self,
        docker_compose_environment,
        gcp_monitoring_client,
        gcp_project_id,
        wait_for_telemetry,
        cleanup_collector_exports,
    ):
        """Test vendor model bridge completes full E2E workflow.

        Given: Vendor model with bridge pattern
        When: Bridge polls vendor API and exports metrics
        Then: Vendor-specific attributes are correctly captured
        """
        # Arrange
        client = MCAClient(
            service_name="e2e-vendor-bridge",
            model_id="e2e-vendor-mdl-001",
            team_name="clinical-ai",
            model_category="vendor",
            model_type="api",
            otel_endpoint=docker_compose_environment["collector_otlp"],
            otel_protocol="http",
        )

        start_time = time.time()

        # Act: Record vendor metrics with specific attributes
        expected_vendor_name = "Epic-Risk-Score"
        expected_api_version = "2.1"

        client.record_prediction(
            latency=0.25,
            prediction_id="vendor-pred-001",
            vendor_model_name=expected_vendor_name,
            vendor_api_version=expected_api_version,
            vendor_request_id="req-abc123",
        )

        client.shutdown()
        wait_for_telemetry(10, "vendor bridge telemetry")

        # Assert: Verify vendor metrics with attributes
        end_time = time.time()
        metrics = gcp_monitoring_client.query_metrics(
            project_id=gcp_project_id,
            metric_type="vendor",
            start_time=start_time,
            end_time=end_time,
        )

        assert len(metrics) > 0, "Vendor metrics not found"

        # Verify vendor-specific attributes
        found_vendor_metric = False
        for metric in metrics:
            attrs = metric.get("attributes", {})

            vendor_name = (
                attrs.get("vendor.model.name")
                or attrs.get("vendor_model_name")
                or attrs.get("vendor.name")
            )

            api_version = (
                attrs.get("vendor.api.version")
                or attrs.get("vendor_api_version")
                or attrs.get("api.version")
            )

            if vendor_name == expected_vendor_name and api_version == expected_api_version:
                found_vendor_metric = True
                break

        assert found_vendor_metric, (
            f"Vendor metric with name='{expected_vendor_name}' and version='{expected_api_version}' not found. "
            f"Available metrics: {metrics}"
        )
