"""End-to-end pipeline audit tests for an 8-channel Pro-tier dataset.

Tests trace data through every stage of the MixLift pipeline:
  1. CSV loading & format detection
  2. Wide-format loading and preprocessing
  3. Tier enforcement (free vs pro)
  4. Model building (no MCMC)
  5. Result extraction with mocked posterior
  6. Budget optimization with mocked model
  7. Recommendation generation with 8 channels

These tests do NOT run actual MCMC sampling -- they mock the fitted model
to validate extraction, optimization, and recommendation logic in isolation.
"""

from __future__ import annotations

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

import numpy as np
import pandas as pd
import pytest

# ---- Paths ----

FIXTURE_DIR = Path(__file__).parent / "fixtures"
SYNTHETIC_8CH_CSV = FIXTURE_DIR / "synthetic_8ch.csv"

CHANNEL_COLUMNS = [
    "facebook_spend",
    "google_spend",
    "tiktok_spend",
    "email_spend",
    "tv_spend",
    "podcast_spend",
    "affiliate_spend",
    "seo_spend",
]

DISPLAY_NAMES = [
    "Facebook",
    "Google",
    "TikTok",
    "Email",
    "TV",
    "Podcast",
    "Affiliate",
    "SEO",
]


# ---- Fixtures ----


@pytest.fixture
def raw_df():
    """Load the raw 8-channel CSV into a DataFrame."""
    return pd.read_csv(SYNTHETIC_8CH_CSV)


@pytest.fixture
def loaded_wide(raw_df):
    """Run _load_wide_format and return (X, y, channel_columns)."""
    from mixlift_mcp.server import _load_wide_format

    return _load_wide_format(raw_df)


# ---- Stage 1: CSV Loading & Format Detection ----


class TestCsvLoadingAndFormatDetection:
    """Verify that the 8-channel CSV is correctly detected and loaded."""

    def test_fixture_exists(self):
        assert SYNTHETIC_8CH_CSV.exists(), f"Fixture not found: {SYNTHETIC_8CH_CSV}"

    def test_fixture_has_365_rows(self, raw_df):
        assert len(raw_df) == 365

    def test_fixture_has_correct_columns(self, raw_df):
        expected = {"date"} | set(CHANNEL_COLUMNS) | {"revenue"}
        assert set(raw_df.columns) == expected

    def test_no_missing_values(self, raw_df):
        assert raw_df.isna().sum().sum() == 0, "Fixture has unexpected NaN values"

    def test_all_spend_positive(self, raw_df):
        for col in CHANNEL_COLUMNS:
            assert (raw_df[col] >= 0).all(), f"{col} has negative spend"

    def test_revenue_positive(self, raw_df):
        assert (raw_df["revenue"] > 0).all()

    def test_detected_as_wide_format(self, raw_df):
        from mixlift_mcp.server import _detect_csv_format

        fmt = _detect_csv_format(raw_df)
        assert fmt == "wide", f"Expected 'wide', got '{fmt}'"

    def test_detect_format_requires_at_least_2_spend_cols(self):
        """Edge case: 1 spend col + revenue should NOT be detected as wide."""
        from mixlift_mcp.server import _detect_csv_format

        df = pd.DataFrame({
            "date": ["2024-01-01"],
            "only_spend": [100],
            "revenue": [500],
        })
        assert _detect_csv_format(df) != "wide"

    def test_detect_format_case_insensitive_revenue(self):
        """Wide detection checks lowercase; verify Revenue/REVENUE columns."""
        from mixlift_mcp.server import _detect_csv_format

        df = pd.DataFrame({
            "date": ["2024-01-01"],
            "a_spend": [1],
            "b_spend": [2],
            "Revenue": [100],
        })
        assert _detect_csv_format(df) == "wide"


# ---- Stage 2: Wide Format Loading & Preprocessing ----


