#!/usr/bin/env python3
"""Quick test to query CloudKit records."""

import json, hashlib, base64, requests, os, sys
from datetime import datetime, timezone

try:
    from cryptography.hazmat.primitives import hashes, serialization
    from cryptography.hazmat.primitives.asymmetric import ec
    from cryptography.hazmat.backends import default_backend
except ImportError:
    sys.exit("Error: pip install cryptography")

CONTAINER = "iCloud.com.sportstime.app"
HOST = "https://api.apple-cloudkit.com"

def sign(key_data, date, body, path):
    key = serialization.load_pem_private_key(key_data, None, default_backend())
    body_hash = base64.b64encode(hashlib.sha256(body.encode()).digest()).decode()
    sig = key.sign(f"{date}:{body_hash}:{path}".encode(), ec.ECDSA(hashes.SHA256()))
    return base64.b64encode(sig).decode()

def query(key_id, key_data, record_type, env='development'):
    path = f"/database/1/{CONTAINER}/{env}/public/records/query"
    body = json.dumps({
        'query': {'recordType': record_type},
        'resultsLimit': 10
    })
    date = datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ')
    headers = {
        'Content-Type': 'application/json',
        'X-Apple-CloudKit-Request-KeyID': key_id,
        'X-Apple-CloudKit-Request-ISO8601Date': date,
        'X-Apple-CloudKit-Request-SignatureV1': sign(key_data, date, body, path),
    }
    r = requests.post(f"{HOST}{path}", headers=headers, data=body, timeout=30)
    return r.status_code, r.json()

if __name__ == '__main__':
    key_id = os.environ.get('CLOUDKIT_KEY_ID') or (sys.argv[1] if len(sys.argv) > 1 else None)
    key_file = os.environ.get('CLOUDKIT_KEY_FILE') or (sys.argv[2] if len(sys.argv) > 2 else 'eckey.pem')
    
    if not key_id:
        sys.exit("Usage: python test_cloudkit.py KEY_ID [KEY_FILE]")
    
    key_data = open(key_file, 'rb').read()
    
    print("Testing CloudKit connection...\n")
    
    for record_type in ['Stadium', 'Team', 'Game']:
        status, result = query(key_id, key_data, record_type)
        count = len(result.get('records', []))
        print(f"{record_type}: status={status}, records={count}")
        if count > 0:
            print(f"  Sample: {result['records'][0].get('recordName', 'N/A')}")
        if 'serverErrorCode' in result:
            print(f"  Error: {result.get('serverErrorCode')}: {result.get('reason')}")
    
    print("\nFull response for Stadium query:")
    status, result = query(key_id, key_data, 'Stadium')
    print(json.dumps(result, indent=2)[:1000])