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_request (headers, body):
    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 send_prompt(api_key, messages):
    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": messages,
        "thinking": {
            "type": "enabled",
            "budget_tokens": 1024
        }
    }

    return send_request(headers, body)


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


def run():
    messages = []

    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]
    messages.append({
        "role": "user",
        "content": prompt
    })

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

    # We can't just append the whole first response, otherwise we get 400 bad req:
    # "Extra inputs are not permitted"
    messages.append({
        'role': 'assistant',
        'content': first_response.get('content', []),
    })

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

    # print first response from Claude
    for block in first_response.get("content", []):
        print(f'{block['type']}:')
        if block["type"] == "text":
            print(f'\t{block["text"]}')
        elif block["type"] == "thinking":
            print(f'\t{block["thinking"]}')
        elif block["type"] == "tool_use":
            print(f'\tTool use requested: {block["name"]}')
            print(f'\tInput: {block["input"]}')

    # this is a stub for the tool execution
    fake_result_message = {
        "role": "user",
        "content": [
            {
                "type": "tool_result",
                "tool_use_id": tool_use["id"],
                "content": "(Stub) Executed tool successfully"
            }
        ]
    }
    messages.append(fake_result_message)

    # send the fake result back to Claude
    headers = HEADERS.copy()
    headers["x-api-key"] = api_key

    tool_result_msg = {
        "model": MODEL,
        "max_tokens": 2048,
        "tools": [],
        "messages": messages,
        "thinking": {
            "type": "enabled",
            "budget_tokens": 1024
        }
    }

    final_response = send_request(headers, tool_result_msg)

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


if __name__ == "__main__":
    run()