"""Shared fixtures for meridian MCP server tests.

Builds a fake Meridian model with known values so every tool can be tested
deterministically without TensorFlow or a real pickle.
"""

from datetime import datetime, timezone
from types import SimpleNamespace
from unittest.mock import MagicMock

import numpy as np
import pytest
import xarray as xr

# ---------------------------------------------------------------------------
# Constants — known values for assertions
# ---------------------------------------------------------------------------

CHANNELS = ["channel_a", "channel_b", "channel_c"]
N_CHAINS = 2
N_DRAWS = 100
N_TIME = 52
N_GEO = 1
GEO = ["national"]
TIME = [f"2024-{w:02d}-01" for w in range(1, N_TIME + 1)]

# Posterior values (deterministic — same value repeated so median == mean)
ROI_VALUES = {"channel_a": 3.0, "channel_b": 1.5, "channel_c": 0.5}
ALPHA_VALUES = {"channel_a": 0.5, "channel_b": 0.25, "channel_c": 0.75}
EC_VALUES = {"channel_a": 0.5, "channel_b": 0.8, "channel_c": 0.3}
SLOPE_VALUES = {"channel_a": 1.0, "channel_b": 2.0, "channel_c": 1.5}
BETA_VALUES = {"channel_a": 1.0, "channel_b": 0.5, "channel_c": 2.0}

# Spend: average weekly spend per channel
SPEND_PER_CHANNEL = [1000.0, 2000.0, 500.0]
SCALE_FACTORS = [100.0, 200.0, 50.0]
KPI_STDEV = 10.0


def _make_posterior_var(
    values_by_channel: dict[str, float], extra_noise: float = 0.0,
) -> np.ndarray:
    """Create a (n_chains, n_draws, n_channels) array from per-channel values."""
    arr = np.zeros((N_CHAINS, N_DRAWS, len(CHANNELS)))
    rng = np.random.default_rng(42)
    for i, ch in enumerate(CHANNELS):
        base = values_by_channel[ch]
        if extra_noise > 0:
            arr[:, :, i] = base + rng.normal(0, extra_noise, (N_CHAINS, N_DRAWS))
        else:
            arr[:, :, i] = base
    return arr


def _make_beta_var(values_by_channel: dict[str, float]) -> np.ndarray:
    """Create a (n_chains, n_draws, n_geo, n_channels) array."""
    arr = np.zeros((N_CHAINS, N_DRAWS, N_GEO, len(CHANNELS)))
    for i, ch in enumerate(CHANNELS):
        arr[:, :, :, i] = values_by_channel[ch]
    return arr


def build_posterior() -> xr.Dataset:
    """Build a fake xarray posterior Dataset with known values."""
    coords = {
        "chain": np.arange(N_CHAINS),
        "draw": np.arange(N_DRAWS),
        "media_channel": CHANNELS,
        "geo": GEO,
        "time": TIME,
    }

    data_vars = {
        "roi_m": (
            ["chain", "draw", "media_channel"],
            _make_posterior_var(ROI_VALUES, extra_noise=0.1),
        ),
        "alpha_m": (["chain", "draw", "media_channel"], _make_posterior_var(ALPHA_VALUES)),
        "ec_m": (["chain", "draw", "media_channel"], _make_posterior_var(EC_VALUES)),
        "slope_m": (["chain", "draw", "media_channel"], _make_posterior_var(SLOPE_VALUES)),
        "beta_gm": (["chain", "draw", "geo", "media_channel"], _make_beta_var(BETA_VALUES)),
    }

    ds = xr.Dataset(data_vars, coords=coords)
    ds.attrs["created_at"] = "2026-01-01T00:00:00Z"
    return ds


def build_sample_stats() -> xr.Dataset:
    """Build fake sample_stats with no divergences."""
    coords = {"chain": np.arange(N_CHAINS), "draw": np.arange(N_DRAWS)}
    return xr.Dataset(
        {
            "diverging": (["chain", "draw"], np.zeros((N_CHAINS, N_DRAWS), dtype=bool)),
            "accept_ratio": (["chain", "draw"], np.full((N_CHAINS, N_DRAWS), 0.85)),
        },
        coords=coords,
    )


class FakeTensor:
    """Mimics a TensorFlow tensor with .numpy() support."""

    def __init__(self, array: np.ndarray):
        self._array = array

    def numpy(self) -> np.ndarray:
        return self._array


