import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '../tools/local/liblibai_controlnet'))
from liblibai_client import LibLibAIClient

client = LibLibAIClient()
TEMPLATE_UUID = "e10adc3949ba59abbe56e057f20f883e"
DEFAULT_CHECKPOINT_ID = "0ea388c7eb854be3ba3c6f65aac6bfd3" 

import requests

def test_key(preprocessor, key_name):
    payload = {
        "templateUuid": TEMPLATE_UUID,
        "generateParams": {
            "checkPointId": DEFAULT_CHECKPOINT_ID,
            "prompt": "cat",
            "width": 512,
            "height": 512,
            "imgCount": 1,
            "controlNet": [{
                "unitOrder": 1,
                "sourceImage": "https://liblibai-airship-temp.oss-cn-beijing.aliyuncs.com/aliyun-cn-prod/73ed6ae42b144d21bf566e05b5a6c138.png",
                "width": 512,
                "height": 512,
                "preprocessor": preprocessor,
                "model": "b6806516962f4e1599a93ac4483c3d23", # Random UUID
                "controlWeight": 1.0,
                "startingControlStep": 0.0,
                "endingControlStep": 1.0,
                "pixelPerfect": 1,
                "controlMode": 0,
                "annotationParameters": {
                    key_name: {
                        "preprocessorResolution": 512
                    }
                }
            }]
        }
    }
    url = client.generate_auth_url("/api/generate/webui/text2img")
    resp = requests.post(url, json=payload).json()
    if resp.get("code") == 100050:
        return False
    return True

print("=== SoftEdge (preprocessor 5) ===")
for k in ["softedge", "hed", "pidinet", "softedge_hed", "softedge_pidinet", "hedSafe", "hedBase"]:
    if test_key(5, k):
        print(f"FOUND SOFTEDGE KEY: {k}")
        break
        
print("=== Lineart (preprocessor 32) ===")
for k in ["lineart", "lineartStandard", "lineart_standard", "lineart_anime", "lineartAnime", "lineartRealistic", "lineartBase"]:
    if test_key(32, k):
        print(f"FOUND LINEART KEY: {k}")
        break