"""E2E integration tests for MixLift MCP server — new fields.

Verifies that `marginal_roas` and `diagnostics` are present and correctly
structured in the output of `mixlift_analyze`, while ensuring all existing
top-level fields remain intact.
"""

import json
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from mixlift_mcp.license import LicenseStatus
from mixlift_mcp.server import mixlift_analyze


# ---------------------------------------------------------------------------
# Shared fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
def mock_pro_license():
    return LicenseStatus(
        is_pro=True, max_channels=999, max_rows=999_999, message="Pro"
    )


@pytest.fixture
def convergence_warning():
    return "R-hat > 1.05 for channel Google Search"


@pytest.fixture
def mock_model_results(convergence_warning):
    """Full model-results mock including the new marginal_roas and diagnostics fields."""
    results = MagicMock()

    # --- Existing ROAS fields ---
    results.roas_estimates = {"Google Search": 3.2, "Meta": 1.8}
    results.roas_ci = {
        "Google Search": (2.1, 4.3),
        "Meta": (1.1, 2.5),
    }

    # --- Channel contributions ---
    results.channel_contributions = {"Google Search": 52000.0, "Meta": 31000.0}
    results.channel_contributions_ci = {
        "Google Search": (41000.0, 63000.0),
        "Meta": (22000.0, 40000.0),
    }

    # --- Spend / timing ---
    results.total_spend = {"Google Search": 16000.0, "Meta": 17000.0}
    results.duration_secs = 12.7
    results.mmm = MagicMock()

    # --- New: marginal ROAS ---
    results.marginal_roas_estimates = {"Google Search": 2.9, "Meta": 1.6}
    results.marginal_roas_ci = {
        "Google Search": (1.9, 3.9),
        "Meta": (1.0, 2.2),
    }

    # --- New: convergence diagnostics ---
    results.convergence_warnings = [convergence_warning]

    return results


@pytest.fixture
def mock_optimization():
    return {
        "current_allocation": {"Google Search": 8000.0, "Meta": 8500.0},
        "optimal_allocation": {"Google Search": 9500.0, "Meta": 7000.0},
        "total_budget": 16500.0,
        "expected_lift_pct": 9.3,
    }


@pytest.fixture
def mock_saturation():
    return {
        "Google Search": {"spend": [0, 5000, 10000], "response": [0, 28000, 50000]},
        "Meta": {"spend": [0, 5000, 10000], "response": [0, 16000, 29000]},
    }


# ---------------------------------------------------------------------------
# Helper: run mixlift_analyze with all heavy dependencies mocked
# ---------------------------------------------------------------------------


async def _run_analyze_with_mocks(
    mock_model_results,
    mock_optimization,
    mock_saturation,
    mock_pro_license,
) -> dict:
    """Invoke mixlift_analyze (demo data path) with all MCMC/license mocks applied."""
    with (
        patch("mixlift_mcp.server.validate_license", new_callable=AsyncMock) as mock_lic,
        patch("mixlift_mcp.server._run_pipeline") as mock_pipeline,
        patch("mixlift_mcp.server._run_optimizer") as mock_opt,
        patch("mixlift_mcp.server._run_saturation") as mock_sat,
    ):
        mock_lic.return_value = mock_pro_license
        mock_pipeline.return_value = mock_model_results
        mock_opt.return_value = mock_optimization
        mock_sat.return_value = mock_saturation

        raw = await mixlift_analyze(csv_path="", license_key="pro-key")

    return json.loads(raw)


# ---------------------------------------------------------------------------
# Integration tests
# ---------------------------------------------------------------------------


