import anthropic
import json
import sys

api_key = sys.argv[1]
model = sys.argv[2] if len(sys.argv) > 2 else 'claude-sonnet-4-6'
print(f'API Key: {api_key[:12]}...{api_key[-4:]}')
print(f'Model: {model}')

client = anthropic.Anthropic(api_key=api_key)

print()
print('--- Test 1: Basic API Call ---')
try:
    r = client.messages.create(model=model, max_tokens=100, messages=[{'role': 'user', 'content': 'Say hello in one sentence.'}])
    print(f'OK: {r.content[0].text}')
    print(f'Tokens: {r.usage.input_tokens} in, {r.usage.output_tokens} out')
except Exception as e:
    print(f'FAILED: {e}')
    sys.exit(1)

print()
print('--- Test 2: Tool Calling ---')
try:
    tools = [{'name': 'get_account_balance', 'description': 'Get balance of an account by code', 'input_schema': {'type': 'object', 'properties': {'account_code': {'type': 'string', 'description': 'Account code'}}, 'required': ['account_code']}}]
    r = client.messages.create(model=model, max_tokens=300, system='You are an accounting AI. Always use tools to look up data before answering.', messages=[{'role': 'user', 'content': 'Look up the balance on account 2005.'}], tools=tools)
    print(f'Stop reason: {r.stop_reason}')
    tool_id = None
    for b in r.content:
        if b.type == 'text':
            print(f'Text: {b.text}')
        elif b.type == 'tool_use':
            print(f'TOOL CALL: {b.name}({json.dumps(b.input)}) id={b.id}')
            tool_id = b.id
    if r.stop_reason == 'tool_use':
        print('RESULT: Tool calling WORKING')
    else:
        print('RESULT: No tool call (model answered directly)')
except Exception as e:
    print(f'FAILED: {e}')
    sys.exit(1)

print()
print('--- Test 3: Tool Result Round-Trip ---')
try:
    if not tool_id:
        tool_id = 'toolu_test123'
    msgs = [
        {'role': 'user', 'content': 'Look up account 2005 balance.'},
        {'role': 'assistant', 'content': [{'type': 'tool_use', 'id': tool_id, 'name': 'get_account_balance', 'input': {'account_code': '2005'}}]},
        {'role': 'user', 'content': [{'type': 'tool_result', 'tool_use_id': tool_id, 'content': json.dumps({'balance': -15234.56, 'name': 'HST Collected'})}]}
    ]
    r2 = client.messages.create(model=model, max_tokens=200, system='You are an accounting AI. Report findings in Canadian dollars.', messages=msgs, tools=tools)
    for b in r2.content:
        if b.type == 'text':
            print(f'AI: {b.text}')
    print(f'Tokens: {r2.usage.input_tokens} in, {r2.usage.output_tokens} out')
    print('RESULT: Multi-turn tool flow WORKING')
except Exception as e:
    print(f'FAILED: {e}')
    sys.exit(1)

print()
print('=== ALL TESTS PASSED ===')