from .LiblibUtil import LiblibUtil


def get_model_version_info(version_uuid):
    """获取模型版本信息"""
    liblib_util = LiblibUtil()
    uri = "/api/model/version/get"
    payload = {"versionUuid": version_uuid}
    
    model_info = liblib_util.post_request(uri, payload)
    if model_info:
        return {
            'modelName': model_info.get('modelName'),
            'versionName': model_info.get('versionName'),
            'commercialUse': model_info.get('commercialUse'),
            'modelUrl': model_info.get('modelUrl')
        }
    return None


def generate_text_to_image(template_uuid, generate_params):
    """调用text2img接口生成图片"""
    liblib_util = LiblibUtil()
    uri = "/api/generate/webui/text2img"
    payload = {
        "templateUuid": template_uuid,
        "generateParams": generate_params
    }
    return liblib_util.post_request(uri, payload)


if __name__ == '__main__':
    # 测试获取模型版本信息
    print("===== 测试获取模型版本信息 =====")
    version_uuid = "390e7df48aed45d4bc3f0bcbc89fa44c"
    model_info = get_model_version_info(version_uuid)
    
    if model_info:
        print(f"模型名称：{model_info['modelName']}")
        print(f"版本名称：{model_info['versionName']}")
        print(f"是否允许商业用途：{model_info['commercialUse']}")
        print(f"模型下载地址：{model_info['modelUrl']}")
    else:
        print("获取模型信息失败")
    
    # 测试生成图片
    print("\n===== 测试生成图片 =====")
    template_uuid = "7d888009f81d4252a7c458c874cd017f"
    generate_params = {
        "checkPointId": "0ea388c7eb854be3ba3c6f65aac6bfd3",
        "prompt": "Asian portrait,A young woman wearing a green baseball cap",
        "width": 768,
        "height": 1152,
        "sampler": 20,
        "steps": 35,
        "cfgScale": 2.0,
        "imgCount": 1
    }
    
    image_data = generate_text_to_image(template_uuid, generate_params)
    if image_data:
        print(f"图片生成成功，任务UUID: {image_data.get('generateUuid')}")
    else:
        print("图片生成失败")