import requests
import time
import hmac
import hashlib
import uuid
import base64
import os
import sys
from dotenv import load_dotenv

# 设置 UTF-8 输出编码（Windows 兼容）
if sys.platform == 'win32':
    sys.stdout.reconfigure(encoding='utf-8')

load_dotenv()

class LibLibControlNet:
    def __init__(self):
        self.ak = os.getenv("LIBLIBAI_ACCESS_KEY") or "你的AccessKey"
        self.sk = os.getenv("LIBLIBAI_SECRET_KEY") or "你的SecretKey"
        self.domain = "https://openapi.liblibai.cloud"
        # 先替换为更通用的模板ID（避免旧ID失效）
        # 文生图模板（必须是 1.5 体系的模板）
        self.TEMPLATE_UUID = "e10adc3949ba59abbe56e057f20f883e"

        # SD1.5 底模
        self.CHECKPOINT_ID = "0ea388c7eb854be3ba3c6f65aac6bfd3"

        # SD1.5 Canny（来自你给的官方表）
        self.CANNY_MODEL_ID = "b6806516962f4e1599a93ac4483c3d23"

    def _generate_signature(self, uri):
        ts = str(int(time.time() * 1000))
        nonce = uuid.uuid4().hex
        sign_str = f"{uri}&{ts}&{nonce}"
        dig = hmac.new(self.sk.encode(), sign_str.encode(), hashlib.sha1).digest()
        signature = base64.urlsafe_b64encode(dig).rstrip(b"=").decode()
        return f"{self.domain}{uri}?AccessKey={self.ak}&Timestamp={ts}&SignatureNonce={nonce}&Signature={signature}"

    def submit_canny_task(self, image_url, prompt):
        uri = "/api/generate/webui/text2img"
        url = self._generate_signature(uri)
        
        payload = {
            "templateUuid": self.TEMPLATE_UUID,
            "generateParams": {
                "checkPointId": self.CHECKPOINT_ID,
                "prompt": prompt,
                "negativePrompt": "lowres, bad anatomy, text, error, extra digit, fewer digits, cropped, worst quality, low quality",
                "sampler": 15,
                "steps": 20,
                "cfgScale": 7.0,  # 修复：改为浮点数（部分版本要求）
                "width": 512,
                "height": 512,
                "imgCount": 1,
                "seed": -1,
                "controlNet": [
                    {
                        "unitOrder": 1,
                        "sourceImage": image_url,
                        "width": 512,
                        "height": 512,
                        "preprocessor": 1,
                        "annotationParameters": {
                            "canny": {
                                "preprocessorResolution": 512,
                                "lowThreshold": 100,
                                "highThreshold": 200
                            }
                        },
                        "model": self.CANNY_MODEL_ID,
                        "controlWeight": 1.0,
                        "startingControlStep": 0.0,
                        "endingControlStep": 1.0,
                        "pixelPerfect": 1,
                        "controlMode": 0,
                        "resizeMode": 1
                    }
                ]
            }
        }

        try:
            print(f"📤 提交ControlNet生图任务...")
            print(f"🔍 请求URL: {url}")
            print(f"🔍 请求Payload: {payload}")  # 打印完整参数
            resp = requests.post(
                url,
                headers={"Content-Type": "application/json"},
                json=payload,
                timeout=10
            )
            # 打印完整响应（关键调试信息）
            print(f"🔍 响应状态码: {resp.status_code}")
            print(f"🔍 响应内容: {resp.text}")
            
            result = resp.json()
            if result.get("code") == 0:
                task_id = result["data"]["generateUuid"]
                print(f"✅ 任务提交成功！任务ID: {task_id}")
                return task_id
            else:
                print(f"❌ 任务提交失败: {result.get('msg')}")
                return None
        except Exception as e:
            print(f"❌ 提交任务异常: {str(e)}")
            return None

    def poll_task_status(self, task_id):
        uri = "/api/generate/webui/status"
        max_retry = 60
        retry_count = 0

        while retry_count < max_retry:
            try:
                url = self._generate_signature(uri)
                resp = requests.post(
                    url,
                    headers={"Content-Type": "application/json"},
                    json={"generateUuid": task_id},
                    timeout=10
                )
                resp.raise_for_status()
                result = resp.json()

                if result.get("code") != 0:
                    print(f"❌ 查询状态失败: {result.get('msg')}")
                    return None

                data = result.get("data", {})
                status = data.get("generateStatus")
                if status == 5:
                    img_url = data["images"][0]["imageUrl"]
                    print(f"🎉 生图成功！最终图片URL: {img_url}")
                    return img_url
                elif status in [6, 7]:
                    print(f"❌ 任务失败: {data.get('generateMsg', '未知错误')}")
                    return None
                else:
                    print(f"⏳ 任务处理中（状态码: {status}），已等待{retry_count*5}秒...")
            except Exception as e:
                print(f"❌ 轮询异常: {str(e)}")
            
            retry_count += 1
            time.sleep(5)
        
        print(f"❌ 任务轮询超时（超过5分钟）")
        return None

if __name__ == "__main__":
    client = LibLibControlNet()
    # 替换为更稳定的公开测试图
    PUBLIC_IMAGE_URL = "https://liblibai-airship-temp.oss-cn-beijing.aliyuncs.com/aliyun-cn-prod/73ed6ae42b144d21bf566e05b5a6c138.png"
    PROMPT = "simple white line art, cat, black background, 512x512"  # 简化提示词（避免过长）
    
    task_id = client.submit_canny_task(PUBLIC_IMAGE_URL, PROMPT)
    if task_id:
        client.poll_task_status(task_id)