class TestWideFormatLoading:
    """Verify _load_wide_format produces correct X, y, and channel_columns."""

    def test_returns_8_channels(self, loaded_wide):
        X, y, channels = loaded_wide
        assert len(channels) == 8
        assert set(channels) == set(CHANNEL_COLUMNS)

    def test_x_has_date_and_trend(self, loaded_wide):
        X, _, _ = loaded_wide
        assert "date" in X.columns
        assert "trend" in X.columns

    def test_x_has_all_spend_columns(self, loaded_wide):
        X, _, channels = loaded_wide
        for ch in channels:
            assert ch in X.columns

    def test_y_length_matches_x(self, loaded_wide):
        X, y, _ = loaded_wide
        assert len(X) == len(y)

    def test_y_is_float(self, loaded_wide):
        _, y, _ = loaded_wide
        assert y.dtype == np.float64 or y.dtype == float

    def test_date_column_is_datetime(self, loaded_wide):
        X, _, _ = loaded_wide
        assert pd.api.types.is_datetime64_any_dtype(X["date"])

    def test_trend_is_sequential(self, loaded_wide):
        X, _, _ = loaded_wide
        expected = np.arange(len(X), dtype=float)
        np.testing.assert_array_equal(X["trend"].values, expected)

    def test_add_control_columns_adds_trend_and_holiday(self):
        """Verify add_control_columns adds 'trend' and 'is_holiday_week'."""
        from mixlift.modeling.preprocessing import add_control_columns

        df = pd.DataFrame({
            "date": pd.date_range("2024-01-01", periods=10),
            "x_spend": range(10),
        })
        result = add_control_columns(df)
        new_cols = set(result.columns) - set(df.columns)
        assert new_cols == {"trend", "is_holiday_week", "month_sin", "month_cos", "is_q4"}, f"Unexpected new columns: {new_cols}"


# ---- Stage 3: Tier Enforcement ----


class TestTierEnforcement:
    """Verify free/pro tier limits work correctly with 8 channels."""

    def test_free_tier_rejects_8_channels(self, loaded_wide):
        from mixlift_mcp.license import LicenseStatus
        from mixlift_mcp.server import _enforce_tier_limits

        free = LicenseStatus(is_pro=False, max_channels=2, max_rows=2000, message="Free")
        X, _, channels = loaded_wide

        with pytest.raises(ValueError, match="channels"):
            _enforce_tier_limits(free, n_channels=len(channels), n_rows=len(X))

    def test_free_tier_rejects_365_rows(self, loaded_wide):
        """Free tier max_rows=2000 but 365 daily rows pass. Verify that's OK."""
        from mixlift_mcp.license import LicenseStatus
        from mixlift_mcp.server import _enforce_tier_limits

        free = LicenseStatus(is_pro=False, max_channels=10, max_rows=2000, message="Free")
        X, _, channels = loaded_wide
        # 365 rows < 2000, so this should NOT raise on rows
        _enforce_tier_limits(free, n_channels=len(channels), n_rows=len(X))

    def test_pro_tier_accepts_8_channels(self, loaded_wide):
        from mixlift_mcp.license import LicenseStatus
        from mixlift_mcp.server import _enforce_tier_limits

        pro = LicenseStatus(is_pro=True, max_channels=999, max_rows=999_999, message="Pro")
        X, _, channels = loaded_wide
        # Should not raise
        _enforce_tier_limits(pro, n_channels=len(channels), n_rows=len(X))

    def test_error_message_includes_upgrade_link(self, loaded_wide):
        from mixlift_mcp.license import LicenseStatus
        from mixlift_mcp.server import _enforce_tier_limits

        free = LicenseStatus(is_pro=False, max_channels=2, max_rows=2000, message="Free")
        X, _, channels = loaded_wide
        with pytest.raises(ValueError, match="mixlift.io/pricing"):
            _enforce_tier_limits(free, n_channels=len(channels), n_rows=len(X))

    def test_error_message_uses_correct_tier_name(self):
        """FIX: _enforce_tier_limits now uses dynamic tier name instead of
        hardcoded 'Free tier'."""
        from mixlift_mcp.license import LicenseStatus
        from mixlift_mcp.server import _enforce_tier_limits

        # Free tier should say "Free tier"
        free = LicenseStatus(is_pro=False, max_channels=5, max_rows=5000, message="Free")
        with pytest.raises(ValueError, match="Free tier"):
            _enforce_tier_limits(free, n_channels=8, n_rows=100)

        # Pro tier should say "Pro tier"
        pro = LicenseStatus(is_pro=True, max_channels=5, max_rows=5000, message="Pro")
        with pytest.raises(ValueError, match="Pro tier"):
            _enforce_tier_limits(pro, n_channels=8, n_rows=100)


# ---- Stage 4: Model Building ----


class TestModelBuilding:
    """Verify build_model works with 8 channels (no MCMC)."""

    def test_build_model_8_channels(self):
        from mixlift.modeling.engine import build_model

        mmm = build_model(CHANNEL_COLUMNS)
        assert mmm.channel_columns == CHANNEL_COLUMNS

    def test_build_model_default_seasonality(self):
        from mixlift.modeling.engine import build_model

        mmm = build_model(CHANNEL_COLUMNS)
        assert mmm.yearly_seasonality == 6

    def test_build_model_custom_adstock_lag(self):
        from mixlift.modeling.engine import build_model

        mmm = build_model(CHANNEL_COLUMNS, adstock_l_max=12)
        assert mmm.adstock.l_max == 12


