#!/usr/bin/env python3
"""
Core XMPP Agent — shared logic for zhiwei / mohe / xxm bots.
Imports by xmpp_zhiwei_bot.py / xmpp_mohe_bot.py with --agent flag.
"""
import os, sys, json, time, logging, threading, traceback
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse, parse_qs
from hashlib import md5

import slixmpp
from slixmpp import JID
import asyncio

# ── Per-agent configuration ──
PER_AGENT = {
    "mohe": {
        "jid": "mohe@yoin.fun",
        "password": "hermes123",
        "nick": "mohe",
        "http_port": 5808,
        "gateway_url": "http://localhost:8642/v1/chat/completions",
        "gateway_api_key": "hermes123",
        "session_id": "xmpp-mohe-v2",
        "name_cn": "莫荷",
        "mention": "@mohe/@莫荷",
    },
    "zhiwei": {
        "jid": "zhiwei@yoin.fun",
        "password": "2nw4psra",
        "nick": "zhiwei",
        "http_port": 5805,
        "gateway_url": "http://localhost:8643/v1/chat/completions",
        "gateway_api_key": "hermes123",
        "session_id": "xmpp-zhiwei-v3",
        "name_cn": "知微",
        "mention": "@知微/zhiwei",
    },
}
_DEFAULT_AGENT = "mohe"

# ── Module-level config (populated by _apply_config after agent detection) ──
AGENT_NAME = ""
XMPP_JID = ""
XMPP_PASSWORD = ""
MUC_ROOM = "coregroup@conference.yoin.fun"
MUC_NICK = ""
AGENT_MENTION = ""
HTTP_PORT = 5808
AGENT_NICK = ""
ACK_DELAY = 15
GATEWAY_URL = ""
GATEWAY_API_KEY = ""
GATEWAY_SESSION_ID = ""
GATEWAY_DEADLINE_SECONDS = 180
CALL_HERMES_TIMEOUT = 180
FALLBACK_REPLY = "请稍等，我在处理..."

# ── 全局队列 ──
_outbound_queue = []
_outbound_lock = threading.Lock()
_inbound_queue = []
_inbound_lock = threading.Lock()

RECENT_SENT_MAX = 50


def _rs(p):
    """Parse agent arg from sys.argv, returns agent name string."""
    agent = _DEFAULT_AGENT
    skip_next = False
    for i, a in enumerate(sys.argv[1:]):
        if skip_next:
            skip_next = False
            continue
        if a.startswith('--agent='):
            agent = a.split('=', 1)[1]
        elif a == '--agent' and i + 1 < len(sys.argv[1:]):
            agent = sys.argv[i + 2]
            skip_next = True
    return agent


agent, is_mohe = _rs(None), None

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s %(levelname)s %(message)s',
    stream=sys.stdout,
)
log = logging.getLogger('xmpp_agent')


def _apply_config(agent_name):
    """Set module-level config variables from PER_AGENT dict + env overrides."""
    global AGENT_NAME, XMPP_JID, XMPP_PASSWORD, MUC_NICK, AGENT_MENTION
    global HTTP_PORT, AGENT_NICK, GATEWAY_URL, GATEWAY_API_KEY, GATEWAY_SESSION_ID
    cfg = PER_AGENT.get(agent_name, PER_AGENT.get(_DEFAULT_AGENT, {}))
    AGENT_NAME = agent_name
    XMPP_JID = os.environ.get('XMPP_JID', cfg.get('jid', ''))
    XMPP_PASSWORD = os.environ.get('XMPP_PASSWORD', cfg.get('password', ''))
    MUC_NICK = os.environ.get('MUC_NICK', cfg.get('nick', agent_name))
    AGENT_MENTION = os.environ.get('AGENT_MENTION', cfg.get('mention', ''))
    HTTP_PORT = int(os.environ.get('HTTP_PORT', cfg.get('http_port', 5808)))
    AGENT_NICK = os.environ.get('AGENT_NICK', cfg.get('nick', agent_name))
    GATEWAY_URL = os.environ.get('GATEWAY_URL', cfg.get('gateway_url', ''))
    GATEWAY_API_KEY = os.environ.get('GATEWAY_API_KEY', cfg.get('gateway_api_key', ''))
    GATEWAY_SESSION_ID = os.environ.get('GATEWAY_SESSION_ID', cfg.get('session_id', ''))


