from flask import Flask, request, make_response
import hashlib
import time
import xml.etree.ElementTree as ET
import os
import json
from openai import OpenAI
from dotenv import load_dotenv
from markdown import markdown
import re
import threading
import logging
from datetime import datetime
import asyncio
from concurrent.futures import ThreadPoolExecutor
import queue
import uuid
import base64
from Crypto.Cipher import AES
import struct
import random
import string
import requests

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('wechat_service.log'),
        logging.StreamHandler()
    ]
)

load_dotenv()

app = Flask(__name__)

# Configuration
TOKEN = os.getenv('TOKEN')
ENCODING_AES_KEY = os.getenv('ENCODING_AES_KEY')
APPID = os.getenv('APPID')
APPSECRET = os.getenv('APPSECRET')  # 新增 AppSecret
API_KEY = os.getenv("API_KEY")
BASE_URL = os.getenv("OPENAI_BASE_URL")
IMAGE_MODEL_URL = os.getenv("IMAGE_MODEL_URL")
IMAGE_MODEL_KEY = os.getenv("IMAGE_MODEL_KEY")

client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
executor = ThreadPoolExecutor(max_workers=10)


# Define tools for image generation
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "generate_image",
            "description": "Generate an image based on text description",
            "parameters": {
                "type": "object",
                "properties": {
                    "prompt": {
                        "type": "string",
                        "description": "The description of the image to generate"
                    }
                },
                "required": ["prompt"]
            }
        }
    }
]
class AccessTokenManager:
    def __init__(self):
        self._access_token = None
        self._expires_at = 0
        self._lock = threading.Lock()

    def get_token(self):
        with self._lock:
            now = time.time()
            # 提前5分钟刷新token，确保在调用时token都是有效的
            if self._access_token and now < (self._expires_at - 300):
                return self._access_token

            try:
                url = "https://api.weixin.qq.com/cgi-bin/token"
                params = {
                    "grant_type": "client_credential",
                    "appid": APPID,
                    "secret": APPSECRET
                }
                
                logging.info("开始获取新的access_token")
                response = requests.get(url, params=params)
                response.raise_for_status()
                result = response.json()
                
                if "access_token" not in result:
                    error_msg = f"获取access_token失败: {result}"
                    logging.error(error_msg)
                    raise ValueError(error_msg)
                
                self._access_token = result["access_token"]
                self._expires_at = now + result["expires_in"]
                logging.info("成功获取新的access_token")
                
                return self._access_token
                
            except Exception as e:
                error_msg = f"获取access_token时发生错误: {str(e)}"
                logging.error(error_msg)
                raise

    def refresh_token(self):
        with self._lock:
            self._access_token = None
            return self.get_token()
            
class WeChatCrypto:
    def __init__(self, key, app_id):
        self.key = base64.b64decode(key + '=')
        self.app_id = app_id

    def encrypt(self, text):
        random_str = ''.join(random.choices(string.ascii_letters + string.digits, k=16))
        text_bytes = text.encode('utf-8')
        msg_len = struct.pack('>I', len(text_bytes))
        message = random_str.encode('utf-8') + msg_len + text_bytes + self.app_id.encode('utf-8')
        pad_len = 32 - (len(message) % 32)
        message += chr(pad_len).encode('utf-8') * pad_len
        cipher = AES.new(self.key, AES.MODE_CBC, self.key[:16])
        encrypted = cipher.encrypt(message)
        return base64.b64encode(encrypted).decode('utf-8')

    def decrypt(self, encrypted_text):
        encrypted_data = base64.b64decode(encrypted_text)
        cipher = AES.new(self.key, AES.MODE_CBC, self.key[:16])
        decrypted = cipher.decrypt(encrypted_data)
        pad_len = decrypted[-1]
        if not isinstance(pad_len, int):
            pad_len = ord(pad_len)
        content = decrypted[16:-pad_len]
        msg_len = struct.unpack('>I', content[:4])[0]
        xml_content = content[4:msg_len + 4].decode('utf-8')
        app_id = content[msg_len + 4:].decode('utf-8')
        if app_id != self.app_id:
            raise ValueError('Invalid AppID')
        return xml_content

