"""
Runtime overrides for Kivot Agents execution.

Provides context-based configuration for KivotRunner.
"""

import contextvars
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import timedelta
from typing import Optional, Sequence

from temporalio.common import RetryPolicy

# Context-var storage for current config (works with asyncio tasks)
_config_var: contextvars.ContextVar[Optional["KivotAgentsConfig"]] = contextvars.ContextVar(
    "kivot_config", default=None
)


@dataclass
class ActivityParameters:
    """Configuration for activity execution."""

    start_to_close_timeout: Optional[timedelta] = None
    schedule_to_close_timeout: Optional[timedelta] = None
    retry_policy: Optional[RetryPolicy] = None


@dataclass
class TruncationConfig:
    """Configuration for context window truncation.

    Controls how message history is managed when it grows beyond
    model token limits. When thresholds are exceeded, older messages
    are removed and optionally summarized via an LLM call.

    Attributes:
        max_messages: Maximum non-system messages before truncation triggers.
            0 means unlimited.
        max_tokens: Estimated token budget for the context window.
            Uses approximate 4-chars-per-token heuristic. 0 means unlimited.
        summarize_removed: Whether to summarize removed messages via LLM.
        summary_model_ref: Model reference for summarization (resolved via
            ModelRegistry on Go side). Required when summarize_removed is True.
        summary_max_tokens: Maximum tokens for the generated summary.
        preserve_system_messages: When True, system messages are never truncated.
    """

    max_messages: int = 100
    max_tokens: int = 0
    summarize_removed: bool = False
    summary_model_ref: str = ""
    summary_max_tokens: int = 500
    preserve_system_messages: bool = True


@dataclass
class KivotAgentsConfig:
    """Runtime configuration for Kivot Agents."""

    model_params: ActivityParameters
    tool_params: ActivityParameters
    thread_params: Optional[ActivityParameters] = None
    go_activities_task_queue: str = "core-activities"
    default_policy_refs: Sequence[str] = ()
    thread_history_limit: int = 50
    truncation: Optional[TruncationConfig] = None


@contextmanager
def set_kivot_agent_overrides(config: KivotAgentsConfig):
    """
    Set runtime overrides for Kivot Agents execution.

    Similar to temporalio.contrib.openai_agents.set_open_ai_agent_temporal_overrides.

    This context manager configures KivotRunner behavior including:
    - Activity timeouts and retry policies
    - Task queue for Go activities
    - Default policy refs for all agents

    Args:
        config: Runtime configuration

    Example:
        >>> with set_kivot_agent_overrides(config):
        ...     # KivotRunner will use this config
        ...     result = await KivotRunner.run(agent, input=prompt)
    """
    token = _config_var.set(config)
    try:
        yield config
    finally:
        _config_var.reset(token)


def get_kivot_config() -> Optional[KivotAgentsConfig]:
    """
    Get current Kivot Agents configuration.

    Returns:
        Current config if set via set_kivot_agent_overrides, else None
    """
    return _config_var.get()