# ---- Stage 5: Result Extraction with Mocked Posterior ----


def _make_mock_mmm(channel_columns, n_chains=2, n_draws=100, n_obs=52):
    """Create a MagicMock that mimics a fitted MMM with realistic shapes."""
    import xarray as xr

    n_channels = len(channel_columns)
    rng = np.random.default_rng(42)

    # Mock posterior
    posterior = xr.Dataset({
        "saturation_lam": xr.DataArray(
            rng.uniform(0.0001, 0.001, (n_chains, n_draws, n_channels)),
            dims=["chain", "draw", "channel"],
        ),
        "saturation_beta": xr.DataArray(
            rng.uniform(5000, 20000, (n_chains, n_draws, n_channels)),
            dims=["chain", "draw", "channel"],
        ),
        "adstock_alpha": xr.DataArray(
            rng.uniform(0.1, 0.8, (n_chains, n_draws, n_channels)),
            dims=["chain", "draw", "channel"],
        ),
    })

    # Mock idata
    idata = MagicMock()
    idata.posterior = posterior

    # Mock sample_stats with no divergences
    sample_stats = xr.Dataset({
        "diverging": xr.DataArray(
            np.zeros((n_chains, n_draws), dtype=bool),
            dims=["chain", "draw"],
        ),
    })
    idata.sample_stats = sample_stats

    # Mock compute_channel_contribution_original_scale
    contributions = xr.DataArray(
        rng.uniform(100, 5000, (n_chains, n_draws, n_obs, n_channels)),
        dims=["chain", "draw", "date", "channel"],
        coords={"channel": channel_columns},
    )

    # Mock channel_contribution_forward_pass: (n_chains*n_draws, n_obs, n_channels)
    def mock_forward_pass(channel_data):
        n_samples = n_chains * n_draws
        n_t = channel_data.shape[0]
        n_ch = channel_data.shape[1]
        return rng.uniform(10, 500, (n_samples, n_t, n_ch))

    mmm = MagicMock()
    mmm.idata = idata
    mmm.channel_columns = channel_columns
    mmm.compute_channel_contribution_original_scale.return_value = contributions
    mmm.channel_contribution_forward_pass = mock_forward_pass
    mmm.adstock = MagicMock()
    mmm.adstock.l_max = 8
    mmm.adstock.normalize = True

    return mmm


class TestResultExtraction:
    """Verify extraction functions work with 8 channels."""

    def test_extract_channel_contributions(self, loaded_wide):
        from mixlift.modeling.engine import extract_channel_contributions

        X, _, channels = loaded_wide
        mmm = _make_mock_mmm(channels, n_obs=len(X))

        means, cis = extract_channel_contributions(mmm, X, channels)

        assert len(means) == 8
        assert len(cis) == 8
        for name in DISPLAY_NAMES:
            assert name in means, f"Missing channel: {name}"
            assert name in cis
            low, high = cis[name]
            assert low <= means[name] <= high or np.isclose(low, high)

    def test_extract_roas(self, loaded_wide):
        from mixlift.modeling.engine import extract_roas

        X, _, channels = loaded_wide
        mmm = _make_mock_mmm(channels, n_obs=len(X))

        roas, roas_ci = extract_roas(mmm, X, channels)

        assert len(roas) == 8
        for name in DISPLAY_NAMES:
            assert name in roas
            # ROAS should be a finite number
            assert np.isfinite(roas[name])

    def test_extract_roas_zero_spend_channel(self, loaded_wide):
        """Edge case: channel with all-zero spend should get ROAS=0."""
        from mixlift.modeling.engine import extract_roas

        X, _, channels = loaded_wide
        X = X.copy()
        X["seo_spend"] = 0.0  # zero out one channel
        mmm = _make_mock_mmm(channels, n_obs=len(X))

        roas, roas_ci = extract_roas(mmm, X, channels)
        assert roas["SEO"] == 0.0
        assert roas_ci["SEO"] == (0.0, 0.0)

    def test_extract_marginal_roas(self, loaded_wide):
        from mixlift.modeling.engine import extract_marginal_roas

        X, _, channels = loaded_wide
        mmm = _make_mock_mmm(channels, n_obs=len(X))

        mroas, mroas_ci = extract_marginal_roas(mmm, X, channels)

        assert len(mroas) == 8
        for name in DISPLAY_NAMES:
            assert name in mroas
            assert np.isfinite(mroas[name])

    def test_extract_saturation_curves(self, loaded_wide):
        from mixlift.modeling.engine import extract_saturation_curves

        X, _, channels = loaded_wide
        mmm = _make_mock_mmm(channels, n_obs=len(X))

        curves = extract_saturation_curves(mmm, X, channels)

        assert len(curves) == 8
        for name in DISPLAY_NAMES:
            assert name in curves
            assert "spend" in curves[name]
            assert "response" in curves[name]
            assert len(curves[name]["spend"]) == 100
            assert len(curves[name]["response"]) == 100
            # Spend should start at 0
            assert curves[name]["spend"][0] == 0.0

    def test_display_name_conversion(self):
        """Verify the display name mapping is consistent across all functions."""
        from mixlift.modeling.optimizer import format_channel_name

        for col in CHANNEL_COLUMNS:
            display = format_channel_name(col)
            assert display in DISPLAY_NAMES, f"Unexpected display name: {display}"

    def test_check_convergence_no_warnings(self, loaded_wide):
        from mixlift.modeling.engine import check_convergence

        X, _, channels = loaded_wide
        mmm = _make_mock_mmm(channels, n_obs=len(X))

        # Need to mock az.rhat to return good values
        import xarray as xr

        mock_rhat = xr.Dataset({
            "saturation_lam": xr.DataArray(np.ones(8) * 1.01),
            "saturation_beta": xr.DataArray(np.ones(8) * 1.02),
        })

        with patch("mixlift.modeling.engine.az.rhat", return_value=mock_rhat):
            warnings = check_convergence(mmm)

        assert len(warnings) == 0

    def test_check_convergence_rhat_warning(self, loaded_wide):
        from mixlift.modeling.engine import check_convergence

        X, _, channels = loaded_wide
        mmm = _make_mock_mmm(channels, n_obs=len(X))

        import xarray as xr

        mock_rhat = xr.Dataset({
            "saturation_lam": xr.DataArray(np.ones(8) * 1.10),  # > 1.05
        })

        with patch("mixlift.modeling.engine.az.rhat", return_value=mock_rhat):
            warnings = check_convergence(mmm)

        assert len(warnings) >= 1
        assert "R-hat" in warnings[0]


