"""
Plugin for integrating OpenAI Agents SDK with Kivot.

Provides seamless configuration and activity registration.
"""

import dataclasses
from contextlib import asynccontextmanager
from datetime import timedelta
from typing import Optional, Sequence

from temporalio.common import RetryPolicy
from temporalio.plugin import SimplePlugin
from temporalio.worker import WorkflowRunner
from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner

from .overrides import (
    ActivityParameters,
    KivotAgentsConfig,
    TruncationConfig,
    set_kivot_agent_overrides,
)


class OpenAIAgentsKivotPlugin(SimplePlugin):
    """
    Plugin for integrating OpenAI Agents SDK with Kivot.

    Automatically configures:
    - Runtime overrides for KivotRunner
    - Go activity references and timeouts
    - Workflow sandbox restrictions

    Example:
        >>> from temporalio.client import Client
        >>> from temporalio.worker import Worker
        >>> from kivot_agents import OpenAIAgentsKivotPlugin, ActivityParameters
        >>>
        >>> plugin = OpenAIAgentsKivotPlugin(
        ...     model_params=ActivityParameters(
        ...         start_to_close_timeout=timedelta(seconds=60)
        ...     )
        ... )
        >>>
        >>> client = await Client.connect("localhost:7233", plugins=[plugin])
        >>> worker = Worker(client, task_queue="agent-workflows", workflows=[...])
    """

    def __init__(
        self,
        model_params: Optional[ActivityParameters] = None,
        tool_params: Optional[ActivityParameters] = None,
        thread_params: Optional[ActivityParameters] = None,
        go_activities_task_queue: str = "core-activities",
        default_policy_refs: Optional[Sequence[str]] = None,
        thread_history_limit: int = 50,
        truncation: Optional[TruncationConfig] = None,
    ):
        """
        Initialize the OpenAI Agents Kivot plugin.

        Args:
            model_params: Timeout/retry config for ModelStepActivity
            tool_params: Timeout/retry config for ToolStepActivity
            thread_params: Timeout/retry config for thread activities (Create, Get, Append)
            go_activities_task_queue: Task queue for Go activities
            default_policy_refs: Default policy refs for all agents
            thread_history_limit: Max messages to load from thread history (default 50)
            truncation: Optional context window truncation config
        """
        # Default parameters
        if model_params is None:
            model_params = ActivityParameters(
                start_to_close_timeout=timedelta(seconds=60),
                retry_policy=RetryPolicy(
                    maximum_attempts=3,
                    initial_interval=timedelta(seconds=1),
                    maximum_interval=timedelta(seconds=10),
                ),
            )

        if tool_params is None:
            tool_params = ActivityParameters(
                start_to_close_timeout=timedelta(seconds=30),
                retry_policy=RetryPolicy(
                    maximum_attempts=2,
                    initial_interval=timedelta(seconds=2),
                    maximum_interval=timedelta(seconds=10),
                ),
            )

        self.config = KivotAgentsConfig(
            model_params=model_params,
            tool_params=tool_params,
            thread_params=thread_params,
            go_activities_task_queue=go_activities_task_queue,
            default_policy_refs=list(default_policy_refs or []),
            thread_history_limit=thread_history_limit,
            truncation=truncation,
        )

        # Context manager for runtime overrides
        @asynccontextmanager
        async def run_context():
            with set_kivot_agent_overrides(self.config):
                yield

        # Workflow runner with sandbox restrictions
        def workflow_runner(runner: Optional[WorkflowRunner]) -> WorkflowRunner:
            if not runner:
                raise ValueError(
                    "No WorkflowRunner provided to OpenAIAgentsKivotPlugin"
                )

            # Add passthrough for agents module if sandboxed
            if isinstance(runner, SandboxedWorkflowRunner):
                return dataclasses.replace(
                    runner,
                    restrictions=runner.restrictions.with_passthrough_modules(
                        "agents", "openai"
                    ),
                )
            return runner

        # Initialize SimplePlugin
        super().__init__(
            name="OpenAIAgentsKivotPlugin",
            run_context=lambda: run_context(),
            workflow_runner=workflow_runner,
        )