class AsyncResponse:
    def __init__(self):
        self.status = "processing"
        self.result = None
        self.error = None
        self.create_time = time.time()
        self.timeout = 3600
        self.response_type = "text"  # Can be "text" or "image"
        self.media_id = None  # For image responses

    def is_expired(self):
        return time.time() - self.create_time > self.timeout

class UserSession:
    def __init__(self):
        self.messages = [{"role": "system", "content": "你是HXIAO公众号的智能助手，这一个用来分享与学习人工智能的公众号，我们的目标是专注AI应用的简单研究与实践。致力于分享切实可行的技术方案，希望让复杂的技术变得简单易懂。也喜欢用通俗的语言来解释专业概念，让技术真正服务于每个学习者"}]
        self.pending_parts = []
        self.last_active = time.time()
        self.current_task = None
        self.response_queue = {}
        self.session_timeout = 3600

    def is_expired(self):
        return time.time() - self.last_active > self.session_timeout

    def cleanup_expired_tasks(self):
        expired_tasks = [
            task_id for task_id, response in self.response_queue.items()
            if response.is_expired()
        ]
        for task_id in expired_tasks:
            del self.response_queue[task_id]
            if self.current_task == task_id:
                self.current_task = None

class SessionManager:
    def __init__(self):
        self.sessions = {}
        self._lock = threading.Lock()
        self.crypto = WeChatCrypto(ENCODING_AES_KEY, APPID)

    def get_session(self, user_id):
        with self._lock:
            current_time = time.time()
            if user_id in self.sessions:
                session = self.sessions[user_id]
                if session.is_expired():
                    session = UserSession()
                else:
                    session.cleanup_expired_tasks()
            else:
                session = UserSession()
            session.last_active = current_time
            self.sessions[user_id] = session
            return session

    def clear_session(self, user_id):
        with self._lock:
            if user_id in self.sessions:
                self.sessions[user_id] = UserSession()

    def cleanup_expired_sessions(self):
        with self._lock:
            current_time = time.time()
            expired_users = [
                user_id for user_id, session in self.sessions.items()
                if session.is_expired()
            ]
            for user_id in expired_users:
                del self.sessions[user_id]
                logging.info(f"已清理过期会话: {user_id}")

def convert_markdown_to_wechat(md_text):
    if not md_text:
        return md_text
        
    md_text = re.sub(r'^# (.*?)$', r'【标题】\1', md_text, flags=re.MULTILINE)
    md_text = re.sub(r'^## (.*?)$', r'【子标题】\1', md_text, flags=re.MULTILINE)
    md_text = re.sub(r'^### (.*?)$', r'【小标题】\1', md_text, flags=re.MULTILINE)
    md_text = re.sub(r'\*\*(.*?)\*\*', r'『\1』', md_text)
    md_text = re.sub(r'\*(.*?)\*', r'「\1」', md_text)
    md_text = re.sub(r'`(.*?)`', r'「\1」', md_text)
    md_text = re.sub(r'^\- ', '• ', md_text, flags=re.MULTILINE)
    md_text = re.sub(r'^\d\. ', '○ ', md_text, flags=re.MULTILINE)
    md_text = re.sub(r'```[\w]*\n(.*?)```', r'【代码开始】\n\1\n【代码结束】', md_text, flags=re.DOTALL)
    md_text = re.sub(r'^> (.*?)$', r'▎\1', md_text, flags=re.MULTILINE)
    md_text = re.sub(r'^-{3,}$', r'—————————', md_text, flags=re.MULTILINE)
    md_text = re.sub(r'\[(.*?)\]\((.*?)\)', r'\1（\2）', md_text)
    md_text = re.sub(r'\n{3,}', '\n\n', md_text)
    
    return md_text

def verify_signature(signature, timestamp, nonce, token):
    items = [token, timestamp, nonce]
    items.sort()
    temp_str = ''.join(items)
    hash_sha1 = hashlib.sha1(temp_str.encode('utf-8')).hexdigest()
    return hash_sha1 == signature

def verify_msg_signature(msg_signature, timestamp, nonce, token, encrypt_msg):
    items = [token, timestamp, nonce, encrypt_msg]
    items.sort()
    temp_str = ''.join(items)
    hash_sha1 = hashlib.sha1(temp_str.encode('utf-8')).hexdigest()
    return hash_sha1 == msg_signature