# ── Periodic ACK task ──
class AckManager:
    def __init__(self):
        self._active = {}
        self._lock = threading.Lock()

    def start(self, session_id, to_jid, msg_body):
        """Record an active LLM analysis and schedule the ACK."""
        with self._lock:
            self._active[session_id] = {
                'to_jid': to_jid,
                'body': msg_body[:80],
                'started': time.time(),
                'acked': False,
            }

    def ack(self, session_id):
        with self._lock:
            self._active.pop(session_id, None)

    def tick(self, bot):
        now = time.time()
        to_send = []
        with self._lock:
            for sid, info in list(self._active.items()):
                if not info['acked'] and now - info['started'] >= ACK_DELAY:
                    info['acked'] = True
                    to_send.append((info['to_jid'], FALLBACK_REPLY))
        for jid, msg in to_send:
            try:
                bot.send_message(mto=jid, mbody=msg, mtype='chat')
            except Exception:
                pass


ack_mgr = AckManager()


# ── Slixmpp Bot ──
class XmppAgent(slixmpp.ClientXMPP):
    def __init__(self, jid, password, room, nick):
        super().__init__(jid, password)
        self._room = room
        self._nick = nick
        self._muc_joined = False
        self._recent_sent = []
        self._reconnecting = False  # 防递归重连
        self.add_event_handler('session_start', self.on_start)
        self.add_event_handler('message', self.on_msg)
        self.add_event_handler('disconnected', self.on_disconnect)
        self.register_plugin('xep_0030')
        self.register_plugin('xep_0045')
        self.register_plugin('xep_0199')

    async def on_start(self, event):
        self.send_presence()
        await self.get_roster()
        try:
            await self.plugin['xep_0045'].join_muc(self._room, self._nick)
            self._muc_joined = True
            log.info(f"{AGENT_NAME} XMPP 就绪 (已加入 {self._room})")
        except Exception as e:
            log.error(f"{AGENT_NAME} MUC加入失败: {e}")

    def on_disconnect(self, event):
        self._muc_joined = False
        log.info(f"{AGENT_NAME} XMPP 断开")
        if self._reconnecting:
            log.warning(f"{AGENT_NAME} 已在重连中，跳过递归重连")
            return
        self._reconnecting = True
        try:
            self.reconnect(wait=5.0, reason="断线自动重连")
        except Exception as e:
            log.warning(f"{AGENT_NAME} 重连失败: {e}")
        finally:
            self._reconnecting = False

    def on_msg(self, msg):
        if msg['type'] in ('chat', 'groupchat'):
            body = str(msg['body']).strip()
            if not body:
                return
            log.info(f"📩 收到: from={msg['from']} type={msg['type']} body={body[:60]}")
            if ('[executor]' in body and 'gateway_zhiwei' in body):
                log.info(f"过滤 executor 消息: {body[:80]}...")
                return
            sender = str(msg['from'])
            msg_type = msg['type']
            for s in self._recent_sent:
                if body[:50] in s or s in body[:50]:
                    return
            if msg_type == 'groupchat':
                nick = sender.split('/')[-1] if '/' in sender else ''
                if nick == AGENT_NICK:
                    return
                mention_list = AGENT_MENTION.replace('@', '').split('/')
                is_for_me = any(m in body for m in ['@' + m for m in mention_list] + mention_list)
                if not is_for_me:
                    return
            with _inbound_lock:
                _inbound_queue.append((sender, body, msg_type))

    def mark_sent(self, body: str):
        self._recent_sent.append(body[:80])
        if len(self._recent_sent) > RECENT_SENT_MAX * 2:
            self._recent_sent = self._recent_sent[-RECENT_SENT_MAX:]


# ── Deliver loop ──
def _deliver_loop(bot):
    global _outbound_queue
    while True:
        try:
            items = []
            with _outbound_lock:
                items, _outbound_queue = _outbound_queue[:], []
            for target, text, msg_type in items:
                try:
                    async def _send(to, body, mtype):
                        bot.send_message(mto=to, mbody=body, mtype=mtype)
                    asyncio.run(_send(target, text, msg_type))
                    bot.mark_sent(text)
                    log.info(f"  已发送到 {target}: {text[:80]}")
                except Exception as e:
                    log.error(f"  发送到 {target} 失败: {e}")
            time.sleep(0.3)
        except Exception as e:
            log.error(f"_deliver_loop error: {e}")
            time.sleep(1)


# ── Inbound processing loop ──
def _inbound_loop(bot):
    global _inbound_queue
    while True:
        try:
            time.sleep(0.2)
            with _inbound_lock:
                if not _inbound_queue:
                    continue
                sender, body, msg_type = _inbound_queue.pop(0)
            log.info(f"🔄 inbound处理: {body[:40]}")
            ack_mgr.start(body[:40], sender, body)
            reply = call_hermes(body)
            ack_mgr.ack(body[:40])
            if reply:
                with _outbound_lock:
                    _outbound_queue.append((sender, reply, 'chat'))
        except Exception as e:
            log.error(f"_inbound_loop error: {e}")
            time.sleep(1)