# ---- Stage 6: Budget Optimization ----


class TestBudgetOptimization:
    """Verify optimize_budget works with 8 channels."""

    def test_optimize_budget_returns_all_channels(self, loaded_wide):
        from mixlift.modeling.optimizer import optimize_budget

        X, _, channels = loaded_wide
        mmm = _make_mock_mmm(channels, n_obs=len(X))

        # Force fallback to gradient optimizer
        mmm.optimize_budget = MagicMock(side_effect=AttributeError)

        result = optimize_budget(mmm, X, channels)

        assert "current_allocation" in result
        assert "optimal_allocation" in result
        assert "total_budget" in result
        assert "expected_lift_pct" in result

        assert len(result["current_allocation"]) == 8
        assert len(result["optimal_allocation"]) == 8

        for name in DISPLAY_NAMES:
            assert name in result["current_allocation"]
            assert name in result["optimal_allocation"]

    def test_optimal_allocation_sums_to_total_budget(self, loaded_wide):
        from mixlift.modeling.optimizer import optimize_budget

        X, _, channels = loaded_wide
        mmm = _make_mock_mmm(channels, n_obs=len(X))
        mmm.optimize_budget = MagicMock(side_effect=AttributeError)

        result = optimize_budget(mmm, X, channels)

        total = sum(result["optimal_allocation"].values())
        expected = result["total_budget"]
        assert abs(total - expected) < 1.0, (
            f"Optimal allocation sums to {total}, expected {expected}"
        )

    def test_current_allocation_is_weekly_average(self, loaded_wide):
        """Verify current allocation = total_spend / n_weeks."""
        from mixlift.modeling.optimizer import optimize_budget

        X, _, channels = loaded_wide
        mmm = _make_mock_mmm(channels, n_obs=len(X))
        mmm.optimize_budget = MagicMock(side_effect=AttributeError)

        result = optimize_budget(mmm, X, channels)
        n_weeks = len(X)

        from mixlift.modeling.optimizer import format_channel_name

        for col in channels:
            display = format_channel_name(col)
            expected = float(X[col].sum() / n_weeks)
            actual = result["current_allocation"][display]
            assert abs(actual - expected) < 0.01, (
                f"{display}: expected {expected}, got {actual}"
            )

    def test_gradient_optimizer_n_steps(self, loaded_wide):
        """The gradient optimizer runs 100 steps by default.
        With 8 channels, that's 8*2=16 forward passes per step = 1600 total.
        This should still be fast since they're numpy-only."""
        from mixlift.modeling.optimizer import _optimize_gradient

        X, _, channels = loaded_wide
        mmm = _make_mock_mmm(channels, n_obs=len(X))
        total_budget = sum(float(X[c].sum() / len(X)) for c in channels)

        result = _optimize_gradient(mmm, X, channels, total_budget, n_steps=10)
        assert len(result) == 8


