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": 2048,
        "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": 1024
        }
    }

    with httpx.Client(timeout=60.0) as client:
        try:
            response = client.post(ANTHROPIC_API_URL, headers=headers, json=body)
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as e:
            print(f"[ERROR] HTTP {e.response.status_code} - {e.response.reason_phrase}")
            try:
                print("[DETAIL] Response JSON:")
                print(json.dumps(e.response.json(), indent=2))
            except Exception:
                print("[DETAIL] Raw Response Text:")
                print(e.response.text)
            sys.exit(1)


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)
    print(first_response)

    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']}")
    print(tool_use)

    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": 2048,
        "tools": [],
        "messages": [
            {
                "role": "user",
                "content": prompt
            },
            {
                "role": "assistant",
                "content": [
                    {
                        'type': 'thinking',
                        'thinking': (
                            "The user is asking me to take a screenshot of their desktop. I can do this "
                            "using the \"computer\" function with the \"screenshot\" action.\n\n"
                            "I need to call the following:\n"
                            "- Function: computer\n"
                            "- Action: screenshot"
                        ),
                        'signature': (
                            "ErUBCkYIBBgCIkDPXNmcKH0varjvxjdtV6bzHp8OfBbJ3Jg+rJZvBGpXIE1L4RrroreDvBLcoC4uO76tNHDikWlnwv7UAgUtHEFMEgy6tLHT/"
                            "oFQ4tczh00aDMND7yTs4mXU4utOjiIw08TVLfADRoZLfX5ookvmFuDfN0m2l7KPcgAfQ1Af7AH+T2YHqi+DRtnuGxt35vAeKh1sTRvojVlHU"
                            "4en+7+TxdQ+yEC3x70IkuQgF5DHQhgC"
                        )
                    },
                    {
                        'type': 'text',
                        'text': "I'll take a screenshot of your desktop right away."
                    },
                    {
                        'type': 'tool_use',
                        'id': tool_use["id"],
                        'name': 'computer',
                        'input': {
                            'action': 'screenshot'
                        }
                    }
                ]
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "tool_result",
                        "tool_use_id": tool_use["id"],
                        "content": "(Stub) Executed tool successfully"
                    }
                ]
            }
        ],
        "thinking": {
            "type": "enabled",
            "budget_tokens": 1024
        }
    }

    with httpx.Client(timeout=60.0) as client:
        try:
            response = client.post(ANTHROPIC_API_URL, headers=headers, json=tool_result_msg)
            response.raise_for_status()
            final_response = response.json()
        except httpx.HTTPStatusError as e:
            print(f"[ERROR] HTTP {e.response.status_code} - {e.response.reason_phrase}")
            try:
                print("[DETAIL] Response JSON:")
                print(json.dumps(e.response.json(), indent=2))
            except Exception:
                print("[DETAIL] Raw Response Text:")
                print(e.response.text)
            sys.exit(1)

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


if __name__ == "__main__":
    run()