#
#  Copyright (c) 2023-2024 - Restate Software, Inc., Restate GmbH
#
#  This file is part of the Restate SDK for Python,
#  which is released under the MIT license.
#
#  You can find a copy of the license in file LICENSE in the root
#  directory of this repository or package, or at
#  https://github.com/restatedev/sdk-typescript/blob/main/LICENSE
#
"""
This module contains the optional OpenAI integration for Restate.
"""

import dataclasses

from agents import (
    Usage,
    AgentsException,
)
from agents.items import TResponseOutputItem
from agents.items import TResponseInputItem
from datetime import timedelta
from typing import Optional
from pydantic import BaseModel

from restate.ext.turnstile import Turnstile
from restate import TerminalError


class State:
    __slots__ = ("turnstile",)

    def __init__(self) -> None:
        self.turnstile = Turnstile([])


@dataclasses.dataclass
class LlmRetryOpts:
    max_attempts: Optional[int] = 10
    """Max number of attempts (including the initial), before giving up.

    When giving up, the LLM call will throw a `TerminalError` wrapping the original error message."""
    max_duration: Optional[timedelta] = None
    """Max duration of retries, before giving up.

    When giving up, the LLM call will throw a `TerminalError` wrapping the original error message."""
    initial_retry_interval: Optional[timedelta] = timedelta(seconds=1)
    """Initial interval for the first retry attempt.
    Retry interval will grow by a factor specified in `retry_interval_factor`.

    If any of the other retry related fields is specified, the default for this field is 50 milliseconds, otherwise restate will fallback to the overall invocation retry policy."""
    max_retry_interval: Optional[timedelta] = None
    """Max interval between retries.
    Retry interval will grow by a factor specified in `retry_interval_factor`.

    The default is 10 seconds."""
    retry_interval_factor: Optional[float] = None
    """Exponentiation factor to use when computing the next retry delay.

    If any of the other retry related fields is specified, the default for this field is `2`, meaning retry interval will double at each attempt, otherwise restate will fallback to the overall invocation retry policy."""


# The OpenAI ModelResponse class is a dataclass with Pydantic fields.
# The Restate SDK cannot serialize this. So we turn the ModelResponse int a Pydantic model.
class RestateModelResponse(BaseModel):
    output: list[TResponseOutputItem]
    """A list of outputs (messages, tool calls, etc) generated by the model"""

    usage: Usage
    """The usage information for the response."""

    response_id: str | None
    """An ID for the response which can be used to refer to the response in subsequent calls to the
    model. Not supported by all model providers.
    If using OpenAI models via the Responses API, this is the `response_id` parameter, and it can
    be passed to `Runner.run`.
    """

    def to_input_items(self) -> list[TResponseInputItem]:
        return [it.model_dump(exclude_unset=True) for it in self.output]  # type: ignore


class AgentsTerminalException(AgentsException, TerminalError):
    """Exception that is both an AgentsException and a restate.TerminalError."""

    def __init__(self, *args: object) -> None:
        super().__init__(*args)