# ---- Stage 7: Recommendations ----


class TestRecommendations:
    """Verify recommendation generation with 8 channels."""

    def _make_allocations(self):
        """Create realistic current/optimal allocations for 8 channels."""
        current = {
            "Facebook": 5000.0,
            "Google": 8000.0,
            "TikTok": 2000.0,
            "Email": 500.0,
            "TV": 10000.0,
            "Podcast": 1500.0,
            "Affiliate": 3000.0,
            "SEO": 800.0,
        }
        # Optimal: increase high-ROAS channels, decrease low-ROAS
        optimal = {
            "Facebook": 4000.0,   # decrease
            "Google": 10000.0,    # increase
            "TikTok": 1500.0,     # decrease
            "Email": 2000.0,      # big increase
            "TV": 5000.0,         # big decrease
            "Podcast": 1500.0,    # no change (within 5%)
            "Affiliate": 5000.0,  # increase
            "SEO": 1800.0,        # increase
        }
        roas = {
            "Facebook": 2.5,
            "Google": 3.0,
            "TikTok": 1.2,
            "Email": 10.0,
            "TV": 0.8,
            "Podcast": 2.0,
            "Affiliate": 4.0,
            "SEO": 12.0,
        }
        return current, optimal, roas

    def test_generate_recommendations_returns_ranked_list(self):
        from mixlift.modeling.recommendations import generate_recommendations

        current, optimal, roas = self._make_allocations()
        recs, summary = generate_recommendations(current, optimal, roas, 15.0)

        assert isinstance(recs, list)
        assert len(recs) > 0
        # Ranks should be sequential from 1
        for i, rec in enumerate(recs):
            assert rec.rank == i + 1

    def test_skips_changes_under_5_pct(self):
        """Podcast stays at 1500 in both, so it should be skipped."""
        from mixlift.modeling.recommendations import generate_recommendations

        current, optimal, roas = self._make_allocations()
        recs, _ = generate_recommendations(current, optimal, roas, 15.0)

        channels_in_recs = {r.channel for r in recs}
        assert "Podcast" not in channels_in_recs

    def test_increase_and_decrease_actions(self):
        from mixlift.modeling.recommendations import generate_recommendations

        current, optimal, roas = self._make_allocations()
        recs, _ = generate_recommendations(current, optimal, roas, 15.0)

        actions = {r.channel: r.action for r in recs}
        assert actions["Google"] == "INCREASE"
        assert actions["Email"] == "INCREASE"
        assert actions["TV"] == "DECREASE"

    def test_summary_mentions_lift(self):
        from mixlift.modeling.recommendations import generate_recommendations

        current, optimal, roas = self._make_allocations()
        _, summary = generate_recommendations(current, optimal, roas, 15.0)

        assert "15.0%" in summary

    def test_all_8_channels_have_roas_in_recommendations(self):
        """Each recommendation that exists should have a non-negative ROAS."""
        from mixlift.modeling.recommendations import generate_recommendations

        current, optimal, roas = self._make_allocations()
        recs, _ = generate_recommendations(current, optimal, roas, 15.0)

        for rec in recs:
            assert rec.roas >= 0

    def test_format_all_recommendations(self):
        from mixlift.modeling.recommendations import (
            format_all_recommendations,
            generate_recommendations,
        )

        current, optimal, roas = self._make_allocations()
        recs, summary = generate_recommendations(current, optimal, roas, 15.0)
        formatted = format_all_recommendations(recs, summary)

        assert "BUDGET OPTIMIZATION RECOMMENDATIONS" in formatted
        assert "INCREASE" in formatted or "DECREASE" in formatted

    def test_empty_recommendations_when_allocations_match(self):
        """If current == optimal, no recommendations should be generated."""
        from mixlift.modeling.recommendations import generate_recommendations

        current = {name: 1000.0 for name in DISPLAY_NAMES}
        roas = {name: 2.0 for name in DISPLAY_NAMES}
        recs, summary = generate_recommendations(current, current, roas, 0.0)

        assert len(recs) == 0
        assert "close to optimal" in summary.lower()


# ---- Stage 8: Full MCP Tool Integration (with mocks) ----


