#!/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-v4",  # v4=2026-07-21 bump: 刷新冻结指令(部署纪律/识图知识更新)
        "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 = 120  # 真卡死才提示（普通 LLM 冷启动 20-100s 不应触发）
GATEWAY_URL = ""
GATEWAY_API_KEY = ""
GATEWAY_SESSION_ID = ""
GATEWAY_DEADLINE_SECONDS = 600
CALL_HERMES_TIMEOUT = 600  # agent 带工具调用（读文件/查数据）需几分钟
FALLBACK_REPLY = "请稍等，我在处理..."

# ── 图片 OCR 配置（截图消息 → SenseNova vision）──
OCR_CONFIG_FILE = "/home/hmo/.config/mofin/ocr_config.json"
OCR_TIMEOUT = 90
IMAGE_EXTS = ('.jpg', '.jpeg', '.png', '.webp', '.gif')


def _load_ocr_config():
    """从 /home/hmo/.config/mofin/ocr_config.json 读取 SenseNova 配置。
    文件不存在或未配置 → 返回 (None, None, None)，图片功能降级为仅提示收到。"""
    try:
        with open(OCR_CONFIG_FILE, 'r', encoding='utf-8') as f:
            cfg = json.load(f)
        return cfg.get('key'), cfg.get('base_url'), cfg.get('model', 'sensenova-6.7-flash-lite')
    except Exception:
        return None, None, None


def _download_image(url, timeout=30, retries=3):
    """下载图片字节，带重试。失败返回 None。

    重试原因：客户端上传和发消息是并发的，bot 收到消息时上传可能还没写完，
    立即 GET 会 404（<100 字节错误页）。等 2 秒再试即可。
    """
    import urllib.request
    import ssl
    import time as _t
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    for attempt in range(1, retries + 1):
        try:
            req = urllib.request.Request(url, headers={'User-Agent': 'curl/8.5.0'})
            with urllib.request.urlopen(req, timeout=timeout, context=ctx) as resp:
                data = resp.read()
            if len(data) > 100:
                return data
            log.warning(f"图片下载第{attempt}次: 响应过小({len(data)}B, 可能上传未完成)，重试...")
        except Exception as e:
            log.warning(f"图片下载第{attempt}次失败: {str(e)[:100]}")
        if attempt < retries:
            _t.sleep(2)
    log.error(f"图片下载最终失败({retries}次尝试): {url}")
    return None


def _ocr_image(img_data, prompt=None):
    """调 SenseNova vision OCR。返回 (ok, text)。"""
    key, base, model = _load_ocr_config()
    if not (key and base):
        return False, "OCR未配置"
    import base64
    import urllib.request
    b64 = base64.b64encode(img_data).decode()
    text_prompt = prompt or "请识别这张图片中的所有文字内容，包括数字、股票名称、金额、日期。用中文回复。"
    payload = json.dumps({
        "model": model,
        "messages": [{
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}},
                {"type": "text", "text": text_prompt},
            ]
        }],
        "max_tokens": 1500,
    }).encode()
    req = urllib.request.Request(
        f"{base.rstrip('/')}/chat/completions", data=payload,
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {key}"})
    try:
        with urllib.request.urlopen(req, timeout=OCR_TIMEOUT) as resp:
            data = json.loads(resp.read().decode())
        msg = data.get('choices', [{}])[0].get('message', {})
        text = msg.get('content', '') or msg.get('reasoning', '')
        return True, text.strip() if text.strip() else "(OCR无内容)"
    except Exception as e:
        return False, f"OCR失败: {str(e)[:120]}"


def _is_image_url(url):
    u = url.lower().split('?')[0]
    return ('upload.yoin.fun' in u or '/upload/' in u or u.endswith(IMAGE_EXTS))


def _process_image_message(url):
    """下载+OCR 一张截图，返回注入 LLM 的上下文文本。"""
    img = _download_image(url)
    if not img:
        return (f"[老爸发来一张截图，bot 下载失败（HTTP错误），无法识别。"
                f"图片URL: {url}\n"
                f"注：bot 的识图管道是 SenseNova（商汤云OCR），不是小果也不是 tesseract。"
                f"下载失败可能是 URL 被截断或 ejabberd 上传过期，请老爸重发一次试试。")
    ok, text = _ocr_image(img)
    if not ok:
        return (f"[老爸发来一张截图，SenseNova OCR 识别失败: {text}。"
                f"注：识图管道是 SenseNova，如持续失败请检查 OCR key 配额。")
    return (f"[老爸发来一张截图，SenseNova OCR 识别内容如下]\n{text}\n"
            f"[截图内容结束] 请基于截图内容回应老爸。")

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

RECENT_SENT_MAX = 50

# ── XMPP 消息日志（dashboard 健康Tab「最近对话」数据源）──
try:
    sys.path.insert(0, '/home/hmo/MoFin')  # 绝对路径（红线#7），相对解析会指到 /
    from xmpp_logger import log_xmpp as _log_xmpp
    _LOG_XMPP_OK = True
except Exception as _e:
    _LOG_XMPP_OK = False
    def _log_xmpp(*a, **kw):
        pass
    print(f"[WARN] xmpp_logger import 失败，消息日志降级为 no-op: {_e}", file=sys.stderr)


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')
        self.register_plugin('xep_0066')  # OOB 附件（截图 URL）

    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()
            # body 为空时检查 OOB 附件（XEP-0066：截图/文件以 OOB url 送达）
            if not body:
                try:
                    oob_url = str(msg['oob']['url'] or '').strip()
                except Exception:
                    oob_url = ''
                if oob_url and _is_image_url(oob_url):
                    body = f"[IMAGE] {oob_url}"
                    log.info(f"📩 收到图片(OOB): from={msg['from']} url={oob_url}")
                else:
                    return
            # body 本身是上传 URL（部分客户端把 URL 直接放 body）
            elif body.startswith('http') and _is_image_url(body):
                log.info(f"📩 收到图片(URL body): from={msg['from']} url={body}")
                body = f"[IMAGE] {body}"
            log.info(f"📩 收到: from={msg['from']} type={msg['type']} body={body[:60]}")
            _log_xmpp('in', sender_str := str(msg['from']), f"{AGENT_NAME}@yoin.fun", body)
            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_xmpp('out', f"{AGENT_NAME}@yoin.fun", target, text)
                    log.info(f"  已发送到 {target}: {text[:80]}")
                except Exception as e:
                    _log_xmpp('out', f"{AGENT_NAME}@yoin.fun", target, text, status='error', error=str(e)[:150])
                    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)
            # 截图消息：下载 + OCR → 注入上下文
            if body.startswith('[IMAGE]'):
                img_url = body[len('[IMAGE]'):].strip()
                log.info(f"🖼️ 图片消息处理: {img_url[:80]}")
                body = _process_image_message(img_url)
                log.info(f"🖼️ OCR完成: {body[:100]}")
            reply = call_hermes(body)
            if not reply:
                # 网关不可用/超时：绝不静默吞掉用户消息——重试一次，仍失败则兜底告知
                log.warning("call_hermes 无响应（网关可能重启中），10s后重试一次")
                time.sleep(10)
                reply = call_hermes(body)
            if not reply:
                reply = "⚠️ 网关暂时不可用（可能在重启/网络波动），你的消息我已收到但未能处理。请几分钟后重发一次，或稍候我会自行恢复。"
                log.error("call_hermes 重试仍失败，已回兜底消息")
            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()