# ── HTTP SendHandler ──
class SendHandler(BaseHTTPRequestHandler):
    def do_POST(self):
        content_len = int(self.headers.get('Content-Length', 0))
        post_body = self.rfile.read(content_len)
        try:
            data = json.loads(post_body)
            target = data.get('to', '')
            text = data.get('body', '')
            if not target or not text:
                self.send_response(400)
                self.end_headers()
                self.wfile.write(b'{"error":"missing to or body"}')
                return
            if '修复失败' in text and 'gateway_zhiwei' in text:
                self.send_response(200)
                self.end_headers()
                self.wfile.write(b'{"ok":true,"filtered":true}')
                return
            msg_type = data.get('type', 'chat')
            if text:
                _outbound_queue.append((target, text, msg_type))
                self.send_response(200)
                self.end_headers()
                self.wfile.write(b'{"ok":true}')
                log.info(f"SendHandler: enqueued -> {target}: {text[:80]}")
        except Exception as e:
            traceback.print_exc()
            self.send_response(500)
            self.end_headers()
            self.wfile.write(str({'error': str(e)}).encode())

    def log_message(self, format, *args):
        pass


def _run_http_server():
    server = HTTPServer(('127.0.0.1', HTTP_PORT), SendHandler)
    log.info(f"HTTP SendHandler listening on 127.0.0.1:{HTTP_PORT}")
    server.serve_forever()


# ── call_hermes ──
def call_hermes(content: str, session_id=None) -> str:
    now_str = time.strftime("[%Y-%m-%d %H:%M %A]", time.localtime())
    timed_content = f"{now_str}\n{content}"
    payload = {
        'model': 'hermes-agent',
        'messages': [
            {'role': 'user', 'content': timed_content},
        ],
        'stream': False,
    }
    if session_id is None:
        session_id = GATEWAY_SESSION_ID
    headers = {
        'Content-Type': 'application/json',
        'X-Hermes-Session-Id': session_id,
    }
    if GATEWAY_API_KEY:
        headers['Authorization'] = f'Bearer {GATEWAY_API_KEY}'

    import urllib.request
    data_bytes = json.dumps(payload).encode('utf-8')
    req = urllib.request.Request(GATEWAY_URL, data=data_bytes, headers=headers, method='POST')
    try:
        resp = urllib.request.urlopen(req, timeout=CALL_HERMES_TIMEOUT)
        resp_data = json.loads(resp.read().decode('utf-8'))
        reply = ''
        if 'choices' in resp_data and len(resp_data['choices']) > 0:
            choice = resp_data['choices'][0]
            if 'message' in choice and 'content' in choice['message']:
                reply = choice['message']['content']
            elif 'delta' in choice and 'content' in choice['delta']:
                reply = choice['delta']['content']
        if not reply:
            reply = resp_data.get('response', '')
        if not reply:
            reply = str(resp_data)
        return reply.strip()
    except urllib.request.HTTPError as e:
        err_body = e.read().decode('utf-8', errors='replace')
        log.error(f"call_hermes HTTP {e.code}: {err_body[:200]}")
        return ''
    except Exception as e:
        log.error(f"call_hermes error: {type(e).__name__}: {e}")
        return ''


# ── ACK manager tick ──
def _ack_tick(bot):
    while True:
        try:
            ack_mgr.tick(bot)
            time.sleep(5)
        except Exception:
            time.sleep(5)


# ── Main ──
def main():
    global is_mohe
    agent_name = _rs(None)
    _apply_config(agent_name)
    is_mohe = (agent_name == 'mohe')

    log.info(f"Starting XMPP Agent: {agent_name} (mohe={is_mohe})")
    log.info(f"  JID={XMPP_JID}  HTTP_PORT={HTTP_PORT}")
    log.info(f"  GATEWAY={GATEWAY_URL}")
    log.info(f"  SESSION_ID={GATEWAY_SESSION_ID}")

    bot = XmppAgent(XMPP_JID, XMPP_PASSWORD, MUC_ROOM, MUC_NICK)
    bot.connect(host='127.0.0.1', port=5222)
    t_deliver = threading.Thread(target=_deliver_loop, args=(bot,), daemon=True)
    t_deliver.start()
    t_inbound = threading.Thread(target=_inbound_loop, args=(bot,), daemon=True)
    t_inbound.start()
    t_ack = threading.Thread(target=_ack_tick, args=(bot,), daemon=True)
    t_ack.start()
    t_http = threading.Thread(target=_run_http_server, daemon=True)
    t_http.start()

    bot.loop.run_forever()


if __name__ == '__main__':
    main()