"""B11: Verify that raw saturation curve arrays are stripped from MCP output.

The MCP output should contain saturation_analysis (traffic lights) but NOT
raw saturation curves (spend/response arrays). Claude cannot render charts,
so only the traffic-light summary is sent.
"""

from __future__ import annotations

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

import numpy as np
import pandas as pd
import pytest

from mixlift_mcp.license import LicenseStatus, PRO_MAX_CHANNELS, PRO_MAX_ROWS
from mixlift_mcp.server import _build_result, mixlift_analyze


def _make_mock_model_results(channel_columns: list[str]):
    """Build a mock that mimics run_full_pipeline output."""
    rng = np.random.default_rng(42)
    display_names = [
        c.replace("_spend", "").replace("_", " ").title() for c in channel_columns
    ]
    results = MagicMock()
    results.roas_estimates = {n: round(rng.uniform(0.5, 5.0), 2) for n in display_names}
    results.roas_ci = {
        n: (round(v - 0.5, 2), round(v + 0.5, 2))
        for n, v in results.roas_estimates.items()
    }
    results.channel_contributions = {
        n: round(rng.uniform(10000, 80000), 2) for n in display_names
    }
    results.channel_contributions_ci = {
        n: (round(v * 0.8, 2), round(v * 1.2, 2))
        for n, v in results.channel_contributions.items()
    }
    results.marginal_roas_estimates = {
        n: round(rng.uniform(0.3, 3.0), 2) for n in display_names
    }
    results.marginal_roas_ci = {
        n: (round(v - 0.3, 2), round(v + 0.3, 2))
        for n, v in results.marginal_roas_estimates.items()
    }
    results.total_spend = {n: round(rng.uniform(5000, 20000), 2) for n in display_names}
    results.convergence_warnings = []
    results.duration_secs = 10.0
    results.raw_contributions = MagicMock()
    results.mmm = MagicMock()
    results.loo_cv = None
    results.intercept_decomposition = None
    return results, display_names


def _make_optimization(display_names: list[str]) -> dict:
    rng = np.random.default_rng(99)
    current = {n: round(rng.uniform(3000, 10000), 2) for n in display_names}
    optimal = {n: round(v * rng.uniform(0.7, 1.4), 2) for n, v in current.items()}
    return {
        "current_allocation": current,
        "optimal_allocation": optimal,
        "total_budget": round(sum(current.values()), 2),
        "expected_lift_pct": round(rng.uniform(5, 20), 1),
    }


def _make_saturation_curves(display_names: list[str]) -> dict:
    """Raw saturation curves with spend/response arrays."""
    return {
        name: {
            "spend": list(range(0, 10001, 500)),
            "response": [float(i * 0.7) for i in range(0, 10001, 500)],
        }
        for name in display_names
    }


