"""
Sentrial Experiment Context

Provides context-aware configuration for experiments. When running experiments
via the CLI, the system prompt and other experiment config is automatically
available to your agent code via simple function calls.

Usage in your agent:

    from sentrial import get_system_prompt

    def my_agent(query: str) -> str:
        # When running normally: returns your default
        # When running experiment: returns the variant's system prompt
        system_prompt = get_system_prompt("You are a helpful assistant.")

        llm = ChatOpenAI(...)
        return llm.invoke([SystemMessage(system_prompt), HumanMessage(query)])

That's it! No other changes needed. The CLI handles setting the context.
"""

from contextvars import ContextVar
from dataclasses import dataclass
from typing import Optional, Any, Dict


@dataclass
class ExperimentContext:
    """Context for the current experiment run."""
    system_prompt: str
    variant_name: str
    experiment_id: str
    test_case_id: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None


# The context variable - holds experiment state for the current execution context
_experiment_context: ContextVar[Optional[ExperimentContext]] = ContextVar(
    'sentrial_experiment',
    default=None
)


def get_system_prompt(default: str = None) -> str:
    """
    Get the current system prompt.

    When running an experiment via `sentrial experiment run`, this returns
    the variant's system prompt. Otherwise, returns your default.

    Args:
        default: The default system prompt to use when not in experiment mode.
                 This should be your normal production prompt.

    Returns:
        The system prompt to use.

    Example:
        from sentrial import get_system_prompt

        def my_agent(query: str) -> str:
            # Just wrap your existing prompt with get_system_prompt()
            prompt = get_system_prompt("You are a helpful weather assistant.")

            llm = ChatOpenAI(...)
            return llm.invoke([SystemMessage(prompt), HumanMessage(query)])
    """
    ctx = _experiment_context.get()
    if ctx is not None and ctx.system_prompt:
        return ctx.system_prompt
    return default


def get_experiment_context() -> Optional[ExperimentContext]:
    """
    Get the full experiment context if running in experiment mode.

    Returns None if not running an experiment.

    Returns:
        ExperimentContext with variant info, or None.

    Example:
        from sentrial import get_experiment_context

        ctx = get_experiment_context()
        if ctx:
            print(f"Running variant: {ctx.variant_name}")
    """
    return _experiment_context.get()


def is_experiment_mode() -> bool:
    """
    Check if currently running in experiment mode.

    Returns:
        True if running via `sentrial experiment run`, False otherwise.
    """
    return _experiment_context.get() is not None


def get_variant_name() -> Optional[str]:
    """
    Get the current variant name if in experiment mode.

    Returns:
        The variant name (e.g., "Control", "Variant A"), or None.
    """
    ctx = _experiment_context.get()
    return ctx.variant_name if ctx else None


# === Internal functions (used by CLI) ===

def set_experiment_context(
    system_prompt: str,
    variant_name: str,
    experiment_id: str,
    test_case_id: str = None,
    metadata: Dict[str, Any] = None,
) -> None:
    """
    Set the experiment context. Called by the CLI before running agent.

    This is an internal function - users don't need to call this directly.
    """
    ctx = ExperimentContext(
        system_prompt=system_prompt,
        variant_name=variant_name,
        experiment_id=experiment_id,
        test_case_id=test_case_id,
        metadata=metadata,
    )
    _experiment_context.set(ctx)


def clear_experiment_context() -> None:
    """
    Clear the experiment context. Called by the CLI after agent completes.

    This is an internal function - users don't need to call this directly.
    """
    _experiment_context.set(None)
