from askui import VisionAgent

with VisionAgent() as agent:
    # AskUI used as default model

    agent.click("search field")

    # Use Anthropic (Claude 3.5 Sonnet V2) as model
    agent.click("search field", model="anthropic-claude-3-5-sonnet-20241022")

    agent.click("search field", model="OS-Copilot/OS-Atlas-Base-7B")

from askui import VisionAgent

# Initialize your agent context manager
with VisionAgent() as agent:
    # Use the webbrowser tool to start browsing
    agent.tools.webbrowser.open_new("http://www.google.com")

    # Start to automate individual steps
    agent.click("url bar")
    agent.type("http://www.google.com")
    agent.keyboard("enter")

    # Extract information from the screen
    datetime = agent.get("What is the datetime at the top of the screen?")
    print(datetime)

    # Or let the agent work on its own, needs an Anthropic key set
    agent.act("search for a flight from Berlin to Paris in January")

from askui import VisionAgent

# Use AskUI's combo model for all commands
with VisionAgent(model="askui-combo") as agent:
    agent.click("Next")  # Uses askui-combo
    agent.get("What's on screen?")  # Uses askui-combo

# Use different models for different tasks
with VisionAgent(model={
    "act": "anthropic-claude-3-5-sonnet-20241022",  # Use Claude for act()
    "get": "askui",  # Use AskUI for get()
    "locate": "askui-combo",  # Use AskUI combo for locate() (and click(), mouse_move())
}) as agent:
    agent.act("Search for flights")  # Uses Claude
    agent.get("What's the current page?")  # Uses AskUI
    agent.click("Submit")  # Uses AskUI combo

# You can still override the default model for individual commands
with VisionAgent(model="askui-combo") as agent:
    agent.click("Next")  # Uses askui-combo (default)
    agent.click("Previous", model="askui-pta")  # Override with askui-pta
    agent.click("Submit")  # Back to askui-combo (default)

import functools
from askui import (
    ActModel,
    GetModel,
    LocateModel,
    Locator,
    ImageSource,
    MessageParam,
    ModelComposition,
    ModelRegistry,
    OnMessageCb,
    Point,
    ResponseSchema,
    VisionAgent,
)
from typing import Type
from typing_extensions import override

# Define custom models
class MyActModel(ActModel):
    @override
    def act(
        self,
        messages: list[MessageParam],
        model_choice: str,
        on_message: OnMessageCb | None = None,
    ) -> None:
        # Implement custom act logic, e.g.:
        # - Use a different AI model
        # - Implement custom business logic
        # - Call external services
        if len(messages) > 0:
            goal = messages[0].content
            print(f"Custom act model executing goal: {goal}")
        else:
            error_msg = "No messages provided"
            raise ValueError(error_msg)

# Because Python supports multiple inheritance, we can subclass both `GetModel` and `LocateModel` (and even `ActModel`)
# to create a model that can both get and locate elements.
class MyGetAndLocateModel(GetModel, LocateModel):
    @override
    def get(
        self,
        query: str,
        image: ImageSource,
        response_schema: Type[ResponseSchema] | None,
        model_choice: str,
    ) -> ResponseSchema | str:
        # Implement custom get logic, e.g.:
        # - Use a different OCR service
        # - Implement custom text extraction
        # - Call external vision APIs
        return f"Custom response to query: {query}"


    @override
    def locate(
        self,
        locator: str | Locator,
        image: ImageSource,
        model_choice: ModelComposition | str,
    ) -> Point:
        # Implement custom locate logic, e.g.:
        # - Use a different object detection model
        # - Implement custom element finding
        # - Call external vision services
        return (100, 100)  # Example coordinates


# Create model registry
custom_models: ModelRegistry = {
    "my-act-model": MyActModel(),
    "my-get-model": MyGetAndLocateModel(),
    "my-locate-model": MyGetAndLocateModel(),
}

# Initialize agent with custom models
with VisionAgent(models=custom_models) as agent:
    # Use custom models for specific tasks
    agent.act("search for flights", model="my-act-model")

    # Get information using custom model
    result = agent.get(
        "what's the current page title?",
        model="my-get-model"
    )

    # Click using custom locate model
    agent.click("submit button", model="my-locate-model")

    # Mix and match with default models
    agent.click("next", model="askui")  # Uses default AskUI model

class DynamicActModel(ActModel):
    @override
    def act(
        self,
        messages: list[MessageParam],
        model_choice: str,
        on_message: OnMessageCb | None = None,
    ) -> None:
        pass

# going to be called each time model is chosen using `model` parameter
def create_custom_model(api_key: str) -> ActModel:
    return DynamicActModel()


# if you don't want to recreate a new model on each call but rather just initialize
# it lazily
@functools.cache
def create_custom_model_cached(api_key: str) -> ActModel:
    return DynamicActModel()


# Register model factory
custom_models: ModelRegistry = {
    "dynamic-model": lambda: create_custom_model("your-api-key"),
    "dynamic-model-cached": lambda: create_custom_model_cached("your-api-key"),
    "askui": lambda: create_custom_model_cached("your-api-key"), # overrides default model
    "anthropic-claude-3-5-sonnet-20241022": lambda: create_custom_model_cached("your-api-key"), # overrides model
}


