import os
import re
import time
import uuid
import requests
import curl_cffi
import json_repair as json
import puremagic


model_map = {
    "Qwen3-Max": "tongyi-qwen3-max-model",
    "Qwen3-Max-Thinking-Preview": "tongyi-qwen3-max-thinking",
    "Qwen3-Plus": "tongyi-qwen-plus-latest",
    "Qwen3-Coder": "tongyi-qwen3-coder",
    "Qwen3-VL-32B": "tongyi-qwen3-vl-32b",
    "Qwen3-VL-30B-A3B": "tongyi-qwen3-vl-30b-a3b-instruct",
    "Qwen3-Omni-Flash": "tongyi-qwen3-omni-flash",
    "Qwen3-Next-80B-A3B": "tongyi-qwen3-next-80b-a3b",
    "Qwen3-235B-A22B-2507": "tongyi-qwen3-235b-a22b-instruct-2507",
    "Qwen3-VL-235B-A22B": "tongyi-qwen3-vl-235b-a22b",
    "Qwen3-30B-A3B-2507": "tongyi-qwen3-30b-a3b-instruct-2507",
    "Qwen3-Coder-Flash": "tongyi-qwen3-coder-flash",
    "Qwen3.5-Plus":"",
}

def parse_sse_stream(sse_text):
    """
    解析原始SSE数据流，返回最后一个 complete 事件的 data 字段解析后的字典。
    如果未找到有效的 complete 事件或 data 为 "true"，则返回 None。
    """
    lines = sse_text.splitlines()
    last_data_dict = None
    i = 0
    while i < len(lines):
        line = lines[i].rstrip()
        if line.startswith('event:'):
            event_type = line[6:].strip()
            i += 1
            # 收集紧接着的 data 行（可能有多个 data 行？SSE 通常每事件一个 data 行）
            data_lines = []
            while i < len(lines) and lines[i].startswith('data:'):
                data_lines.append(lines[i][5:].strip())
                i += 1
            # 跳过可能存在的空行（事件分隔）
            while i < len(lines) and lines[i] == '':
                i += 1

            if event_type == 'complete' and data_lines:
                data_str = '\n'.join(data_lines)
                # 如果 data 是 "true" 则忽略（流结束标记）
                if data_str == 'true':
                    continue
                try:
                    last_data_dict = json.loads(data_str)
                except json.JSONDecodeError:
                    # 如果解析失败，继续尝试后续事件（理论上不应该发生）
                    continue
        else:
            i += 1
    return last_data_dict

def extract_sse_response(data_dict):
    """
    从 parse_sse_stream 返回的字典中提取 AI 回答、思考内容和 reqid。
    返回格式：{'text': ..., 'think': ..., 'reqid': ...}
    若某字段不存在，则对应值为空字符串。
    """
    result = {'text': '', 'think': '', 'reqid': ''}

    # 提取 reqid（从 communication 对象中）
    communication = data_dict.get('communication', {})
    result['reqid'] = communication.get('reqid', '')

    # 提取 messages 数组
    messages = data_dict.get('data', {}).get('messages', [])
    if not messages:
        return result

    # 最后一个 message 通常是最终回答
    last_msg = messages[-1]
    content = last_msg.get('content', '')
    # 移除开头的标记，如 [(multimodal_chat_think_1)]
    clean_content = re.sub(r'^\[\([^)]+\)\]', '', content).strip()
    result['text'] = clean_content

    # 提取思考内容（可能有多段，用换行拼接）
    think_parts = []
    meta_data = last_msg.get('meta_data', {})
    multi_load = meta_data.get('multi_load', [])
    for item in multi_load:
        if item.get('type') == 'multimodal_chat_think':
            think = item.get('content', {}).get('think_content', '')
            if think:
                think_parts.append(think)
    result['think'] = '\n'.join(think_parts).strip()
    return result



def parse_sse_data(sse_data):
    """
    解析SSE格式数据，提取JSON内容
    """
    messages = []
    sse = sse_data.strip().split('\n')
    for line in sse:
        line = re.findall(r"\{.*\}", line)
        if line:
            line = line[0]
        else:
            continue
        messages.append(json.loads(line))
    return messages