class TestSaturationCurveStripping:
    """Verify raw saturation curve arrays do not appear in MCP output."""

    def test_build_result_has_no_raw_curves(self):
        """_build_result should NOT include raw spend/response arrays."""
        channels = ["google_spend", "meta_spend"]
        model_results, display_names = _make_mock_model_results(channels)
        optimization = _make_optimization(display_names)
        saturation_curves = _make_saturation_curves(display_names)

        pro_status = LicenseStatus(
            is_pro=True,
            max_channels=PRO_MAX_CHANNELS,
            max_rows=PRO_MAX_ROWS,
            message="Pro",
        )

        result = _build_result(
            model_results=model_results,
            optimization=optimization,
            saturation_curves=saturation_curves,
            channel_columns=channels,
            license_status=pro_status,
            csv_format="wide",
        )

        # The result must contain saturation_analysis (traffic lights)
        assert "saturation_analysis" in result
        assert len(result["saturation_analysis"]) > 0

        # The result must NOT contain raw curve arrays
        result_json = json.dumps(result)

        # Check no key called "spend" or "response" with array values
        # exists at any nesting level
        def _find_raw_arrays(obj, path=""):
            """Recursively check for raw curve arrays."""
            violations = []
            if isinstance(obj, dict):
                for key, value in obj.items():
                    current_path = f"{path}.{key}"
                    if key in ("spend", "response") and isinstance(value, list) and len(value) > 5:
                        violations.append(
                            f"Raw curve array found at {current_path} with {len(value)} elements"
                        )
                    violations.extend(_find_raw_arrays(value, current_path))
            elif isinstance(obj, list):
                for i, item in enumerate(obj):
                    violations.extend(_find_raw_arrays(item, f"{path}[{i}]"))
            return violations

        violations = _find_raw_arrays(result)
        assert violations == [], (
            f"Raw saturation curve arrays found in MCP output: {violations}"
        )

    def test_traffic_lights_have_expected_structure(self):
        """Each traffic light entry must have saturation_pct, status, and label."""
        channels = ["google_spend", "meta_spend", "tiktok_spend"]
        model_results, display_names = _make_mock_model_results(channels)
        optimization = _make_optimization(display_names)
        saturation_curves = _make_saturation_curves(display_names)

        pro_status = LicenseStatus(True, PRO_MAX_CHANNELS, PRO_MAX_ROWS, "Pro")

        result = _build_result(
            model_results=model_results,
            optimization=optimization,
            saturation_curves=saturation_curves,
            channel_columns=channels,
            license_status=pro_status,
            csv_format="wide",
        )

        sat_analysis = result["saturation_analysis"]
        for channel_name, light in sat_analysis.items():
            assert "saturation_pct" in light, f"Missing saturation_pct for {channel_name}"
            assert "status" in light, f"Missing status for {channel_name}"
            assert light["status"] in ("green", "yellow", "red"), (
                f"Invalid status '{light['status']}' for {channel_name}"
            )
            assert "label" in light, f"Missing label for {channel_name}"

    @pytest.mark.asyncio
    async def test_full_flow_no_raw_curves_in_json_output(self, tmp_path):
        """E2E: run mixlift_analyze and verify the JSON output has no raw curves."""
        channels = ["google_spend", "meta_spend"]
        model_results, display_names = _make_mock_model_results(channels)
        optimization = _make_optimization(display_names)
        saturation_curves = _make_saturation_curves(display_names)

        pro_status = LicenseStatus(True, PRO_MAX_CHANNELS, PRO_MAX_ROWS, "Pro")

        # Create test CSV
        rng = np.random.default_rng(42)
        n = 52
        df = pd.DataFrame({
            "date": pd.date_range("2023-01-01", periods=n, freq="W-MON"),
            "google_spend": rng.uniform(1000, 5000, n),
            "meta_spend": rng.uniform(2000, 8000, n),
            "revenue": rng.uniform(30000, 80000, n),
        })
        csv_path = tmp_path / "test.csv"
        df.to_csv(csv_path, index=False)

        mem_dir = tmp_path / "memory"
        with (
            patch("mixlift_mcp.server.validate_license", new_callable=AsyncMock, return_value=pro_status),
            patch("mixlift_mcp.server._run_pipeline", return_value=model_results),
            patch("mixlift_mcp.server._run_optimizer", return_value=optimization),
            patch("mixlift_mcp.server._run_saturation", return_value=saturation_curves),
            patch("mixlift_mcp.memory.MEMORY_DIR", mem_dir),
        ):
            result_json = await mixlift_analyze(
                csv_path=str(csv_path), license_key="pro_key"
            )

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

        # Verify no raw curve arrays leaked through
        def _has_long_numeric_list(obj, min_len=5):
            if isinstance(obj, list) and len(obj) >= min_len:
                if all(isinstance(x, (int, float)) for x in obj):
                    return True
            if isinstance(obj, dict):
                for v in obj.values():
                    if _has_long_numeric_list(v, min_len):
                        return True
            if isinstance(obj, list):
                for item in obj:
                    if isinstance(item, dict) and _has_long_numeric_list(item, min_len):
                        return True
            return False

        # saturation_analysis should exist but should NOT contain numeric arrays
        sat = data.get("saturation_analysis", {})
        for ch, light in sat.items():
            assert not _has_long_numeric_list(light), (
                f"Raw numeric array found in saturation_analysis for {ch}"
            )
