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, TEMPLATE_UUID, INSTANT_ID_TEMPLATE_UUID
except ImportError:
    print(f"Failed to import LibLibAIClient. Make sure the path {tools_dir} is correct.")
    sys.exit(1)

def main():
    client = LibLibAIClient()

    # 1. Search for a checkpoint's UUID
    keyword = "Juggernaut XL"
    print(f"1. Searching for Checkpoint matching '{keyword}'...")
    
    search_res = client.search_models(keyword)
    if not search_res or search_res.get('code') != 0 or not search_res.get('data', {}).get('data'):
        print("Search failed or no models found.")
        print(search_res)
        return
        
    first_model = search_res['data']['data'][0]
    uuid = first_model['uuid']
    version_uuid = first_model['versionUuid']
    model_name = first_model.get('name')
    base_types = first_model.get('baseType', [])
    
    print(f"   -> Found Checkpoint: '{model_name}'")
    print(f"   -> UUID: {uuid}")
    print(f"   -> Version UUID: {version_uuid}")
    print(f"   -> baseType: {base_types} (1=1.5, 2=XL etc.)\n")

    # 2. Match local template_id
    print(f"2. Matching local template_id based on baseType ...")
    print(f"   -> Client is currently configured to use TEMPLATE_UUID: {TEMPLATE_UUID}\n")

    # 3. Read local image and encode it to base64
    local_image_path = os.path.join(script_dir, "..", "depth_map.png")
    print(f"3. Reading local reference image: {local_image_path}")
    if not os.path.exists(local_image_path):
        print(f"Error: {local_image_path} does not exist.")
        return
        
    with open(local_image_path, "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}"

    # 4. Request ControlNet Image Generation (mode='depth')
    print("4. Executing controlnet generation (using Depth + found Checkpoint)...")
    try:
        gen_res = client.generate_advanced(
            mode="depth",  # Changed to depth since it's a depth map
            prompt="A masterpiece, best quality, an easel on the left, a beautiful girl holding a clipboard on the right, highly detailed, photorealistic, cinematic lighting, 8k resolution",
            image=image_payload,
            base_model_uuid=version_uuid,
            width=1024,
            height=1024,
            steps=5,
            cfg_scale=1.5,
            control_nets=[{
                "mode": "depth",
                "image": image_payload,
                "weight": 0.4
            }]
        )
        print("\n=== Generation Task Success ===")
        print(json.dumps(gen_res, indent=2, ensure_ascii=False))
        
    except Exception as e:
        print(f"\n=== Generation Error ===")
        print(f"Error: {e}")

if __name__ == '__main__':
    main()