def parse_xml_message(xml_content):
    root = ET.fromstring(xml_content)
    return {
        'content': root.find('Content').text if root.find('Content') is not None else '',
        'from_user': root.find('FromUserName').text,
        'to_user': root.find('ToUserName').text,
        'msg_id': root.find('MsgId').text if root.find('MsgId') is not None else '',
        'create_time': root.find('CreateTime').text,
        'msg_type': root.find('MsgType').text
    }

def generate_response_xml(to_user, from_user, content, response_type='text', media_id=None, encrypt_type='aes'):
    timestamp = str(int(time.time()))
    nonce = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
    
    if response_type == 'image' and media_id:
        xml_content = f'''
        <xml>
            <ToUserName><![CDATA[{to_user}]]></ToUserName>
            <FromUserName><![CDATA[{from_user}]]></FromUserName>
            <CreateTime>{timestamp}</CreateTime>
            <MsgType><![CDATA[image]]></MsgType>
            <Image>
                <MediaId><![CDATA[{media_id}]]></MediaId>
            </Image>
        </xml>
        '''
    else:
        formatted_content = convert_markdown_to_wechat(content)
        xml_content = f'''
        <xml>
            <ToUserName><![CDATA[{to_user}]]></ToUserName>
            <FromUserName><![CDATA[{from_user}]]></FromUserName>
            <CreateTime>{timestamp}</CreateTime>
            <MsgType><![CDATA[text]]></MsgType>
            <Content><![CDATA[{formatted_content}]]></Content>
        </xml>
        '''
    
    if encrypt_type == 'aes':
        encrypted = session_manager.crypto.encrypt(xml_content)
        signature_list = [TOKEN, timestamp, nonce, encrypted]
        signature_list.sort()
        msg_signature = hashlib.sha1(''.join(signature_list).encode('utf-8')).hexdigest()
        
        response_xml = f'''
        <xml>
            <Encrypt><![CDATA[{encrypted}]]></Encrypt>
            <MsgSignature><![CDATA[{msg_signature}]]></MsgSignature>
            <TimeStamp>{timestamp}</TimeStamp>
            <Nonce><![CDATA[{nonce}]]></Nonce>
        </xml>
        '''
    else:
        response_xml = xml_content
    
    response = make_response(response_xml)
    response.content_type = 'application/xml'
    return response
    
# 创建全局的token管理器实例
token_manager = AccessTokenManager()
def upload_image_to_wechat(image_data):
    """上传图片到微信服务器并获取media_id"""
    try:
        access_token = token_manager.get_token()
        upload_url = f'https://api.weixin.qq.com/cgi-bin/media/upload?access_token={access_token}&type=image'
        files = {'media': ('image.jpg', image_data, 'image/jpeg')}
        
        logging.info("开始上传图片到微信服务器")
        response = requests.post(upload_url, files=files)
        response.raise_for_status()
        result = response.json()
        
        if 'media_id' not in result:
            if 'errcode' in result and result['errcode'] == 40001:
                # access_token 可能过期，尝试刷新并重试
                logging.info("access_token已过期，正在刷新并重试")
                access_token = token_manager.refresh_token()
                upload_url = f'https://api.weixin.qq.com/cgi-bin/media/upload?access_token={access_token}&type=image'
                response = requests.post(upload_url, files=files)
                response.raise_for_status()
                result = response.json()
            
            if 'media_id' not in result:
                error_msg = f"上传图片失败: {result}"
                logging.error(error_msg)
                raise ValueError(error_msg)
        
        logging.info(f"图片上传成功，获取到media_id")
        return result['media_id']
        
    except Exception as e:
        error_msg = f"上传图片过程中发生错误: {str(e)}"
        logging.error(error_msg)
        raise

