import os
import time
import logging
from aiohttp import web
from wechatpy.crypto import WeChatCrypto
from wechatpy.exceptions import InvalidSignatureException
from wechatpy import parse_message
from wechatpy.replies import TextReply
from anthropic import AsyncAnthropic

# --- 配置区域 ---
CONFIG = {
    "WOA_TOKEN": os.environ['WOA_TOKEN'],          # 微信公众号后台设置的 Token
    "WOA_APPID": os.environ['WOA_APPID'],          # kaiyunayun 公众号 AppID
    "WOA_AESKEY": os.environ['WOA_AESKEY']         # 如果开启加密模式则填写，否则留空
}

# 初始化日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 初始化客户端
claude_client = AsyncAnthropic(api_key=CONFIG["ANTHROPIC_API_KEY"])
crypto = None
if CONFIG["WECHAT_ENCODING_AES_KEY"]:
    crypto = WeChatCrypto(CONFIG["WECHAT_TOKEN"], CONFIG["WECHAT_ENCODING_AES_KEY"], CONFIG["WECHAT_APP_ID"])

# --- 内存数据库 (生产环境请替换为 Redis) ---
# 结构: { openid: [ {"role": "user", "content": "..."}, {"role": "assistant", "content": "..."} ] }
user_sessions = {}

def get_system_prompt():
    return "你是一个运行在微信公众号里的智能助手，回答要简洁、友好，适合手机端阅读。"

async def get_claude_reply(openid, user_content):
    """调用 Claude API 获取回复"""
    
    # 1. 初始化或获取用户会话历史
    if openid not in user_sessions:
        user_sessions[openid] = []
    
    history = user_sessions[openid]
    
    # 2. 添加当前用户消息
    history.append({"role": "user", "content": user_content})
    
    # 3. 限制历史记录长度 (防止 token 超限)
    if len(history) > CONFIG["MAX_HISTORY"] * 2:
        history = history[-CONFIG["MAX_HISTORY"] * 2:]
        user_sessions[openid] = history

    try:
        # 4. 调用 Claude
        response = await claude_client.messages.create(
            model="claude-3-5-sonnet-20241022",
            max_tokens=1024,
            system=get_system_prompt(),
            messages=history
        )
        
        ai_reply = response.content[0].text
        
        # 5. 保存 AI 回复到历史
        history.append({"role": "assistant", "content": ai_reply})
        
        return ai_reply

    except Exception as e:
        logger.error(f"Claude API Error: {e}")
        return "抱歉，大脑暂时短路了，请稍后再试。"

# --- HTTP 路由处理 ---

async def handle_wechat(request):
    """处理微信回调 (GET 验证 & POST 消息)"""
    query = request.query
    
    # 1. GET 请求：微信服务器验证
    if request.method == "GET":
        signature = query.get("signature", "")
        timestamp = query.get("timestamp", "")
        nonce = query.get("nonce", "")
        echostr = query.get("echostr", "")
        
        try:
            # 如果开启了加密模式，使用 crypto.check_signature，否则手动验证
            if crypto:
                crypto.check_signature(signature, timestamp, nonce)
            else:
                # 手动验证逻辑 (简单版)
                import hashlib
                tmp_list = sorted([CONFIG["WECHAT_TOKEN"], timestamp, nonce])
                tmp_str = "".join(tmp_list)
                hash_str = hashlib.sha1(tmp_str.encode('utf-8')).hexdigest()
                if hash_str != signature:
                    raise InvalidSignatureException()
            
            return web.Response(text=echostr)
        except InvalidSignatureException:
            logger.warning("Invalid signature from WeChat")
            return web.Response(text="Invalid signature", status=403)

    # 2. POST 请求：接收用户消息
    elif request.method == "POST":
        body = await request.read()
        
        # 如果开启加密，先解密
        if crypto:
            msg_signature = query.get("msg_signature", "")
            timestamp = query.get("timestamp", "")
            nonce = query.get("nonce", "")
            try:
                body = crypto.decrypt_message(body, msg_signature, timestamp, nonce)
            except Exception as e:
                logger.error(f"Decrypt error: {e}")
                return web.Response(text="error")

        # 解析消息
        try:
            msg = parse_message(body)
        except Exception as e:
            logger.error(f"Parse message error: {e}")
            return web.Response(text="success") # 即使解析失败也返回 success 避免微信重试风暴

        # 只处理文本消息
        if msg.type == 'text':
            openid = msg.source
            user_content = msg.content
            
            logger.info(f"User {openid}: {user_content}")
            
            # 异步调用 Claude
            reply_text = await get_claude_reply(openid, user_content)
            
            # 构建回复
            reply = TextReply(content=reply_text, message=msg)
            response_xml = reply.render()
            
            # 如果开启加密，需要加密返回
            if crypto:
                response_xml = crypto.encrypt_message(response_xml, nonce, timestamp)
            
            return web.Response(text=response_xml, content_type='application/xml')
        
        # 非文本消息或事件推送，直接返回 success 避免重试
        return web.Response(text="success")

    return web.Response(status=405)

# --- 启动服务 ---

async def init_app():
    app = web.Application()
    app.router.add_route('*', '/wechat', handle_wechat) # 路径需与公众号后台配置一致
    return app

if __name__ == '__main__':
    app = init_app()
    logger.info(f"Starting WeChat Bot on port {CONFIG['PORT']}...")
    web.run_app(app, host='0.0.0.0', port=CONFIG['PORT'])