class TestMcpToolIntegration:
    """Test the full mixlift_analyze flow with 8-channel data and mocked modeling."""

    def _mock_model_results(self):
        mock = MagicMock()
        mock.roas_estimates = {name: 2.0 + i * 0.5 for i, name in enumerate(DISPLAY_NAMES)}
        mock.roas_ci = {
            name: (1.0 + i * 0.3, 3.0 + i * 0.5)
            for i, name in enumerate(DISPLAY_NAMES)
        }
        mock.channel_contributions = {
            name: 10000 + i * 5000 for i, name in enumerate(DISPLAY_NAMES)
        }
        mock.channel_contributions_ci = {
            name: (5000 + i * 3000, 15000 + i * 7000)
            for i, name in enumerate(DISPLAY_NAMES)
        }
        mock.marginal_roas_estimates = {
            name: 1.5 + i * 0.3 for i, name in enumerate(DISPLAY_NAMES)
        }
        mock.marginal_roas_ci = {
            name: (0.8 + i * 0.2, 2.5 + i * 0.4)
            for i, name in enumerate(DISPLAY_NAMES)
        }
        mock.total_spend = {
            name: 50000 + i * 10000 for i, name in enumerate(DISPLAY_NAMES)
        }
        mock.duration_secs = 42.0
        mock.mmm = MagicMock()
        mock.convergence_warnings = []
        return mock

    def _mock_optimization(self):
        return {
            "current_allocation": {name: 5000.0 + i * 1000 for i, name in enumerate(DISPLAY_NAMES)},
            "optimal_allocation": {name: 4000.0 + i * 1200 for i, name in enumerate(DISPLAY_NAMES)},
            "total_budget": sum(5000.0 + i * 1000 for i in range(8)),
            "expected_lift_pct": 12.3,
        }

    def _mock_saturation(self):
        return {
            name: {"spend": [0.0, 100.0, 200.0], "response": [0.0, 50.0, 80.0]}
            for name in DISPLAY_NAMES
        }

    @pytest.mark.asyncio
    async def test_full_pipeline_success(self):
        from mixlift_mcp.license import LicenseStatus
        from mixlift_mcp.server import mixlift_analyze

        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 = LicenseStatus(True, 999, 999_999, "Pro")
            mock_pipeline.return_value = self._mock_model_results()
            mock_opt.return_value = self._mock_optimization()
            mock_sat.return_value = self._mock_saturation()

            result = await mixlift_analyze(
                csv_path=str(SYNTHETIC_8CH_CSV), license_key="pro-key"
            )

        data = json.loads(result)
        assert data["status"] == "success"

        # Verify all sections present
        assert data["data"]["n_channels"] == 8
        assert data["data"]["format_detected"] == "wide"
        assert len(data["roas"]["estimates"]) == 8
        assert len(data["channel_contributions"]["estimates"]) == 8
        assert len(data["budget_optimization"]["current_allocation"]) == 8
        assert len(data["saturation_analysis"]) == 8
        assert len(data["marginal_roas"]["estimates"]) == 8
        assert len(data["recommendations"]["actions"]) > 0
        assert data["license"]["tier"] == "pro"

    @pytest.mark.asyncio
    async def test_free_tier_rejects_8ch_csv(self):
        from mixlift_mcp.license import LicenseStatus
        from mixlift_mcp.server import mixlift_analyze

        with patch("mixlift_mcp.server.validate_license", new_callable=AsyncMock) as mock_lic:
            mock_lic.return_value = LicenseStatus(False, 2, 2000, "Free")
            result = await mixlift_analyze(csv_path=str(SYNTHETIC_8CH_CSV))

        data = json.loads(result)
        assert data["status"] == "error"
        assert "8" in data["message"]
        assert "channels" in data["message"].lower()

    @pytest.mark.asyncio
    async def test_result_json_serializable(self):
        """Verify the full result dict is JSON-serializable (no numpy types leak)."""
        from mixlift_mcp.license import LicenseStatus
        from mixlift_mcp.server import mixlift_analyze

        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 = LicenseStatus(True, 999, 999_999, "Pro")
            mock_pipeline.return_value = self._mock_model_results()
            mock_opt.return_value = self._mock_optimization()
            mock_sat.return_value = self._mock_saturation()

            result = await mixlift_analyze(
                csv_path=str(SYNTHETIC_8CH_CSV), license_key="pro"
            )

        # This should not raise
        data = json.loads(result)
        # Re-serialize to ensure no hidden non-serializable types
        json.dumps(data)


# ---- Edge Case & Bug-Finding Tests ----


