# -*- coding: utf-8 -*-
import os
import json
import uuid
import logging
import sseclient
import requests
import json
import sys



from pathlib import Path

import requests
from qcloud_cos import CosConfig, CosS3Client
from tencentcloud.common.common_client import CommonClient
from tencentcloud.common import credential
from tencentcloud.common.exception.tencent_cloud_sdk_exception import TencentCloudSDKException
from tencentcloud.common.profile.client_profile import ClientProfile
from tencentcloud.common.profile.http_profile import HttpProfile




# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


def chat_sse(request_data: dict, sse_url: str) -> None:
    """
    通过 SSE 协议与对话服务进行交互。

    :param request_data: 发送到对话服务的请求数据。
    :param sse_url: SSE 服务的 URL。
    """
    logging.info("\n ========== Start Chat ========== \n")
    logging.info(f'SSE Request Data: {request_data}')

    response = requests.post(sse_url, data=json.dumps(request_data), stream=True, headers={"Accept": "text/event-stream"})
    client = sseclient.SSEClient(response)

    for event in client.events():
        data = json.loads(event.data)
        if event.event == "reply":
            if data["payload"]["is_from_self"]:  # 自己发出的消息，需要检查是否包含敏感内容
                logging.info(f'Event: {event.event}, Data: {event.data}')
                if data["payload"]["is_evil"]:
                    logging.warning("问题或答案包含敏感内容，未通过审核！")
            elif data["payload"]["is_final"]:  # 服务端消息传输完成
                logging.info(f'Event: {event.event}, Data: {event.data}')
                logging.info("问答完成！")
        else:
            logging.info(f'Event: {event.event}, Data: {event.data}')


def generate_markdown(content: str, cos_final_url: str) -> str:
    """
    生成包含图片链接的 Markdown 格式文本。

    :param content: 图片描述文本。
    :param cos_final_url: 图片的 URL 链接。
    :return: Markdown 格式的字符串。
    """
    markdown_text = f"{content}![]({cos_final_url})"
    return markdown_text


def document_parse(request_data: dict) -> tuple:
    """
    实时文档解析。

    :param request_data: 发送到文档解析服务的请求数据。
    :return: 文档 ID、状态和错误信息。
    """
    headers = {'Content-Type': 'application/json'}
    response = requests.post(DocParseUrl, data=json.dumps(request_data), headers=headers)
    response.raise_for_status()

    content_str = response.content.decode('utf-8')
    client = sseclient.SSEClient(response)
    for event in client.events():
        data = json.loads(event.data)
        doc_id = data['payload']['doc_id']
        status = data['payload']['status']
        error_message = data['payload']['error_message']
        is_final = data['payload']['is_final']

        if is_final:
            logging.info(f"Document ID: {doc_id}")
            return doc_id, status, error_message
        else:
            logging.info(data)


def get_temporary_credentials(bot_biz_id: str, file_type: str, is_public: bool, type_key: str) -> dict:
    """
    获取临时密钥。

    :param bot_biz_id: 业务 ID。
    :param file_type: 文件类型。
    :param is_public: 是否公开。
    :param type_key: 类型键。
    :return: 临时密钥相关数据。
    """
    try:
        cred = credential.Credential(SecretID, SecretKey)
        http_profile = HttpProfile()
        http_profile.endpoint = EndPoint
        client_profile = ClientProfile()
        client_profile.httpProfile = http_profile

        params = {
            "BotBizId": bot_biz_id,
            "FileType": file_type,
            "TypeKey": type_key,
            "IsPublic": is_public
        }
        common_client = CommonClient("lke", "2023-11-30", cred, Region, profile=client_profile)
        response = common_client.call_json("DescribeStorageCredential", params)
        credentials = response['Response']['Credentials']
        upload_path = response['Response']['UploadPath']
        bucket = response['Response']['Bucket']
        region = response['Response']['Region']
        cos_type = response['Response']['Type']

        logging.info("======== DescribeStorageCredential Success =======")
        logging.info(f"Temporary Secret ID: {credentials['TmpSecretId']}")
        logging.info(f"Temporary Secret Key: {credentials['TmpSecretKey']}")
        logging.info(f"Token: {credentials['Token']}")
        logging.info(f"Upload Path: {upload_path}")
        logging.info(f"Bucket: {bucket}")
        logging.info(f"Region: {region}")
        logging.info(f"Type: {cos_type}")

        return {
            "TmpSecretId": credentials['TmpSecretId'],
            "TmpSecretKey": credentials['TmpSecretKey'],
            "Token": credentials['Token'],
            "UploadPath": upload_path,
            "Bucket": bucket,
            "Region": region,
            "Type": cos_type
        }
    except Exception as err:
        logging.error(err)
        raise


def upload_file_to_cos(file_path: str, credentials: dict) -> dict:
    """
    将文件上传到 COS。

    :param file_path: 文件路径。
    :param credentials: 临时密钥相关数据。
    :return: 文件的 COS URL。
    """
    config = CosConfig(
        Region=credentials['Region'],
        SecretId=credentials['TmpSecretId'],
        SecretKey=credentials['TmpSecretKey'],
        Token=credentials['Token'],
        Scheme='https'
    )
    client = CosS3Client(config)

    file_name = Path(file_path).name
    response = client.upload_file(
        Bucket=credentials['Bucket'],
        Key=credentials['UploadPath'],
        LocalFilePath=file_name,
        EnableMD5=False,
        progress_callback=None
    )

    logging.info(f"Upload Result: {response}")
    e_tag = response.get('ETag')
    cos_hash = response.get('x-cos-hash-crc64ecma')
    logging.info(f"ETag: {e_tag}, COS Hash: {cos_hash}")

    bucket_url = f"https://{credentials['Bucket']}.{credentials['Type']}.{credentials['Region']}.myqcloud.com"
    cos_final_url = f"{bucket_url}{credentials['UploadPath']}"
    logging.info(f"File URL: {cos_final_url}")
    return {
        "cos_final_url": cos_final_url,
        "e_tag": e_tag,
        "cos_hash": cos_hash
    }