class TestMarginalRoasAndDiagnosticsPresent:
    """Verify the new marginal_roas and diagnostics top-level keys are present
    and structurally correct in the response."""

    @pytest.mark.asyncio
    async def test_status_is_success(
        self,
        mock_model_results,
        mock_optimization,
        mock_saturation,
        mock_pro_license,
    ):
        data = await _run_analyze_with_mocks(
            mock_model_results, mock_optimization, mock_saturation, mock_pro_license
        )
        assert data["status"] == "success"

    @pytest.mark.asyncio
    async def test_marginal_roas_estimates_values(
        self,
        mock_model_results,
        mock_optimization,
        mock_saturation,
        mock_pro_license,
    ):
        data = await _run_analyze_with_mocks(
            mock_model_results, mock_optimization, mock_saturation, mock_pro_license
        )

        assert "marginal_roas" in data, "top-level 'marginal_roas' key missing"
        estimates = data["marginal_roas"]["estimates"]
        assert estimates["Google Search"] == pytest.approx(2.9)
        assert estimates["Meta"] == pytest.approx(1.6)

    @pytest.mark.asyncio
    async def test_marginal_roas_confidence_intervals_structure(
        self,
        mock_model_results,
        mock_optimization,
        mock_saturation,
        mock_pro_license,
    ):
        data = await _run_analyze_with_mocks(
            mock_model_results, mock_optimization, mock_saturation, mock_pro_license
        )

        cis = data["marginal_roas"]["confidence_intervals"]

        # Google Search CI
        google_ci = cis["Google Search"]
        assert "low" in google_ci and "high" in google_ci, (
            "CI must have 'low' and 'high' keys"
        )
        assert google_ci["low"] == pytest.approx(1.9)
        assert google_ci["high"] == pytest.approx(3.9)

        # Meta CI
        meta_ci = cis["Meta"]
        assert meta_ci["low"] == pytest.approx(1.0)
        assert meta_ci["high"] == pytest.approx(2.2)

    @pytest.mark.asyncio
    async def test_diagnostics_convergence_warnings_is_list(
        self,
        mock_model_results,
        mock_optimization,
        mock_saturation,
        mock_pro_license,
    ):
        data = await _run_analyze_with_mocks(
            mock_model_results, mock_optimization, mock_saturation, mock_pro_license
        )

        assert "diagnostics" in data, "top-level 'diagnostics' key missing"
        warnings = data["diagnostics"]["convergence_warnings"]
        assert isinstance(warnings, list), "'convergence_warnings' must be a list"

    @pytest.mark.asyncio
    async def test_diagnostics_convergence_warnings_content(
        self,
        mock_model_results,
        mock_optimization,
        mock_saturation,
        mock_pro_license,
        convergence_warning,
    ):
        data = await _run_analyze_with_mocks(
            mock_model_results, mock_optimization, mock_saturation, mock_pro_license
        )

        warnings = data["diagnostics"]["convergence_warnings"]
        assert len(warnings) >= 1, "Expected at least one convergence warning"
        assert convergence_warning in warnings, (
            f"Expected warning '{convergence_warning}' not found in {warnings}"
        )

    @pytest.mark.asyncio
    async def test_existing_roas_field_still_present(
        self,
        mock_model_results,
        mock_optimization,
        mock_saturation,
        mock_pro_license,
    ):
        data = await _run_analyze_with_mocks(
            mock_model_results, mock_optimization, mock_saturation, mock_pro_license
        )

        assert "roas" in data
        assert "estimates" in data["roas"]
        assert "confidence_intervals" in data["roas"]
        assert data["roas"]["estimates"]["Google Search"] == pytest.approx(3.2)
        assert data["roas"]["estimates"]["Meta"] == pytest.approx(1.8)

    @pytest.mark.asyncio
    async def test_existing_budget_optimization_field_still_present(
        self,
        mock_model_results,
        mock_optimization,
        mock_saturation,
        mock_pro_license,
    ):
        data = await _run_analyze_with_mocks(
            mock_model_results, mock_optimization, mock_saturation, mock_pro_license
        )

        assert "budget_optimization" in data
        bo = data["budget_optimization"]
        assert "current_allocation" in bo
        assert "optimal_allocation" in bo
        assert "expected_lift_pct" in bo

    @pytest.mark.asyncio
    async def test_existing_recommendations_field_still_present(
        self,
        mock_model_results,
        mock_optimization,
        mock_saturation,
        mock_pro_license,
    ):
        data = await _run_analyze_with_mocks(
            mock_model_results, mock_optimization, mock_saturation, mock_pro_license
        )

        assert "recommendations" in data
        assert "summary" in data["recommendations"]
        assert "actions" in data["recommendations"]

    @pytest.mark.asyncio
    async def test_existing_saturation_analysis_field_still_present(
        self,
        mock_model_results,
        mock_optimization,
        mock_saturation,
        mock_pro_license,
    ):
        data = await _run_analyze_with_mocks(
            mock_model_results, mock_optimization, mock_saturation, mock_pro_license
        )

        assert "saturation_analysis" in data
        assert "Google Search" in data["saturation_analysis"]
        assert "Meta" in data["saturation_analysis"]

    @pytest.mark.asyncio
    async def test_existing_model_info_field_still_present(
        self,
        mock_model_results,
        mock_optimization,
        mock_saturation,
        mock_pro_license,
    ):
        data = await _run_analyze_with_mocks(
            mock_model_results, mock_optimization, mock_saturation, mock_pro_license
        )

        assert "model_info" in data
        assert "duration_secs" in data["model_info"]
        assert "total_spend" in data["model_info"]