class TestEdgeCases:
    """Tests that probe edge cases and potential bugs in the pipeline."""

    def test_logistic_saturation_at_zero(self):
        """logistic_saturation(0) should be 0."""
        from mixlift.modeling.engine import logistic_saturation

        result = logistic_saturation(np.array([0.0]), lam=0.001)
        assert result[0] == pytest.approx(0.0, abs=1e-10)

    def test_logistic_saturation_monotonic(self):
        """Saturation curve should be monotonically increasing for positive spend."""
        from mixlift.modeling.engine import logistic_saturation

        spend = np.linspace(0, 10000, 1000)
        response = logistic_saturation(spend, lam=0.001)
        diffs = np.diff(response)
        assert (diffs >= 0).all(), "Saturation curve is not monotonically increasing"

    def test_logistic_saturation_clamped_at_extreme_spend(self):
        """FIX: Saturation exponent is now clamped at 30 to prevent exact 1.0,
        which preserves a small gradient for the optimizer."""
        from mixlift.modeling.engine import logistic_saturation

        spend = np.array([1e6])  # very large spend
        response = logistic_saturation(spend, lam=0.001)
        # With clamping, result should be very close to 1.0 but not exactly 1.0
        assert response[0] > 0.999
        assert response[0] < 1.0  # Not exactly 1.0 due to clamping

        # At moderate spend, behavior unchanged
        moderate = np.array([1000.0])
        resp_mod = logistic_saturation(moderate, lam=0.001)
        assert 0 < resp_mod[0] < 1.0

    def test_geometric_adstock_preserves_length(self):
        from mixlift.modeling.engine import geometric_adstock_np

        x = np.random.rand(52)
        result = geometric_adstock_np(x, alpha=0.5, l_max=8)
        assert len(result) == len(x)

    def test_geometric_adstock_normalized_sum(self):
        """With normalize=True, a constant input should map to itself."""
        from mixlift.modeling.engine import geometric_adstock_np

        x = np.ones(100) * 1000.0
        result = geometric_adstock_np(x, alpha=0.5, l_max=8, normalize=True)
        # After warmup period, output should converge to input value
        assert result[-1] == pytest.approx(1000.0, rel=0.01)

    def test_wide_format_target_column_priority(self):
        """Verify revenue > target > sales priority ordering."""
        from mixlift_mcp.server import _load_wide_format

        n = 200  # Need ≥180 days to pass validation
        df = pd.DataFrame({
            "date": pd.date_range("2024-01-01", periods=n),
            "a_spend": range(n),
            "b_spend": range(n),
            "revenue": [100] * n,
            "target": [200] * n,
            "sales": [300] * n,
        })
        _, y, _ = _load_wide_format(df)
        assert y.iloc[0] == 100.0  # should use 'revenue', not 'target' or 'sales'

    def test_channel_display_name_consistency(self):
        """Verify display names are used consistently across all pipeline outputs.

        BUG PROBE: If channel columns contain underscores beyond the pattern
        X_spend, the .title() conversion might produce unexpected names.
        E.g. 'google_search_spend' -> 'Google Search', not 'Google_Search'.
        """
        test_col = "google_search_spend"
        display = test_col.replace("_spend", "").replace("_", " ").title()
        assert display == "Google Search"

        # Our 8-channel data has simple names, so this works fine:
        for col in CHANNEL_COLUMNS:
            display = col.replace("_spend", "").replace("_", " ").title()
            # No unexpected spaces or capitalization
            assert " " not in display or display.istitle()

    def test_validate_data_requires_90_days(self):
        """The canonical-format validator requires >= 90 days.
        Our 365-day CSV in wide format bypasses this since wide format
        skips validate_data(). This is a potential gap."""
        from mixlift.ingest.service import validate_data

        short_df = pd.DataFrame({
            "date": pd.to_datetime(pd.date_range("2024-01-01", periods=60)),
            "channel": ["test"] * 60,
            "spend": [100.0] * 60,
            "impressions": [1000] * 60,
            "clicks": [50] * 60,
            "conversions": [5] * 60,
        })
        errors = validate_data(short_df)
        assert any("90" in e for e in errors)

    def test_wide_format_now_validates_short_date_range(self):
        """FIX: Wide-format CSVs now run validation (was previously skipped).
        A wide CSV with < 90 days should raise ValueError."""
        from mixlift_mcp.server import _load_wide_format

        short_df = pd.DataFrame({
            "date": pd.date_range("2024-01-01", periods=30),
            "a_spend": range(30),
            "b_spend": range(30),
            "revenue": [100] * 30,
        })
        with pytest.raises(ValueError, match="90"):
            _load_wide_format(short_df)

    def test_negative_spend_in_wide_format_rejected(self):
        """FIX: Wide format now rejects negative spend."""
        from mixlift_mcp.server import _load_wide_format

        df = pd.DataFrame({
            "date": pd.date_range("2024-01-01", periods=200),
            "a_spend": [-100] * 200,  # Negative!
            "b_spend": range(200),
            "revenue": [100] * 200,
        })
        with pytest.raises(ValueError, match="negative"):
            _load_wide_format(df)

    def test_nan_spend_in_wide_format_rejected(self):
        """FIX: Wide format now rejects NaN spend."""
        from mixlift_mcp.server import _load_wide_format

        df = pd.DataFrame({
            "date": pd.date_range("2024-01-01", periods=200),
            "a_spend": [float("nan")] * 200,
            "b_spend": range(200),
            "revenue": [100] * 200,
        })
        with pytest.raises(ValueError, match="NaN"):
            _load_wide_format(df)

    def test_wide_format_case_insensitive_spend_columns(self):
        """FIX: Wide format now loads columns case-insensitively (e.g. Facebook_Spend)."""
        from mixlift_mcp.server import _load_wide_format

        df = pd.DataFrame({
            "date": pd.date_range("2024-01-01", periods=200),
            "Facebook_Spend": range(200),
            "Google_Spend": range(200),
            "revenue": [100] * 200,
        })
        X, y, channels = _load_wide_format(df)
        assert len(channels) == 2
        assert "Facebook_Spend" in channels
        assert "Google_Spend" in channels

    def test_enforce_tier_limits_returns_warnings_for_many_channels(self):
        """FIX: _enforce_tier_limits now returns warnings for >15 channels."""
        from mixlift_mcp.license import LicenseStatus
        from mixlift_mcp.server import _enforce_tier_limits

        pro = LicenseStatus(is_pro=True, max_channels=999, max_rows=999_999, message="Pro")
        warnings = _enforce_tier_limits(pro, n_channels=20, n_rows=365)
        assert len(warnings) == 1
        assert "15 channels" in warnings[0]

    def test_enforce_tier_limits_no_warnings_for_few_channels(self):
        """No warnings for ≤15 channels."""
        from mixlift_mcp.license import LicenseStatus
        from mixlift_mcp.server import _enforce_tier_limits

        pro = LicenseStatus(is_pro=True, max_channels=999, max_rows=999_999, message="Pro")
        warnings = _enforce_tier_limits(pro, n_channels=8, n_rows=365)
        assert len(warnings) == 0

    def test_contributions_computed_once_in_pipeline(self, loaded_wide):
        """FIX: compute_channel_contribution_original_scale should be called
        exactly once in run_full_pipeline (not once per extraction function)."""
        from mixlift.modeling.engine import run_full_pipeline

        X, y, channels = loaded_wide
        mmm = _make_mock_mmm(channels, n_obs=len(X))

        # Patch build_model and fit_model to return our mock
        with (
            patch("mixlift.modeling.engine.build_model", return_value=mmm),
            patch("mixlift.modeling.engine.fit_model", return_value=(mmm, "")),
        ):
            result = run_full_pipeline(X, y, channels)

        # compute_channel_contribution_original_scale called exactly once
        assert mmm.compute_channel_contribution_original_scale.call_count == 1
        # raw_contributions should be stored on the result
        assert result.raw_contributions is not None

    def test_saturation_clamp_preserves_gradient(self):
        """FIX: Clamped saturation at extreme spend should not be exactly 1.0."""
        from mixlift.modeling.engine import logistic_saturation

        extreme = np.array([1e10])
        result = logistic_saturation(extreme, lam=1.0)
        assert result[0] < 1.0, "Clamped saturation should not be exactly 1.0"
        assert result[0] > 0.999, "Clamped saturation should still be close to 1.0"

    def test_rhat_skipped_with_single_chain(self):
        """FIX: check_convergence should skip R-hat with 1 chain."""
        from unittest.mock import MagicMock

        import xarray as xr

        from mixlift.modeling.engine import check_convergence

        mmm = MagicMock()
        rng = np.random.default_rng(42)
        # 1 chain, 100 draws
        posterior = xr.Dataset({
            "param_a": xr.DataArray(rng.normal(1.0, 0.1, (1, 100)), dims=["chain", "draw"])
        })
        sample_stats = xr.Dataset({
            "diverging": xr.DataArray(np.zeros((1, 100), dtype=bool), dims=["chain", "draw"])
        })
        import arviz as az

        mmm.idata = az.InferenceData(posterior=posterior, sample_stats=sample_stats)

        warnings = check_convergence(mmm)
        # Should have a warning about R-hat not being computed
        rhat_skip = [w for w in warnings if "R-hat not computed" in w]
        assert len(rhat_skip) == 1
        # Should NOT have attempted R-hat computation (no "R-hat for" warnings)
        rhat_bad = [w for w in warnings if "R-hat for" in w]
        assert len(rhat_bad) == 0
