"""Smoke-test all MCP tools against the last N models in S3."""

import os
import sys
import traceback

os.environ.setdefault("MERIDIAN_S3_BUCKET", "bluealpha-mmm-meridian-object")

from mcp_server_meridian.server import (  # noqa: E402
    get_adstock_parameters,
    get_channel_priors,
    get_channel_roi,
    get_contribution_breakdown,
    get_marginal_roi,
    get_model_settings,
    get_model_summary,
    get_prior_posterior_comparison,
    get_saturation_curves,
    get_weekly_contributions,
    list_models,
    load_model,
    simulate_budget_reallocation,
)

TOOLS = [
    ("get_model_summary", lambda: get_model_summary()),
    ("get_channel_roi", lambda: get_channel_roi(credible_interval=0.8)),
    ("get_marginal_roi", lambda: get_marginal_roi()),
    ("get_contribution_breakdown", lambda: get_contribution_breakdown()),
    ("get_saturation_curves", lambda: get_saturation_curves()),
    ("get_adstock_parameters", lambda: get_adstock_parameters()),
    ("get_weekly_contributions", lambda: get_weekly_contributions(last_n_weeks=3)),
    ("get_channel_priors", lambda: get_channel_priors()),
    ("get_prior_posterior_comparison", lambda: get_prior_posterior_comparison()),
    ("get_model_settings", lambda: get_model_settings()),
]

N_MODELS = int(sys.argv[1]) if len(sys.argv) > 1 else 5


def main():
    print("=== Listing models from S3 ===")
    catalog = list_models()
    if "error" in catalog:
        print(f"ERROR: {catalog['error']}")
        return

    models = catalog["models"]
    # Sort by last_modified descending, pick last N
    models.sort(key=lambda m: m["last_modified"], reverse=True)
    selected = models[:N_MODELS]
    print(f"Found {catalog['count']} models, testing last {len(selected)}:\n")
    for m in selected:
        print(f"  {m['name']}  ({m['size_mb']} MB, {m['last_modified'][:10]})")
    print()

    for model_info in selected:
        name = model_info["name"]
        print(f"\n{'='*60}")
        print(f"MODEL: {name}")
        print(f"{'='*60}")

        try:
            summary = load_model(name)
            if "error" in summary:
                print(f"  LOAD FAILED: {summary['error']}")
                continue
            channels = summary.get("media_channels", [])
            n_weeks = summary.get("time_range", {}).get("n_weeks", "?")
            print(f"  Loaded: {len(channels)} channels, {n_weeks} weeks")
        except Exception:
            print(f"  LOAD EXCEPTION:\n{traceback.format_exc()}")
            continue

        for tool_name, tool_fn in TOOLS:
            try:
                result = tool_fn()
                has_error = "error" in result if isinstance(result, dict) else False
                status = "FAIL" if has_error else "OK"
                detail = ""
                if has_error:
                    detail = f" — {result['error'][:80]}"
                elif tool_name == "get_weekly_contributions":
                    ch0 = result.get("channels", [{}])[0] if result.get("channels") else {}
                    detail = (
                        f" — {result.get('n_weeks', '?')}wk, "
                        f"has_spend={bool(ch0.get('spend'))}, "
                        f"has_impressions={bool(ch0.get('impressions'))}, "
                        f"has_pct={bool(ch0.get('contribution_pct'))}, "
                        f"baseline={result.get('baseline', {}).get('weekly_baseline', 'N/A')}"
                    )
                print(f"  {status:4s} {tool_name}{detail}")
            except Exception as e:
                print(f"  ERR  {tool_name} — {type(e).__name__}: {e}")

        # Test simulate_budget_reallocation with first two channels
        try:
            if len(channels) >= 2:
                budgets = {channels[0]: 10000, channels[1]: 10000}
                result = simulate_budget_reallocation(budgets)
                has_error = "error" in result if isinstance(result, dict) else False
                status = "FAIL" if has_error else "OK"
                print(f"  {status:4s} simulate_budget_reallocation")
        except Exception as e:
            print(f"  ERR  simulate_budget_reallocation — {type(e).__name__}: {e}")

    print(f"\n{'='*60}")
    print("Done.")


if __name__ == "__main__":
    main()
