from __future__ import annotations
import abc

from typing import (
    TypedDict, Dict, List, Any, Optional, Literal,
    Callable, Tuple, Awaitable, Union, Protocol,
    TypeVar, Type,
)

from typing_extensions import NotRequired


class ToolCallFunction(TypedDict):
    name: str
    arguments: str   # <-- JSON-encoded arguments (kept as a string for reproducibility reasons)


class ToolCall(TypedDict):
    call_id: str
    call_type: Literal['function']   # <-- functions are the only type of "tool use" right now!
    function: ToolCallFunction


class ToolResultBase(TypedDict):
    call_id: str
    is_error: NotRequired[bool]


class ToolResultAny(ToolResultBase):
    type: Literal['any']
    result: Any


class ToolResultLayeredAgent(ToolResultBase):
    type: Literal['layered_agent']
    run: AgentRun


ToolResult = Union[
    ToolResultAny,
    ToolResultLayeredAgent,
]


class Cost(TypedDict):
    input_tokens: Optional[int]
    output_tokens: Optional[int]
    total_tokens: Optional[int]


class MediaImage(TypedDict):
    type: Literal['image']
    image: str   # <-- either a local filepath or a remote URL


Media = MediaImage   # note: this will become a Union[MediaImage, ...] in the future as needed


class MessageBase(TypedDict):
    cost: NotRequired[Optional[Cost]]  # <-- cost of generating this message (for messages generated by an external service)
    raw: NotRequired[Any]              # <-- the raw payload(s) from the API (before being converted our our schema; only for messages generated by an external service)


class MessageContent(MessageBase):
    role: Literal['system', 'human', 'ai']
    text: Optional[str]
    media: NotRequired[List[Media]]


class MessageToolCall(MessageBase):
    role: Literal['tool_call']
    tools: List[ToolCall]


class MessageToolResult(MessageBase):
    role: Literal['tool_res']
    tools: List[ToolResult]


Message = Union[
    MessageContent,
    MessageToolCall,
    MessageToolResult,
]


class AgentRunBase(TypedDict):
    agent: str
    provider: NotRequired[str]
    model: NotRequired[str]
    model_kwargs: NotRequired[Dict[str, Any]]


class AgentRunMessageList(AgentRunBase):
    type: Literal['messages']
    messages: List[Message]


class AgentRunParallel(AgentRunBase):
    type: Literal['parallel']
    runs: List[AgentRun]


class AgentRunChained(AgentRunBase):
    type: Literal['chain']
    runs: List[AgentRun]


class AgentRunExtraction(AgentRunBase):
    type: Literal['extraction']
    messages: List[Message]
    result: Any    # ideally, it's `ExtractionType`, and we'd inherit from Generic[ExtractionType], but that's not supported until python3.11


AgentRun = Union[
    AgentRunMessageList,
    AgentRunParallel,
    AgentRunChained,
    AgentRunExtraction,
]


EventPayload = Union[
    Tuple[Literal['human'],       Literal['echo_event'],      MessageContent],

    Tuple[Literal['ai'],          Literal['text_event'],      str],

    Tuple[Literal['tool_call'],   Literal['text_event'],      str],
    Tuple[Literal['tool_call'],   Literal['tool_call_event'], ToolCall],

    Tuple[Literal['tool_res'],    Literal['tool_res_event'],  ToolResult],

    Tuple[Literal['progress'],    Literal['start'],           Tuple[str, str]],    # payload is `(key, details)`
    Tuple[Literal['progress'],    Literal['update'],          Tuple[str, float]],  # payload is `(key, progress_0_to_1)`
    Tuple[Literal['progress'],    Literal['end'],             str],                # payload is `key`

    Tuple[Literal['transaction'], Literal['start'],           Tuple[str, str]],    # payload is `(provider, model)`
    Tuple[Literal['transaction'], Literal['rollback'],        None],
    Tuple[Literal['transaction'], Literal['commit'],          None],

    Tuple[Literal['agent'],       Literal['start'],           str],                # payload is `agent_name`
    Tuple[Literal['agent'],       Literal['end'],             AgentRun],
]


EventCallback = Callable[[EventPayload], Awaitable[None]]


ExtractionType = TypeVar('ExtractionType')  # <-- for generic programming


class Model(abc.ABC):
    """
    Interface for an AI model.

    In many cases this will be an LLM (when you're in pure-text use-cases). But
    this interface also supports AI models that operate on multimodal content,
    thus we'll use the phrase "AI model" to be more general.
    """

    @abc.abstractmethod
    async def run(
        self,
        event_callback: EventCallback,
        messages: List[Message],
        tools: List[Callable],
        force_tool: bool = False,
        max_tool_iters: int = 5,
    ) -> List[Message]:
        """
        Generate one-or-more responses from the AI. More than one response
        is generated when tools are used. This method will respond to tool-
        use requests until the AI generates a non-tool response (up to
        `max_tool_iters`, then it halts).
        """
        pass

    @abc.abstractmethod
    async def extract(
        self,
        event_callback: EventCallback,
        messages: List[Message],
        extraction_type: Type[ExtractionType],
    ) -> Tuple[Message, ExtractionType]:
        """
        Use the AI to extract structured output from the `messages`. The
        schema of the extracted data will follow `extraction_type`.
        """
        pass

    @abc.abstractmethod
    def config_hash(self) -> str:
        """
        Returns a hash of this models configuration. This likely contains
        the provider's name, model's name, and any model_kwargs. You can use
        any algorithm you'd like as long as it's stable from one Python version
        to the next; e.g., we recommend you use `util.recursive_hash()`.
        """
        pass


AgentCallable = Callable[[Model, EventCallback, List[AgentRun]], Awaitable[AgentRun]]

BoundAgentCallable = Callable[[EventCallback, List[AgentRun]], Awaitable[AgentRun]]


class AgentRecord(TypedDict):
    name: str
    runner: AgentCallable


class ModelRecord(TypedDict):
    formal_name: str
    display_name: str
    outdated: NotRequired[bool]  # <-- if true, there is another model which replaces this one


class ModelFactory(Protocol):
    def __call__(self, model: str, **model_kwargs: Any) -> Model: ...


class ProviderRecord(TypedDict):
    name: str
    factory: ModelFactory
    models: List[ModelRecord]


class ModelSpec(TypedDict):
    provider: Union[str, ModelFactory]
    model: Union[str, ModelRecord]
    model_kwargs: NotRequired[Dict[str, Any]]


class AgentSpec(ModelSpec):
    agent: Union[str, AgentCallable]


class ToolParam(TypedDict):
    name: str
    type: str
    description: NotRequired[str]
    optional: NotRequired[bool]  # default = False (i.e. "required")


class CacheEventPayload(TypedDict):
    delta_time: float   # <-- when this event arrived (as the number of second since the start of the agent's execution)
    event: EventPayload


CacheKey = str


class CacheRecord(TypedDict):
    events: List[CacheEventPayload]
    run: AgentRun


ImageMimeTypes = Literal['image/jpeg', 'image/png', 'image/gif', 'image/webp']