if __name__ == "__main__":

    Region = "ap-jakarta"
    EndPoint = "lke.tencentcloudapi.com"
    SecretID = ""
    SecretKey = ""
    BotBizID = ""  # BotBizID通常是19位的一个数字,如果不知道如何获取请参考：https://www.tencentcloud.com/document/product/1254/69984 第三项
    BotAppKey = ""            # BotAppKey是一个8位的字符串,如果不知道如何获取请参考：https://www.tencentcloud.com/document/product/1254/69984 第三项
    TypeKeyRealtime = "realtime"
    DocParseUrl = "https://wss.lke.tencentcloud.com/v1/qbot/chat/docParse"
    SSEUrl      = "https://wss.lke.tencentcloud.com/v1/qbot/chat/sse"


    # 请使用 python 3.9，使用前需要安装依赖: pip3 install -r requirements.txt<br>
    # 图片路径
    # file_path = "./小楷.jpeg"
    # 文件路径
    file_path = "./致橡树.txt"
    file_name = Path(file_path).name
    file_ext = Path(file_path).suffix[1:]
    file_size = str(os.path.getsize(file_path))

    logging.info(f"File Path: {file_path}")
    logging.info(f"File Name: {file_name}")
    logging.info(f"File Extension: {file_ext}")
    logging.info(f"File Size: {file_size}")


    # 临时密钥的获取，请注意，参数组合不同得到的临时密钥权限不同，会影响后面上传cos和文件解析的结果。常见问题如： 上传cos报错403， 实时文档解析报错 Invalid-URL
	# 可参考 https://www.tencentcloud.com/document/product/1254/69981 的参数组合，或者在遇到需要上传文件的地方F12抓包DescribeStorageCredential接口，看下参数组合
    # 请注意，该场景为对话接口上传，is_public做了特殊判断，其他场景请参考上上面文档确定图片是否需要特殊处理
    # 请注意，对话是上传图片和知识库上传图片限制的图片格式和大小随着产品能力的迭代会增加或者扩大，请参考知识引擎页面调整图片的支持范围
    is_public = file_ext.lower() in ["jpg", "jpeg", "png", "bmp"]
    credentials = get_temporary_credentials(BotBizID, file_ext, is_public, TypeKeyRealtime)

    # 上传文件到 COS
    # 请注意，需要使用临时密钥建立cos_clint, 不同的语言方式不一样，使用其他语言可自行参考cos的sdk
    # 目前已知至少包含三个参数【secret_key,secret_id,token】, 该参数请从获取临时密钥接口返回的数据中获取
    cos_result = upload_file_to_cos(file_path, credentials)
    logging.info(f"Cos Result: {cos_result}")

    # 生成会话 ID
   # session_id很重要，请遵循规则生成，docParse传入的session_id需和对话时传入的session_id保持一致
    session_id = str(uuid.uuid4())

    if file_ext.lower() in ["jpg", "jpeg", "png", "bmp"]:
        # 图片处理逻辑, 图片的类型和大小限制的扩充，请关注知识引擎发版公告或者页面上的限制
        content = ""
        request_data = {
            "content": generate_markdown(content,cos_result['cos_final_url']),
            "bot_app_key": BotAppKey,
            "visitor_biz_id": session_id,
            "session_id": session_id,
            "visitor_labels": [],
            "request_id": session_id,
        }
        # 请注意，这个sseChat主要是快速验证效果，请接入方参考https://www.tencentcloud.com/document/product/1254/69978，结合自己的业务选择http sse或者websocket处理相关逻辑
		# 尤其注意，需要判断reply事件中的is_evil字段， is_evil=true,表示被风控拦截了，当前对用户的输入及大模型的输出都接入了风控策略
        chat_sse(request_data, SSEUrl)
    else:
         # 文档处理逻辑
        try:
            request_data = {
                "session_id": session_id,
                "request_id": session_id,
                "cos_bucket": credentials['Bucket'],
                "file_type": file_ext,
                "file_name": file_name,
                "cos_url": credentials['UploadPath'],
                "e_tag": cos_result["e_tag"],   
                "cos_hash": cos_result["cos_hash"],   
                "size": file_size,
                "bot_app_key": BotAppKey,
            }
            logging.info(f"Document Parse Request Data: {request_data}")
            doc_id, status, error_message = document_parse(request_data)
            if status == "FAILED":
                logging.error(f"Document Parse Failed! Status: {status}, Error Message: {error_message}")
                exit()

            sse_request_data = {
                "content": "",
                "bot_app_key": BotAppKey,
                "visitor_biz_id": session_id,
                "session_id": session_id,
                "visitor_labels": [],
                "request_id": session_id,
                "file_infos": [
                    {
                        "doc_id": doc_id,
                        "file_name": file_name,
                        "file_type": file_ext,
                        "file_size": file_size,
                        "file_url": cos_result["cos_final_url"]
                    }
                ]
            }
            # 请注意，这个sseChat主要是快速验证效果，请接入方参考https://www.tencentcloud.com/document/product/1254/69978，结合自己的业务选择http sse或者websocket处理相关逻辑
		    # 尤其注意，需要判断reply事件中的is_evil字段， is_evil=true,表示被风控拦截了，当前对用户的输入及大模型的输出都接入了风控策略
            chat_sse(sse_request_data, SSEUrl)

        except requests.exceptions.RequestException as e:
            logging.error(e)