def build_mock_model() -> SimpleNamespace:
    """Build a complete fake Meridian model object."""
    posterior = build_posterior()
    sample_stats = build_sample_stats()
    trace = SimpleNamespace(posterior=posterior, sample_stats=sample_stats)

    # Media spend: (1, n_time, n_channels)
    spend = np.zeros((1, N_TIME, len(CHANNELS)))
    for i, s in enumerate(SPEND_PER_CHANNEL):
        spend[0, :, i] = s

    media_transformer = SimpleNamespace(
        _scale_factors_gm=FakeTensor(np.array([SCALE_FACTORS])),  # (1, 3)
    )
    media_tensors = SimpleNamespace(
        media_spend=FakeTensor(spend),
        media_transformer=media_transformer,
    )

    kpi_transformer = SimpleNamespace(
        _population_scaled_stdev=FakeTensor(np.array(KPI_STDEV)),
    )

    # Fake prior distributions (mimic TFP distribution objects)
    def _make_fake_dist(name, **kwargs):
        params = dict(kwargs)
        params.update(name=name, validate_args=False,
                      allow_nan_stats=True, dtype="float32")
        cls = type(name, (), {})
        obj = cls()
        obj.batch_shape = ()
        obj.parameters = params
        for k, v in kwargs.items():
            setattr(obj, k, v)
        return obj

    fake_roi_prior = _make_fake_dist(
        "LogNormal", loc=np.float32(0.2), scale=np.float32(0.9),
    )
    fake_alpha_prior = _make_fake_dist(
        "Uniform", low=np.float32(0.0), high=np.float32(1.0),
    )
    fake_ec_prior = _make_fake_dist(
        "TruncatedNormal",
        loc=np.float32(0.8), scale=np.float32(0.8),
        low=np.float32(0.1), high=np.float32(10.0),
    )
    fake_slope_prior = _make_fake_dist(
        "Deterministic", loc=np.float32(1.0),
    )
    fake_beta_prior = _make_fake_dist(
        "HalfNormal", scale=np.float32(5.0),
    )

    prior = SimpleNamespace(
        roi_m=fake_roi_prior,
        alpha_m=fake_alpha_prior,
        ec_m=fake_ec_prior,
        slope_m=fake_slope_prior,
        beta_m=fake_beta_prior,
        mroi_m=None,
    )

    model_spec = SimpleNamespace(
        max_lag=4,
        hill_before_adstock=False,
        media_prior_type="roi",
        media_effects_dist="log_normal",
        prior=prior,
    )
    model_context = SimpleNamespace(_model_spec=model_spec)

    model = SimpleNamespace(
        _inference_data=trace,
        _model_context=model_context,
        media_tensors=media_tensors,
        kpi_transformer=kpi_transformer,
    )

    return model


@pytest.fixture()
def mock_model(monkeypatch):
    """Patch the server's global state with a fake model.

    Returns the fake model for assertions.
    """
    import mcp_server_meridian.server as srv

    model = build_mock_model()

    monkeypatch.setattr(srv, "_model", model)
    monkeypatch.setattr(srv, "_trace", model._inference_data)
    monkeypatch.setattr(srv, "_posterior", model._inference_data.posterior)
    monkeypatch.setenv("MERIDIAN_MODEL_PATH", "/fake/model.pkl")

    return model


S3_BUCKET = "test-meridian-bucket"
S3_PREFIX = "models/"
S3_MODELS = [
    {
        "Key": "models/client_20260301.pkl",
        "Size": 50 * 1024 * 1024,
        "LastModified": datetime(2026, 3, 1, tzinfo=timezone.utc),
    },
    {
        "Key": "models/client_20260201.pkl",
        "Size": 48 * 1024 * 1024,
        "LastModified": datetime(2026, 2, 1, tzinfo=timezone.utc),
    },
    {
        "Key": "models/readme.txt",
        "Size": 1024,
        "LastModified": datetime(2026, 1, 15, tzinfo=timezone.utc),
    },
]


@pytest.fixture()
def mock_s3(monkeypatch):
    """Patch boto3 to return fake S3 listings and downloads."""
    mock_client = MagicMock()

    # Mock paginator for list_objects_v2
    mock_paginator = MagicMock()
    mock_paginator.paginate.return_value = [{"Contents": S3_MODELS}]
    mock_client.get_paginator.return_value = mock_paginator

    # Mock download_file to be a no-op
    mock_client.download_file = MagicMock()

    mock_boto3 = MagicMock()
    mock_boto3.client.return_value = mock_client

    monkeypatch.setenv("MERIDIAN_S3_BUCKET", S3_BUCKET)
    monkeypatch.setenv("MERIDIAN_S3_PREFIX", S3_PREFIX)

    import sys

    monkeypatch.setitem(sys.modules, "boto3", mock_boto3)

    return mock_client
