import os
import sys
import httpx
import json

ANTHROPIC_API_URL = "https://api.anthropic.com/v1/messages"
MODEL = "claude-3-7-sonnet-20250219"
BETA_FLAG = "computer-use-2025-01-24"
TOOL_VERSION = "20250124"

HEADERS = {
    "content-type": "application/json",
    "anthropic-version": "2023-06-01",
    "anthropic-beta": BETA_FLAG
}


def send_prompt(api_key, user_prompt):
    headers = HEADERS.copy()
    headers["x-api-key"] = api_key

    body = {
        "model": MODEL,
        "max_tokens": 1024,
        "tools": [
            {
                "type": f"computer_{TOOL_VERSION}",
                "name": "computer",
                "display_width_px": 1024,
                "display_height_px": 768,
                "display_number": 1
            },
            {
                "type": f"bash_{TOOL_VERSION}",
                "name": "bash"
            },
            {
                "type": f"text_editor_{TOOL_VERSION}",
                "name": "str_replace_editor"
            }
        ],
        "messages": [
            {
                "role": "user",
                "content": user_prompt
            }
        ],
        "thinking": {
            "type": "enabled",
            "budget_tokens": 512
        }
    }

    with httpx.Client(timeout=60.0) as client:
        response = client.post(ANTHROPIC_API_URL, headers=headers, json=body)
        response.raise_for_status()
        return response.json()


def extract_tool_use(response):
    for block in response.get("content", []):
        if block.get("type") == "tool_use":
            return block
    return None


def run():
    if len(sys.argv) < 3:
        print("Usage: python claude_computer_use_cli.py <api_key> '<your_prompt>'")
        return

    api_key = sys.argv[1]
    prompt = sys.argv[2]

    print(f"Sending prompt: {prompt}")
    first_response = send_prompt(api_key, prompt)

    tool_use = extract_tool_use(first_response)
    if not tool_use:
        print("No tool use required. Response:")
        print(first_response["content"])
        return

    print(f"Tool requested: {tool_use['name']}")
    print(f"Input: {tool_use['input']}")

    fake_result = {"result": f"(Stub) Executed {tool_use['name']} successfully."}

    headers = HEADERS.copy()
    headers["x-api-key"] = api_key

    tool_result_msg = {
        "model": MODEL,
        "max_tokens": 512,
        "tools": [],
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "tool_result",
                        "tool_use_id": tool_use["id"],
                        "content": fake_result
                    }
                ]
            }
        ]
    }

    with httpx.Client(timeout=60.0) as client:
        response = client.post(ANTHROPIC_API_URL, headers=headers, json=tool_result_msg)
        response.raise_for_status()
        final_response = response.json()

    print("Final Claude response:")
    for block in final_response.get("content", []):
        if block["type"] == "text":
            print(block["text"])


if __name__ == "__main__":
    run()