class TestMarginalRoasChannelConsistency:
    """Verify that marginal_roas channels match the overall channel set."""

    @pytest.mark.asyncio
    async def test_marginal_roas_channels_match_roas_channels(
        self,
        mock_model_results,
        mock_optimization,
        mock_saturation,
        mock_pro_license,
    ):
        data = await _run_analyze_with_mocks(
            mock_model_results, mock_optimization, mock_saturation, mock_pro_license
        )

        roas_channels = set(data["roas"]["estimates"].keys())
        marginal_channels = set(data["marginal_roas"]["estimates"].keys())
        assert marginal_channels == roas_channels, (
            f"marginal_roas channels {marginal_channels} differ from "
            f"roas channels {roas_channels}"
        )

    @pytest.mark.asyncio
    async def test_marginal_roas_ci_channels_match_estimates(
        self,
        mock_model_results,
        mock_optimization,
        mock_saturation,
        mock_pro_license,
    ):
        data = await _run_analyze_with_mocks(
            mock_model_results, mock_optimization, mock_saturation, mock_pro_license
        )

        estimate_channels = set(data["marginal_roas"]["estimates"].keys())
        ci_channels = set(data["marginal_roas"]["confidence_intervals"].keys())
        assert ci_channels == estimate_channels, (
            "marginal_roas CI channels must match estimate channels"
        )

    @pytest.mark.asyncio
    async def test_marginal_roas_ci_low_below_high(
        self,
        mock_model_results,
        mock_optimization,
        mock_saturation,
        mock_pro_license,
    ):
        data = await _run_analyze_with_mocks(
            mock_model_results, mock_optimization, mock_saturation, mock_pro_license
        )

        for channel, ci in data["marginal_roas"]["confidence_intervals"].items():
            assert ci["low"] < ci["high"], (
                f"Channel '{channel}': CI low ({ci['low']}) must be < high ({ci['high']})"
            )


class TestDiagnosticsWithMultipleWarnings:
    """Verify diagnostics handles multiple convergence warnings correctly."""

    @pytest.mark.asyncio
    async def test_multiple_convergence_warnings_all_present(
        self,
        mock_optimization,
        mock_saturation,
        mock_pro_license,
    ):
        warnings = [
            "R-hat > 1.05 for channel Google Search",
            "R-hat > 1.05 for channel Meta",
            "Low effective sample size for beta[0]",
        ]

        results = MagicMock()
        results.roas_estimates = {"Google Search": 3.2, "Meta": 1.8}
        results.roas_ci = {"Google Search": (2.1, 4.3), "Meta": (1.1, 2.5)}
        results.channel_contributions = {"Google Search": 52000.0, "Meta": 31000.0}
        results.channel_contributions_ci = {
            "Google Search": (41000.0, 63000.0),
            "Meta": (22000.0, 40000.0),
        }
        results.total_spend = {"Google Search": 16000.0, "Meta": 17000.0}
        results.duration_secs = 14.2
        results.mmm = MagicMock()
        results.marginal_roas_estimates = {"Google Search": 2.9, "Meta": 1.6}
        results.marginal_roas_ci = {
            "Google Search": (1.9, 3.9),
            "Meta": (1.0, 2.2),
        }
        results.convergence_warnings = warnings

        data = await _run_analyze_with_mocks(
            results, mock_optimization, mock_saturation, mock_pro_license
        )

        returned_warnings = data["diagnostics"]["convergence_warnings"]
        assert len(returned_warnings) == 3
        for w in warnings:
            assert w in returned_warnings

    @pytest.mark.asyncio
    async def test_empty_convergence_warnings_returns_empty_list(
        self,
        mock_optimization,
        mock_saturation,
        mock_pro_license,
    ):
        results = MagicMock()
        results.roas_estimates = {"Google Search": 3.2, "Meta": 1.8}
        results.roas_ci = {"Google Search": (2.1, 4.3), "Meta": (1.1, 2.5)}
        results.channel_contributions = {"Google Search": 52000.0, "Meta": 31000.0}
        results.channel_contributions_ci = {
            "Google Search": (41000.0, 63000.0),
            "Meta": (22000.0, 40000.0),
        }
        results.total_spend = {"Google Search": 16000.0, "Meta": 17000.0}
        results.duration_secs = 8.1
        results.mmm = MagicMock()
        results.marginal_roas_estimates = {"Google Search": 2.9, "Meta": 1.6}
        results.marginal_roas_ci = {
            "Google Search": (1.9, 3.9),
            "Meta": (1.0, 2.2),
        }
        results.convergence_warnings = []

        data = await _run_analyze_with_mocks(
            results, mock_optimization, mock_saturation, mock_pro_license
        )

        warnings = data["diagnostics"]["convergence_warnings"]
        assert warnings == [], (
            f"Expected empty list for convergence_warnings, got {warnings}"
        )