def extract_final_response(messages):
    """
    从解析后的消息中提取最终回复和相关信息
    """
    sessionId = ""
    parentMsgId = ""
    text = ""
    think = ""

    for msg in messages:
        # 提取最终文本回复
        if 'contents' in msg:
            msgList = msg.get("contents")
            for i in msgList:
                if i and type(i) is dict:
                    if i.get("incremental"):
                        if i.get("contentType") == "think":
                            content = i.get("content", dict())
                            think += json.loads(content).get("content", "")
                        elif i.get("contentType") == "text":
                            text += i.get("content", "")
        if 'sessionId' in msg:
            sessionId = msg.get("sessionId")
        if 'parentMsgId' in msg:
            parentMsgId = msg.get("parentMsgId")

    return {
        'text': text,
        'think': think,
        'sessionId': sessionId,
        'parentMsgId': parentMsgId
    }


def qwen(cookies, msg, model="tongyi-qwen3-max-thinking", sessionId="", parentMsgId="", imgPath="", timeout=120):
    x_deviceid = str(uuid.uuid4().hex)
    x_xsrf_token = str(uuid.uuid4())
    requestId = str(uuid.uuid4()).replace('-', '')

    json_data = {
        'sessionId': sessionId,
        'sessionType': 'text_chat',
        'parentMsgId': parentMsgId,
        'model': '',
        'mode': 'chat',
        'userAction': 'new_top',
        'actionSource': '',
        'contents': [
            {
                'content': msg,
                'contentType': 'text',
                'role': 'user',
                'ext': {
                    'deepThink': True,
                },
            },
        ],
        'action': 'next',
        'requestId': requestId,
        'params': {
            'deepThink': True,
            'specifiedModel': model,
            'lastUseModelList': [
                model,
            ],
            'recordModelName': model,
            'bizSceneInfo': {},
        },
    }
    if imgPath:
        headers = {
            'accept': 'application/json, text/plain, */*',
            'accept-language': 'zh-CN,zh;q=0.9,en;q=0.8,zh-TW;q=0.7',
            'content-type': 'application/json',
            'origin': 'https://www.qianwen.com',
            'priority': 'u=1, i',
            'referer': 'https://www.qianwen.com/chat',
            'sec-ch-ua': '"Chromium";v="142", "Google Chrome";v="142", "Not_A Brand";v="99"',
            'sec-ch-ua-mobile': '?0',
            'sec-ch-ua-platform': '"Windows"',
            'sec-fetch-dest': 'empty',
            'sec-fetch-mode': 'cors',
            'sec-fetch-site': 'same-site',
            'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/142.0.0.0 Safari/537.36',
            'x-deviceid': x_deviceid,
            'x-platform': 'pc_tongyi',
            'x-xsrf-token': x_xsrf_token,
        }
        filename = os.path.basename(imgPath)
        response = curl_cffi.post('https://api.qianwen.com/dialog/uploadToken', cookies=cookies, headers=headers,
                                  json={'source': 'dialogue'}, impersonate="chrome124")
        Data = response.json()
        data = Data.get('data')
        OSSAccessKeyId = data.get('accessId')
        policy = data.get('policy')
        signature = data.get('signature')
        dir = data.get('dir')
        key = f"{dir}{filename}"
        print(puremagic.from_file(imgPath, mime=True))
        files = {
            'OSSAccessKeyId': (None, OSSAccessKeyId),
            'policy': (None, policy),
            'signature': (None, signature),
            'key': (None, key),
            'dir': (None, dir),
            'success_action_status': (None, '200'),
            'file': (filename, open(imgPath, "rb"), puremagic.from_file(imgPath, mime=True)),
        }
        headers = {
            'Accept': '*/*',
            'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8,zh-TW;q=0.7',
            'Connection': 'keep-alive',
            'Origin': 'https://www.qianwen.com',
            'Referer': 'https://www.qianwen.com/chat',
            'Sec-Fetch-Dest': 'empty',
            'Sec-Fetch-Mode': 'cors',
            'Sec-Fetch-Site': 'cross-site',
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/142.0.0.0 Safari/537.36',
            'X-Requested-With': 'XMLHttpRequest',
            'sec-ch-ua': '"Chromium";v="142", "Google Chrome";v="142", "Not_A Brand";v="99"',
            'sec-ch-ua-mobile': '?0',
            'sec-ch-ua-platform': '"Windows"',
        }

        requests.post('https://tongyi-main.oss-accelerate.aliyuncs.com/', headers=headers, files=files)

        imgData = {
            'fileKey': filename,
            'fileType': 'image',
            'dir': dir,
            'source': 'dialogue',
        }
        headers = {
            'accept': 'application/json, text/plain, */*',
            'accept-language': 'zh-CN,zh;q=0.9,en;q=0.8,zh-TW;q=0.7',
            'content-type': 'application/json',
            'origin': 'https://www.qianwen.com',
            'priority': 'u=1, i',
            'referer': 'https://www.qianwen.com/chat/1dd5c793985941c48d974152c1a02937',
            'sec-ch-ua': '"Chromium";v="142", "Google Chrome";v="142", "Not_A Brand";v="99"',
            'sec-ch-ua-mobile': '?0',
            'sec-ch-ua-platform': '"Windows"',
            'sec-fetch-dest': 'empty',
            'sec-fetch-mode': 'cors',
            'sec-fetch-site': 'same-site',
            'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/142.0.0.0 Safari/537.36',
            'x-deviceid': x_deviceid,
            'x-platform': 'pc_tongyi',
            'x-xsrf-token': x_xsrf_token,

        }
        response = curl_cffi.post('https://api.qianwen.com/dialog/downloadLink', cookies=cookies, headers=headers,
                                  json=imgData, impersonate="chrome124")
        urlLink = response.json().get('data').get('url')
        json_data.get('contents').append({
            'role': 'user',
            'contentType': 'image',
            'content': urlLink,
        })

    headers = {
        'accept': 'text/event-stream',
        'accept-language': 'zh-CN,zh;q=0.9,en;q=0.8,zh-TW;q=0.7',
        'content-type': 'application/json',
        'origin': 'https://www.qianwen.com',
        'priority': 'u=1, i',
        'referer': 'https://www.qianwen.com/chat',
        'sec-ch-ua': '"Chromium";v="142", "Google Chrome";v="142", "Not_A Brand";v="99"',
        'sec-ch-ua-mobile': '?0',
        'sec-ch-ua-platform': '"Windows"',
        'sec-fetch-dest': 'empty',
        'sec-fetch-mode': 'cors',
        'sec-fetch-site': 'same-site',
        'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/142.0.0.0 Safari/537.36',
        'x-deviceid': x_deviceid,
        'x-platform': 'pc_tongyi',
        'x-xsrf-token': x_xsrf_token,
    }

    response = curl_cffi.post('https://api.qianwen.com/dialog/conversation', cookies=cookies, headers=headers,
                              json=json_data, impersonate="chrome124", timeout=timeout)

    data = parse_sse_data(response.text)
    result = extract_final_response(data)
    return result




