from typing import Annotated, Sequence, List, Optional, Any, Dict

from browsergym.core.env import BrowserEnv
from typing_extensions import TypedDict
import operator
from playwright.async_api import Page
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from pydantic import BaseModel, Field


class Prediction(TypedDict):
    action: str
    args: Optional[List[str]]


class AgentState(TypedDict):
    next: str  # The 'next' field indicates where to route to next
    pages: Annotated[Sequence[str], operator.add]  # List of pages traversed
    page: Page  # The Playwright web page lets us interact with the web environment
    input: str  # User request
    prediction: Prediction  # The Agent's output
    # A system message (or messages) containing the intermediate steps
    scratchpad: List[BaseMessage]
    observation: str  # The most recent response from a tool
    img: str  # b64 encoded screenshot
    annotations: List[str]  # Annotations for the current page
    extension_obj: Any  # The Playwright extension object
    actions: str  # The chosen actions
    url: str  # The URL of the current page
    elements: str  # The elements on the page
    messages: List[AIMessage]  # The messages exchanged between the user and the agent
    env: BrowserEnv  # The environment in which the agent is operating
    pointer_env: BrowserEnv  # The environment in which the agent is operating
    sender: str
    policy: str
    update_policy_reason: str
    read_page: str  # The outer text of the page
    env_policy: List[dict]