with VisionAgent(models=custom_models, model="dynamic-model") as agent:
    agent.act("do something") # creates and uses instance of DynamicActModel
    agent.act("do something") # creates and uses instance of DynamicActModel
    agent.act("do something", model="dynamic-model-cached") # uses new instance of DynamicActModel as it is the first call
    agent.act("do something", model="dynamic-model-cached") # reuses cached instance

from askui import VisionAgent
from askui.models import (
    OpenRouterGetModel,
    OpenRouterSettings,
    ModelRegistry,
)


# Register OpenRouter model in the registry
custom_models: ModelRegistry = {
    "my-custom-model": OpenRouterGetModel(
        OpenRouterSettings(
            model="anthropic/claude-opus-4",
        )
    ),
}

with VisionAgent(models=custom_models, model={"get":"my-custom-model"}) as agent:
    result = agent.get("What is the main heading on the screen?")
    print(result)

agent.tools.os.click("left", 2) # clicking
agent.tools.os.mouse_move(100, 100) # mouse movement
agent.tools.os.keyboard_tap("v", modifier_keys=["control"]) # Paste
# and many more

agent.tools.webbrowser.open_new("http://www.google.com")
# also check out open and open_new_tab

agent.tools.clipboard.copy("...")
result = agent.tools.clipboard.paste()

import logging

with VisionAgent(log_level=logging.DEBUG) as agent:
    agent...

from typing import Optional, Union
from typing_extensions import override
from askui.reporting import SimpleHtmlReporter
from PIL import Image

with VisionAgent(reporters=[SimpleHtmlReporter()]) as agent:
    agent...

from askui.reporting import Reporter

class CustomReporter(Reporter):
    @override
    def add_message(
        self,
        role: str,
        content: Union[str, dict, list],
        image: Optional[Image.Image | list[Image.Image]] = None,
    ) -> None:
        # adding message to the report (see implementation of `SimpleHtmlReporter` as an example)
        pass

    @override
    def generate(self) -> None:
        # generate the report if not generated live (see implementation of `SimpleHtmlReporter` as an example)
        pass


with VisionAgent(reporters=[CustomReporter()]) as agent:
    agent...

with VisionAgent(reporters=[SimpleHtmlReporter(), CustomReporter()]) as agent:
    agent...

with VisionAgent(display=1) as agent:
    agent...

agent.click("Password textfield")
agent.type("********")

from askui import locators as loc

password_textfield_label = loc.Text("Password")
password_textfield = loc.Element("textfield").right_of(password_textfield_label)

agent.click(password_textfield)
agent.type("********")

# Get text from screen
url = agent.get("What is the current url shown in the url bar?")
print(url)  # e.g., "github.com/login"

# Check UI state
# Just as an example, may be flaky if used as is, better use a response schema to check for a boolean value (see below)
is_logged_in = agent.get("Is the user logged in? Answer with 'yes' or 'no'.") == "yes"
if is_logged_in:
    agent.click("Logout")
else:
    agent.click("Login")

from PIL import Image
from askui import VisionAgent

# From PIL Image
with VisionAgent() as agent:
  image = Image.open("screenshot.png")
  result = agent.get("What's in this image?", image)

  # From file path
  result = agent.get("What's in this image?", "screenshot.png")

from askui import ResponseSchemaBase, VisionAgent
from PIL import Image
import json

class UserInfo(ResponseSchemaBase):
    username: str
    is_online: bool

class UrlResponse(ResponseSchemaBase):
    url: str

class NestedResponse(ResponseSchemaBase):
    nested: UrlResponse

class LinkedListNode(ResponseSchemaBase):
    value: str
    next: "LinkedListNode | None"

with VisionAgent() as agent:
    # Get structured data
    user_info = agent.get(
        "What is the username and online status?",
        response_schema=UserInfo
    )
    print(f"User {user_info.username} is {'online' if user_info.is_online else 'offline'}")

    # Get URL as string
    url = agent.get("What is the current url shown in the url bar?")
    print(url)  # e.g., "github.com/login"

    # Get URL as Pydantic model from image at (relative) path
    response = agent.get(
        "What is the current url shown in the url bar?",
        response_schema=UrlResponse,
        image="screenshot.png",
    )

    # Dump whole model
    print(response.model_dump_json(indent=2))
    # or
    response_json_dict = response.model_dump(mode="json")
    print(json.dumps(response_json_dict, indent=2))
    # or for regular dict
    response_dict = response.model_dump()
    print(response_dict["url"])

    # Get boolean response from PIL Image
    is_login_page = agent.get(
        "Is this a login page?",
        response_schema=bool,
        image=Image.open("screenshot.png"),
    )
    print(is_login_page)

    # Get integer response
    input_count = agent.get(
        "How many input fields are visible on this page?",
        response_schema=int,
    )
    print(input_count)

    # Get float response
    design_rating = agent.get(
        "Rate the page design quality from 0 to 1",
        response_schema=float,
    )
    print(design_rating)

    # Get nested response
    nested = agent.get(
        "Extract the URL and its metadata from the page",
        response_schema=NestedResponse,
    )
    print(nested.nested.url)

    # Get recursive response
    linked_list = agent.get(
        "Extract the breadcrumb navigation as a linked list",
        response_schema=LinkedListNode,
    )
    current = linked_list
    while current:
        print(current.value)
        current = current.next