def qwenPlus(cookies, msg, req_id=None, parent_req_id=None, timeout=120):
    x_deviceid = str(uuid.uuid4().hex)
    x_xsrf_token = str(uuid.uuid4())
    x_chat_id = str(uuid.uuid4()).replace('-', '')
    clt_acs_reqt = int(time.time() * 1000)

    headers = {
        'accept': 'application/json, text/event-stream, text/plain, */*',
        'accept-language': 'zh-CN,zh;q=0.9,en;q=0.8,zh-TW;q=0.7',
        'cache-control': 'no-cache',
        'clt-acs-caer': 'vrad',
        'clt-acs-reqt': clt_acs_reqt,
        'clt-acs-request-params': 'biz_id,chat_client,device,fr,pr,ut,la,tz,nonce,timestamp',
        'clt-acs-sign': 'mYxlzMHF92KIzfByJ15xA4dW3cfwsWpVopeFxl1HTaw=',
        'content-type': 'application/json',
        'eo-clt-acs-kp': 'tytk_hash:04984e3698b2305df6236327dd2edec7',
        'eo-clt-acs-ve': '1.0.0',
        'eo-clt-actkn': 'j5Mh/l6al5PMkZT2ZxzgjOgoH8LjQHmjQ+V6U07CA6LzC453vHlbAZBzYpRCHIEvehWhjxhncTTSnYSzAOQO74OazWwW9sSLpPhIU0q+dgBZY5HIkvhHoTc+Aj6ZYH7O7QzmTScGw/rhGzbGJnmO4NhcZjX0TOUSH8/1huBAKEmGoEA1tel8oOlyzq/FuUtm',
        'eo-clt-dvidn': 'eCy#AAOKaWvKlUBQ+PPGUeNnMFSJiFVGLUi8cc8jyEODUeXpZpRPlH4PjuwZnAiu6KfqgX0=',
        'eo-clt-sacsft': 'PhoUNyw0DC8GJDU+MzA/MwATHiEeHA4eGGo1FSIiJRgwMzA4MTQwNw==',
        'eo-clt-snver': 'lv',
        'origin': 'https://www.qianwen.com',
        'pragma': 'no-cache',
        'priority': 'u=1, i',
        'referer': 'https://www.qianwen.com/chat/ca6bd37f74014b2dbb03f0e01b85afc4',
        'sec-ch-ua': '"Not:A-Brand";v="99", "Google Chrome";v="145", "Chromium";v="145"',
        'sec-ch-ua-mobile': '?0',
        'sec-ch-ua-platform': '"Windows"',
        'sec-fetch-dest': 'empty',
        'sec-fetch-mode': 'cors',
        'sec-fetch-site': 'same-site',
        'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/145.0.0.0 Safari/537.36',
        'x-chat-id': x_chat_id,
        'x-deviceid': x_deviceid,
        'x-platform': 'pc_tongyi',
        'x-xsrf-token': x_xsrf_token,
        # 'cookie': 'UM_distinctid=19b37a6d7d79ad-005f2aee3b555d8-26061a51-384000-19b37a6d7d8d48; tongyi_guest_ticket=Qc2oDc8uU_yWEm8QE$qulunHsANBVSWh08baW3*1HpPBZwBSrMYeJOGpn9vcx7IPTjeNsu6hQ53y0; cna=/HjMIXlO1lwCARsuVMDFciN4; _ON_EXT_DVIDN=eCy#AAOKaWvKlUBQ+PPGUeNnMFSJiFVGLUi8cc8jyEODUeXpZpRPlH4PjuwZnAiu6KfqgX0=; _qk_bx_ck_v1=eyJkZXZpY2VJZCI6ImVDeSNBQU9LYVd2S2xVQlErUFBHVWVObk1GU0ppRlZHTFVpOGNjOGp5RU9EVWVYcFpwUlBsSDRQanV3Wm5BaXU2S2ZxZ1gwPSIsImRldmljZUZpbmdlcnByaW50IjoiZTY1N2Q1ZGM5ZmMxMmQwNjVlNzk5MzIyYzQ2N2E1MjYifQ==; tongyi_sso_ticket=ppERXt_atkhS4_SxLgUqoj0I*K7YojXq1HB9Os3pfbGzzg_8yn1gsBH1opRO0qadzXjCc$oyhqMN0; tongyi_sso_ticket_hash=tytk_hash:04984e3698b2305df6236327dd2edec7; _qw_bx_timeout_v1=1000; theme-mode=light; xlly_s=1; _qk_bx_um_v1=eyJ1ZCI6IlQyZ0Fza0daSjFUZldXQ2VfZkd5SHVtU2FTLWRmTkk1TFBfVFp1TE9VOHFjbWpnWi1vQ2ZGTXNSUjhSeG54Mk5nRGc9IiwiZXAiOjE3NzM2NzA2Mjg4NDF9; isg=BJ-fe5mcCuE0vA5rtvN1lZ5bLvMpBPOmEl6-nTFvY41_wN9CctG59vhSglC-2Mse; tfstk=gHcEoRTIoBdEtltmuqVygONadCNdL786cwc7rz0-yaq3t6slqlurRe35O_kzb0EQdWxdz7uYXXTLrLUzjrmfx_VuKP2zzmYJy7qhUVmbryYRZyVPU04m48LLVuqrV4DB5IOjJ2F8Zg-XGI_B3lDI4yq3xNYuzz_H5Jjl1kF7ZntX1O_gpHNoRGDTqR0gXz5hZymuIV4YjzX3Z0qgIr4fx7VoqFzgzrfuq6VujC4ujuVuZ00M7zZgr7VoqVYayAFnqDCabJYH_3TPxNYNyog3b_fotQH703-dNPra8JZm-Xy4wfl8Lo03b_xEYHhQ4lhDfdut0vmQW0R2b7ntE04rjGx7p4DZbyoB0gES92GzA2JHvVPTY5cn3afo7WrLUX2FYhrm9VcYsRODuVctA2hZPaAu5mZiJXVyiEFUtkVaWbtOtukiEXwLwMxbp4DZbyPc4zj8SFk-wvSlUJ4T7oTw7ajEUyDl0i2AeTewBPrXkEBReJVQ7oTTKTB8QRaacehG.',
    }

    params = {
        'biz_id': 'ai_qwen',
        'chat_client': 'h5',
        'device': 'pc',
        'fr': 'pc',
        'pr': 'qwen',
        'ut': '276801d9-fcdb-5535-6ed0-4dfc5a3344a2',
        'la': 'zh-CN',
        'tz': 'Asia/Shanghai',
        'nonce': '2sp1r1kfxyg',
        'timestamp': int(time.time() * 1000),
    }

    if parent_req_id:
        scene_param = 'continue_chat'
    else:
        parent_req_id = "0"
        scene_param = 'first_turn'
    if not req_id:

        req_id = str(uuid.uuid4()).replace('-', '')

    json_data = {
        'deep_search': '1',
        'req_id': req_id,
        'model': 'Qwen3.5-Plus',
        'scene': 'chat',
        'session_id': 'ca6bd37f74014b2dbb03f0e01b85afc4',
        'sub_scene': 'chat',
        'temporary': False,
        'messages': [
            {
                'content': msg,
                'mime_type': 'text/plain',
                'meta_data': {
                    'ori_query': msg,
                },
            },
        ],
        'from': 'default',
        'topic_id': '1248371eb5ce4604b54d05929c2db047',
        'parent_req_id': parent_req_id,
        'enable_search': True,
        'scene_param': scene_param,
        'chat_client': 'h5',
        'client_tm': int(time.time() * 1000),
        'protocol_version': 'v2',
        'biz_id': 'ai_qwen',
    }




    response = curl_cffi.post('https://chat2.qianwen.com/api/v2/chat', params=params, cookies=cookies, headers=headers, json=json_data, impersonate="chrome124",timeout=timeout)
    # logger.debug(response.text)

    result = parse_sse_stream(response.text)

    # logger.debug(result)
    data = extract_sse_response(result)
    data["req_id"] = req_id

    return data

    
    # with open("123.txt","w", encoding="utf-8") as f:
    #     f.write(response.text)



if __name__ == '__main__':
	# https://www.qianwen.com/
    # https://chat2.qianwen.com/api/v2/chat 在这个api下取得cookies
    cookies = {}

    # data = qwen(cookies, "我叫JB6")
    # print(data)

    # data = qwen(cookies, "我叫什么?", sessionId=data["sessionId"], parentMsgId=data["parentMsgId"])
    # print(data)

    # data = qwen(cookies, "提取文字", imgPath="abc.jpg")
    # print(data)

    # data = qwenPlus(cookies=cookies,msg="我叫JB6")
    # print(data)

    # data = qwenPlus(cookies=cookies,msg="我叫什么",req_id=data.get('req_id'),parent_req_id=data.get('parent_req_id'))
    # print(data)
