import sys
import os
import json

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
    # Because liblibai_client defaults to XL controlnet models (self.sdxl_canny),
    # we search for an XL model to ensure the baseType matches (Rule 1).
    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
        
    # Get the first result
    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
    # According to our docs: `e10adc3949ba59abbe56e057f20f883e` is for 1.5/XL.
    print(f"2. Matching local template_id based on baseType ...")
    print(f"   -> Client is currently configured to use TEMPLATE_UUID: {TEMPLATE_UUID}")
    print(f"   -> Client INSTANT_ID_TEMPLATE_UUID: {INSTANT_ID_TEMPLATE_UUID}\n")

    # 3. Request ControlNet Image Generation
    print("3. Executing controlnet generation (using Canny + found Checkpoint)...")
    
    # We use a valid test image (has to be standard or accessible)
    # This is an image URL hosted on liblib's OSS or accessible publicly.
    # Using a clean realistic portrait to avoid Canny extracting text/watermarks.
    test_image_url = "https://images.unsplash.com/photo-1544005313-94ddf0286df2?ixlib=rb-4.0.3&auto=format&fit=crop&w=1024&q=80"
    
    try:
        # Request advanced generation with Canny controlnet
        gen_res = client.generate_advanced(
            mode="canny",
            prompt="A masterpiece, best quality, beautiful girl, highly detailed, photorealistic, 8k resolution",
            image=test_image_url,
            base_model_uuid=version_uuid,  # Override CheckPointId with the found version_uuid
            width=1024,
            height=1024,
            steps=5,
            cfg_scale=1.5
        )
        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()