import hmac
from hashlib import sha1
import base64
import time
import uuid
import requests
import Config.Config


class LiblibUtil:
    def __init__(self):
        self.base_url = Config.Config.LIBLIB_URL
        self.access_key = Config.Config.LIBLIB_ACCESSKEY
        self.secret_key = Config.Config.LIBLIB_SECRETKEY
        self.timeout = 30

    def make_sign(self, uri):
        """生成API请求签名"""
        timestamp = str(int(time.time() * 1000))
        signature_nonce = str(uuid.uuid4())
        content = '&'.join((uri, timestamp, signature_nonce))
        
        digest = hmac.new(self.secret_key.encode(), content.encode(), sha1).digest()
        sign = base64.urlsafe_b64encode(digest).rstrip(b'=').decode()
        return sign, timestamp, signature_nonce

    def post_request(self, uri, payload):
        """发送POST请求到Liblib API"""
        try:
            sign, timestamp, signature_nonce = self.make_sign(uri)
            url = f'{self.base_url}{uri}?AccessKey={self.access_key}&Signature={sign}&Timestamp={timestamp}&SignatureNonce={signature_nonce}'
            
            headers = {'Content-Type': 'application/json'}
            response = requests.post(url, json=payload, headers=headers, timeout=self.timeout)
            response.raise_for_status()
            
            response_data = response.json()
            if response_data.get('code') == 0:
                return response_data.get('data')
            else:
                print(f"API错误: {response_data.get('message')}")
                return None
        except requests.exceptions.RequestException as e:
            print(f"请求异常: {str(e)}")
            return None
        except Exception as e:
            print(f"处理异常: {str(e)}")
            return None

    # ===== 从T1_VersionGet.py迁移的方法 =====
    def get_model_version_info(self, version_uuid):
        """获取模型版本信息"""
        uri = "/api/model/version/get"
        payload = {"versionUuid": version_uuid}
        
        model_info = self.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(self, template_uuid, generate_params):
        """调用text2img接口生成图片"""
        uri = "/api/generate/webui/text2img"
        payload = {
            "templateUuid": template_uuid,
            "generateParams": generate_params
        }
        return self.post_request(uri, payload)

    # ===== 从T2.py迁移的方法 =====
    def get_generation_status(self, generate_uuid):
        """查询生图任务状态和结果"""
        uri = "/api/generate/webui/status"
        payload = {"generateUuid": generate_uuid}
        return self.post_request(uri, payload)

    def wait_for_generation_completion(self, generate_uuid, check_interval=5, max_wait_time=300):
        """等待生图任务完成并获取最终结果"""
        start_time = time.time()
        
        while time.time() - start_time < max_wait_time:
            status_data = self.get_generation_status(generate_uuid)
            
            if status_data:
                status = status_data.get('status')
                progress = status_data.get('percentCompleted', 0) * 100
                print(f"生成进度：{progress:.2f}%")
                
                if status == 'completed':
                    print("生图任务已完成")
                    images = status_data.get('images', [])
                    if images:
                        print("生成的最终图片：")
                        for img in images:
                            print(f"- {img.get('imageUrl')}")
                    return status_data
                elif status == 'failed':
                    print(f"生图任务失败: {status_data.get('errorMessage', '未知错误')}")
                    return None
                else:
                    print(f"生图任务处理中，当前状态: {status}")
            
            time.sleep(check_interval)
        
        print(f"生图任务超时（{max_wait_time}秒）")
        return None