import sys
import os
import json
import base64

sys.stdout.reconfigure(encoding='utf-8')

script_dir = os.path.dirname(os.path.abspath(__file__))
tools_dir = os.path.join(script_dir, '..', 'tools', 'local', 'liblibai_controlnet')
sys.path.append(tools_dir)

try:
    from liblibai_client import LibLibAIClient
except ImportError:
    print(f"Failed to import LibLibAIClient. Make sure the path {tools_dir} is correct.")
    sys.exit(1)

def load_and_upload(client, filepath):
    with open(filepath, "rb") as f:
        image_bytes = f.read()
    b64_image = base64.b64encode(image_bytes).decode('utf-8')
    image_payload = f"data:image/png;base64,{b64_image}"
    return client.process_image_url(image_payload)

def main():
    client = LibLibAIClient()

    print("1. Searching for Juggernaut XL Checkpoint...")
    search_res = client.search_models('Juggernaut XL')
    first_model = search_res['data']['data'][0]
    version_uuid = first_model['versionUuid']
    print(f"   -> Version UUID: {version_uuid}")

    print("2. Uploading local images to LibLib OSS...")
    depth_map_path = os.path.join(script_dir, "..", "depth_map.png")
    character_path = os.path.join(script_dir, "..", "character_ref_main.png")
    
    depth_url = load_and_upload(client, depth_map_path)
    print(f"   -> Depth Map uploaded: {depth_url}")
    
    char_url = load_and_upload(client, character_path)
    print(f"   -> Character Ref uploaded: {char_url}")

    print("3. Submitting Dual ControlNet (Depth + IP-Adapter) Task...")
    payload = {
        'templateUuid': 'e10adc3949ba59abbe56e057f20f883e',
        'generateParams': {
            'checkPointId': version_uuid,
            'prompt': 'A masterpiece, best quality, a beautiful girl on the right holding a clipboard, an easel on the left, photorealistic, 8k resolution, cinematic lighting',
            'width': 1024,
            'height': 1024,
            'steps': 5,
            'cfgScale': 1.5,
            'imgCount': 1,
            'controlNet': [
                {
                    'unitOrder': 1,
                    'sourceImage': depth_url,
                    'width': 1024,
                    'height': 1024,
                    'preprocessor': 3,
                    'model': '6349e9dae8814084bd9c1585d335c24c', # SDXL Depth Model
                    'annotationParameters': {
                        'depthLeres': {
                            'preprocessorResolution': 1024,
                            'removeNear': 0,
                            'removeBackground': 0
                        }
                    },
                    'controlWeight': 1.0,
                    'startingControlStep': 0.0,
                    'endingControlStep': 1.0,
                    'pixelPerfect': 1,
                    'controlMode': 0
                },
                {
                    'unitOrder': 2,
                    'sourceImage': char_url,
                    'width': 1024,
                    'height': 1024,
                    'preprocessor': 0, 
                    'model': '8ea2538fdd7dcdea52b2da6b5151f875', # SDXL IP-Adapter Model
                    'annotationParameters': {}, # Usually empty for preprocessor=0
                    'controlWeight': 0.8,
                    'startingControlStep': 0.0,
                    'endingControlStep': 0.8,
                    'pixelPerfect': 1,
                    'controlMode': 0
                }
            ]
        }
    }

    try:
        task_id = client.submit_task_payload(payload)
        print(f"   -> Task submitted successfully! Task ID: {task_id}")
        print("4. Waiting for generation result...")
        res = client.wait_for_result(task_id)
        print("\n=== Generation Task Success ===")
        print(json.dumps(res, indent=2, ensure_ascii=False))
    except Exception as e:
        print(f"\n=== Generation Error ===")
        print(f"Error: {e}")

if __name__ == '__main__':
    main()