"""
Experiments & Advanced Testing Tools
Wraps the AdvancedTestingEngine to provide tools for Lift Studies, A/B tests, and Bandit tests.
"""

from typing import Dict, Any, Tuple
from functools import partial
import mcp.types as mcp_types

from ..meta_client import MetaClient
from ..schemas import ToolResult
from ..advanced_testing import get_advanced_testing_engine
from .base import ToolHandler

async def run_multi_cell_test(
    client: MetaClient,
    arguments: Dict[str, Any]
) -> ToolResult:
    """Run a multi-cell lift study."""
    engine = get_advanced_testing_engine()
    engine.client = client
    return await engine.run_multi_cell_test(arguments)

async def run_sequential_test(
    client: MetaClient,
    arguments: Dict[str, Any]
) -> ToolResult:
    """Run sequential tests."""
    engine = get_advanced_testing_engine()
    engine.client = client
    return await engine.run_sequential_test(arguments)

async def run_adaptive_test(
    client: MetaClient,
    arguments: Dict[str, Any]
) -> ToolResult:
    """Run an adaptive (bandit) test."""
    engine = get_advanced_testing_engine()
    engine.client = client
    return await engine.run_adaptive_test(arguments)

async def design_test(
    client: MetaClient,
    arguments: Dict[str, Any]
) -> ToolResult:
    """Calculate sample size."""
    engine = get_advanced_testing_engine()
    engine.client = client
    return await engine.design_test(arguments)

def get_tools(client: MetaClient) -> Dict[str, Tuple[mcp_types.Tool, ToolHandler]]:
    """Register experiment tools"""
    return {
        "meta.experiments.multi_cell": (
            mcp_types.Tool(
                name="meta.experiments.multi_cell",
                description="Launch multi-cell split tests or lift studies. Compare 2+ variants with holdout groups. Use for incrementality measurement and conversion lift testing.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "test_name": {"type": "string"},
                        "variant_count": {"type": "integer", "description": "Number of cells"}
                    },
                    "required": ["test_name"]
                }
            ),
            partial(run_multi_cell_test, client)
        ),
        "meta.experiments.sequential": (
            mcp_types.Tool(
                name="meta.experiments.sequential",
                description="Launch sequential elimination tests that drop poor performers early. Efficient testing with auto-stopping rules. Use for reduced sample size testing.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "test_name": {"type": "string"}
                    },
                    "required": ["test_name"]
                }
            ),
            partial(run_sequential_test, client)
        ),
        "meta.experiments.adaptive": (
            mcp_types.Tool(
                name="meta.experiments.adaptive",
                description="Launch adaptive multi-armed bandit tests. Auto-shifts traffic to winning variants during test. Use for learning while earning optimization.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "test_name": {"type": "string"}
                    },
                    "required": ["test_name"]
                }
            ),
            partial(run_adaptive_test, client)
        ),
        "meta.experiments.design": (
            mcp_types.Tool(
                name="meta.experiments.design",
                description="Calculate sample size and statistical power for tests. Design experiments with minimum detectable effect. Use for test planning and duration estimation.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "baseline_conversion_rate": {"type": "number"},
                        "minimum_detectable_effect": {"type": "number"},
                        "variant_count": {"type": "integer"}
                    },
                    "required": ["baseline_conversion_rate"]
                }
            ),
            partial(design_test, client)
        )
    }