def process_long_running_task(messages):
    try:
        logging.info("开始调用AI服务")
        response = client.chat.completions.create(
            model="o3-mini",
            messages=messages,
            tools=TOOLS,
            tool_choice="auto",
            timeout=60
        )
        logging.info("AI服务响应成功")

        if response.choices[0].message.tool_calls:
            logging.info("检测到tool调用")
            tool_call = response.choices[0].message.tool_calls[0]
            if tool_call.function.name == "generate_image":
                logging.info("开始处理图片生成请求")
                args = json.loads(tool_call.function.arguments)
                image_response = requests.post(
                    IMAGE_MODEL_URL,
                    headers={
                        'Content-Type': 'application/json',
                        'Authorization': f'Bearer {IMAGE_MODEL_KEY}'
                    },
                    json={
                        "model": "grok-latest-image",
                        "messages": [{
                            "role": "user",
                            "content": args['prompt']
                        }]
                    }
                )
                image_response.raise_for_status()
                result = image_response.json()
                logging.info("图片生成成功，准备下载图片")

                # 从markdown格式中提取URL
                markdown_content = result['choices'][0]['message']['content']
                image_url = re.search(r'\!\[image\]\((.*?)\)', markdown_content).group(1)
                logging.info(f"提取到图片URL: {image_url}")
                
                img_response = requests.get(image_url)
                img_response.raise_for_status()
                media_id = upload_image_to_wechat(img_response.content)
                
                return {
                    "type": "image",
                    "media_id": media_id
                }

        logging.info("返回文本响应")
        return {
            "type": "text",
            "content": response.choices[0].message.content
        }
    except Exception as e:
        logging.error(f"API调用错误: {str(e)}")
        raise

def handle_async_task(session, task_id, messages):
    try:
        logging.info(f"开始处理异步任务: {task_id}")
        if task_id not in session.response_queue:
            return
            
        result = process_long_running_task(messages)
        
        if task_id in session.response_queue and not session.response_queue[task_id].is_expired():
            session.response_queue[task_id].status = "completed"
            session.response_queue[task_id].response_type = result.get("type", "text")
            if result["type"] == "image":
                session.response_queue[task_id].media_id = result["media_id"]
                session.response_queue[task_id].result = None
                messages.append({"role": "assistant", "content": "图片已生成"})
            else:
                session.response_queue[task_id].result = result["content"]
                messages.append({"role": "assistant", "content": result["content"]})
    except Exception as e:
        logging.error(f"异步任务处理失败: {str(e)}")
        if task_id in session.response_queue:
            session.response_queue[task_id].status = "failed"
            session.response_queue[task_id].error = str(e)

def generate_initial_response():
    return "您的请求正在处理中，请回复'查询'获取结果（生图需要时间）"

def split_message(message, max_length=500):
    return [message[i:i+max_length] for i in range(0, len(message), max_length)]

def append_status_message(content, has_pending_parts=False, is_processing=False):
    if "您的请求正在处理中" in content:
        return content + "\n\n-------------------\n发送'新对话'开始新的对话"
        
    status_message = "\n\n-------------------"
    if is_processing:
        status_message += "\n请回复'查询'获取结果"
    elif has_pending_parts:
        status_message += "\n当前消息已截断，发送'继续'查看后续内容"
    status_message += "\n发送'新对话'开始新的对话"
    return content + status_message

session_manager = SessionManager()

