import base64
import requests
import json
import sys

def test_gen(style_id, output_path):
    with open('public/test_skeleton_kid.jpg', 'rb') as f:
        img_bytes = f.read()
    b64 = 'data:image/jpeg;base64,' + base64.b64encode(img_bytes).decode('utf-8')
    
    payload = {
        'imageBase64': b64,
        'styleId': style_id,
        'clientId': 'dev_test_batch'
    }
    
    print(f"Calling http://localhost:3001/api/generate for {style_id}...")
    try:
        res = requests.post('http://localhost:3001/api/generate', json=payload, timeout=90)
        print(f"Status: {res.status_code}")
        data = res.json()
        if res.status_code == 200 and 'imageUrl' in data:
            img_url = data['imageUrl']
            print(f"Success! Provider: {data.get('providerInfo')}")
            if img_url.startswith('data:'):
                header, b64data = img_url.split(',', 1)
                with open(output_path, 'wb') as out_f:
                    out_f.write(base64.b64decode(b64data))
            else:
                full_url = 'http://localhost:3001' + img_url
                img_res = requests.get(full_url)
                with open(output_path, 'wb') as out_f:
                    out_f.write(img_res.content)
            print(f"Saved to {output_path}")
            return True
        else:
            print("Error response:", data)
            return False
    except Exception as e:
        print("Exception:", e)
        return False

if __name__ == '__main__':
    style = sys.argv[1]
    out = sys.argv[2]
    test_gen(style, out)