@app.route('/api/wx', methods=['GET', 'POST'])
def wechatai():
    if request.method == 'GET':
        signature = request.args.get('signature')
        timestamp = request.args.get('timestamp')
        nonce = request.args.get('nonce')
        echostr = request.args.get('echostr')
        
        if verify_signature(signature, timestamp, nonce, TOKEN):
            return echostr
        return 'error', 403
    
    try:
        encrypt_type = request.args.get('encrypt_type', '')
        
        if encrypt_type == 'aes':
            msg_signature = request.args.get('msg_signature')
            timestamp = request.args.get('timestamp')
            nonce = request.args.get('nonce')
            
            xml_tree = ET.fromstring(request.data)
            encrypted_text = xml_tree.find('Encrypt').text
            
            if not verify_msg_signature(msg_signature, timestamp, nonce, TOKEN, encrypted_text):
                return 'Invalid signature', 403
            
            decrypted_xml = session_manager.crypto.decrypt(encrypted_text)
            message_data = parse_xml_message(decrypted_xml)
        else:
            message_data = parse_xml_message(request.data)
        
        user_content = message_data['content'].strip()
        from_user = message_data['from_user']
        to_user = message_data['to_user']
        
        logging.info(f"收到用户({from_user})消息: {user_content}")
        session = session_manager.get_session(from_user)

        if user_content == '新对话':
            session_manager.clear_session(from_user)
            return generate_response_xml(
                from_user,
                to_user,
                append_status_message('已开始新的对话。请描述您的问题。'),
                encrypt_type=encrypt_type
            )

        if user_content == '继续':
            if session.pending_parts:
                next_part = session.pending_parts.pop(0)
                has_more = bool(session.pending_parts)
                return generate_response_xml(
                    from_user,
                    to_user,
                    append_status_message(next_part, has_more),
                    encrypt_type=encrypt_type
                )
            return generate_response_xml(
                from_user,
                to_user,
                append_status_message('没有更多内容了。请继续您的问题。'),
                encrypt_type=encrypt_type
            )

        if user_content == '查询':
            if session.current_task:
                task_response = session.response_queue.get(session.current_task)
                if task_response:
                    if task_response.is_expired():
                        del session.response_queue[session.current_task]
                        session.current_task = None
                        return generate_response_xml(
                            from_user,
                            to_user,
                            append_status_message('请求已过期，请重新提问。'),
                            encrypt_type=encrypt_type
                        )
                        
                    if task_response.status == "completed":
                        if task_response.response_type == "image":
                            logging.info("返回图片响应")
                            del session.response_queue[session.current_task]
                            session.current_task = None
                            return generate_response_xml(
                                from_user,
                                to_user,
                                "",
                                response_type="image",
                                media_id=task_response.media_id,
                                encrypt_type=encrypt_type
                            )
                        else:
                            response = task_response.result
                            del session.response_queue[session.current_task]
                            session.current_task = None
                            
                            if len(response) > 500:
                                parts = split_message(response)
                                first_part = parts.pop(0)
                                session.pending_parts = parts
                                return generate_response_xml(
                                    from_user,
                                    to_user,
                                    append_status_message(first_part, True),
                                    encrypt_type=encrypt_type
                                )
                            return generate_response_xml(
                                from_user,
                                to_user,
                                append_status_message(response),
                                encrypt_type=encrypt_type
                            )
                    elif task_response.status == "failed":
                        error_message = '处理过程中出现错误，请重新提问。'
                        del session.response_queue[session.current_task]
                        session.current_task = None
                        return generate_response_xml(
                            from_user,
                            to_user,
                            append_status_message(error_message),
                            encrypt_type=encrypt_type
                        )
                    else:
                        return generate_response_xml(
                            from_user,
                            to_user,
                            append_status_message('正在处理中，请稍后再次查询。（生图需要时间）', is_processing=True),
                            encrypt_type=encrypt_type
                        )
            return generate_response_xml(
                from_user,
                to_user,
                append_status_message('没有正在处理的请求。'),
                encrypt_type=encrypt_type
            )

        session.messages.append({"role": "user", "content": user_content})
        
        task_id = str(uuid.uuid4())
        session.current_task = task_id
        session.response_queue[task_id] = AsyncResponse()
        
        executor.submit(handle_async_task, session, task_id, session.messages.copy())
        
        logging.info("返回初始响应")
        return generate_response_xml(
            from_user,
            to_user,
            append_status_message(generate_initial_response(), is_processing=True),
            encrypt_type=encrypt_type
        )

    except Exception as e:
        logging.error(f"处理请求时出错: {str(e)}")
        return generate_response_xml(
            message_data['from_user'],
            message_data['to_user'],
            append_status_message('抱歉，系统暂时出现问题，请稍后重试。'),
            encrypt_type if 'encrypt_type' in locals() else ''
        )

def cleanup_sessions():
    while True:
        time.sleep(3600)  # 每小时清理一次
        try:
            session_manager.cleanup_expired_sessions()
        except Exception as e:
            logging.error(f"清理会话时出错: {str(e)}")

if __name__ == '__main__':
    cleanup_thread = threading.Thread(target=cleanup_sessions, daemon=True)
    cleanup_thread.start()
    
    app.run(host='0.0.0.0', port=7860, debug=True)