hendrik/services/xmpp_client.py

import asyncio
import random
import signal
import threading
from datetime import datetime

from slixmpp import ClientXMPP
from services.session_manager import SessionManager
from lib.agent_loop import run_agent_loop

import config
from tools.roleplayer import should_respond
from lib import personality
from lib import ragroleplay

# Anti-ban: delay constants for MUC rejoin behavior
MUC_REJOIN_INITIAL_DELAY   = 5.0   # detik, delay awal sebelum rejoin
MUC_REJOIN_BACKOFF_MULT    = 2.0   # multiplier exponential backoff
MUC_REJOIN_MAX_DELAY       = 300.0 # detik, batas max backoff (5 menit)
MUC_REJOIN_COOLDOWN        = 10.0  # detik, cooldown minimum antar rejoin attempt
MUC_NICK_SUFFIX_MAX        = 3     # max coba nick alternatif (anti-ban: jangan terlalu banyak)


def _ts():
    return datetime.now().strftime('%H:%M:%S')


def _typing_delay(text: str) -> float:
    """Hitung delay mengetik (detik) proporsional dengan panjang teks."""
    char_count = len(text) if text else 0
    delay = char_count / config.TYPING_SPEED
    return max(1.0, min(delay, config.TYPING_MAX))


async def _read_delay():
    """Delay simulasi membaca pesan user."""
    delay = random.uniform(config.READ_DELAY_MIN, config.READ_DELAY_MAX)
    await asyncio.sleep(delay)


class XMPPClient(ClientXMPP):
    def __init__(self, jid, password, llm_client, tools_definition, TOOLS,
                 TOOL_HANDLERS, build_system_prompt, agent_max_iterations,
                 muc_rooms=None):
        super().__init__(jid, password)

        self._llm = llm_client
        self._tools_def = tools_definition
        self._TOOLS = TOOLS
        self._TOOL_HANDLERS = TOOL_HANDLERS
        self._build_system_prompt = build_system_prompt
        self._max_iterations = agent_max_iterations
        self._skill = config.AGENT_SKILL
        self._muc_rooms = muc_rooms or []
        # Custom nick dari config, fallback ke username JID
        self._muc_nick = config.XMPP_NICKNAME.strip() or jid.split('@')[0]
        self._muc_nick_suffix = 0  # counter untuk nick alternatif saat 409
        self._muc_ready: set[str] = set()

        self._session_mgr = SessionManager()
        self._loop = None
        self._stopped: asyncio.Event | None = None

        # Anti-ban: MUC rejoin tracking per room
        self._muc_rejoin_attempts: dict[str, int] = {}   # room -> jumlah attempt
        self._muc_rejoin_tasks: dict[str, asyncio.Task] = {}  # room -> pending rejoin task
        self._muc_last_join: dict[str, datetime] = {}    # room -> terakhir join (cooldown)

        self.auto_reconnect = True

        self.register_plugin('xep_0030')
        self.register_plugin('xep_0045')
        self.register_plugin('xep_0199')

        self.add_event_handler('session_start', self._on_session_start)
        self.add_event_handler('message', self._on_message)
        self.add_event_handler('groupchat_message', self._on_groupchat_message)
        self.add_event_handler('disconnected', self._on_disconnected)
        self.add_event_handler('connected', self._on_connected)
        self.add_event_handler('groupchat_presence', self._on_muc_presence)

    def _get_muc_nick(self, room: str) -> str:
        """Anti-ban: resolve nick untuk room, coba nick alternatif kalau conflict."""
        base = config.XMPP_NICKNAME.strip() or self._muc_nick
        suffix = self._muc_rejoin_attempts.get("_nick_" + room, 0)
        if suffix == 0:
            return base
        # Anti-ban: append suffix untuk menghindari 409 Conflict
        return f"{base}_{suffix}"

    def _calc_rejoin_delay(self, room: str) -> float:
        """Anti-ban: hitung delay rejoin dengan exponential backoff."""
        attempts = self._muc_rejoin_attempts.get(room, 0)
        delay = MUC_REJOIN_INITIAL_DELAY * (MUC_REJOIN_BACKOFF_MULT ** attempts)
        return min(delay, MUC_REJOIN_MAX_DELAY)

    def _schedule_muc_rejoin(self, room: str):
        """Anti-ban: schedule rejoin room dengan backoff & cooldown."""
        # Cancel pending rejoin task untuk room yang sama (anti-ban: avoid duplicate rejoin)
        pending = self._muc_rejoin_tasks.get(room)
        if pending and not pending.done():
            pending.cancel()
            print(f'[{_ts()}] MUC [{room}] Cancelled pending rejoin (new trigger)', flush=True)

        # Check cooldown: jangan rejoin terlalu cepat berturut-turut
        now = datetime.now()
        last_join = self._muc_last_join.get(room)
        if last_join:
            elapsed = (now - last_join).total_seconds()
            if elapsed < MUC_REJOIN_COOLDOWN:
                # Anti-ban: too soon, schedule delayed rejoin instead of immediate
                cooldown_left = MUC_REJOIN_COOLDOWN - elapsed
                print(f'[{_ts()}] MUC [{room}] Cooldown active ({cooldown_left:.0f}s left), delaying rejoin', flush=True)
                delay = cooldown_left + self._calc_rejoin_delay(room)
            else:
                delay = self._calc_rejoin_delay(room)
        else:
            delay = self._calc_rejoin_delay(room)

        # Increment attempt counter (anti-ban: track for exponential backoff)
        attempts = self._muc_rejoin_attempts.get(room, 0) + 1
        self._muc_rejoin_attempts[room] = attempts

        print(f'[{_ts()}] MUC [{room}] Rejoin scheduled in {delay:.0f}s (attempt #{attempts})', flush=True)

        if self._loop and not self._loop.is_closed():
            task = asyncio.run_coroutine_threadsafe(
                self._muc_rejoin_coro(room, delay), self._loop
            )
            self._muc_rejoin_tasks[room] = task

    async def _muc_rejoin_coro(self, room: str, delay: float):
        """Anti-ban: coroutine untuk rejoin room setelah delay."""
        try:
            await asyncio.sleep(delay)
            # Double-check: jangan rejoin kalau sudah di _muc_ready
            if room in self._muc_ready:
                print(f'[{_ts()}] MUC [{room}] Already ready, skip rejoin', flush=True)
                return
            nick = self._get_muc_nick(room)
            print(f'[{_ts()}] MUC [{room}] Rejoining as {nick}...', flush=True)
            await self.plugin['xep_0045'].join_muc_wait(room, nick, maxstanzas=0)
            self._muc_last_join[room] = datetime.now()
            # _muc_ready akan di-set oleh _on_muc_presence saat join berhasil
            self._muc_rejoin_attempts.pop(room, None)
            self._muc_rejoin_attempts.pop("_nick_" + room, None)
            print(f'[{_ts()}] MUC [{room}] Rejoin successful as {nick}', flush=True)
        except asyncio.CancelledError:
            print(f'[{_ts()}] MUC [{room}] Rejoin cancelled', flush=True)
        except Exception as e:
            print(f'[{_ts()}] MUC [{room}] Rejoin failed: {e}', flush=True)
            # Anti-ban: handle 409 Conflict - nick sudah dipakai orang lain
            if '409' in str(e) or 'conflict' in str(e).lower():
                nick_attempts = self._muc_rejoin_attempts.get("_nick_" + room, 0)
                if nick_attempts < MUC_NICK_SUFFIX_MAX:
                    # Anti-ban: coba nick alternatif (lily_, lily__)
                    self._muc_rejoin_attempts["_nick_" + room] = nick_attempts + 1
                    new_nick = self._get_muc_nick(room)
                    print(f'[{_ts()}] MUC [{room}] Nick conflict, trying alternative: {new_nick}', flush=True)
                    # Retry segera dengan nick baru (tanpa backoff rejoin, tapi tetap ada delay biasa)
                    self._schedule_muc_rejoin(room)
                else:
                    # Anti-ban: semua nick alternativehabis, stop retry untuk avoid ban
                    print(f'[{_ts()}] MUC [{room}] All nick variations exhausted, skipping room', flush=True)
                    print(f'[{_ts()}] MUC [{room}] Set XMPP_NICKNAME in .env to a unique nick', flush=True)
            else:
                # Anti-ban: error biasa (network, dll), retry with backoff
                self._schedule_muc_rejoin(room)

    async def _on_connected(self, event):
        print(f'[{_ts()}] XMPP connected', flush=True)

    async def _on_disconnected(self, event):
        print(f'[{_ts()}] XMPP disconnected', flush=True)
        # Anti-ban: cancel all pending rejoin tasks on disconnect
        for room, task in list(self._muc_rejoin_tasks.items()):
            if not task.done():
                task.cancel()
                print(f'[{_ts()}] MUC [{room}] Cancelled pending rejoin (disconnected)', flush=True)
        self._muc_rejoin_tasks.clear()

    async def _on_session_start(self, event):
        self.send_presence()
        self.get_roster()
        print(f'[{_ts()}] XMPP online as {self.boundjid.full}', flush=True)
        for room in self._muc_rooms:
            # Anti-ban: retry join dengan incremental delay & nick fallback
            success = False
            for attempt in range(1, 4):
                nick = self._get_muc_nick(room)
                try:
                    await self.plugin['xep_0045'].join_muc_wait(room, nick, maxstanzas=0)
                    print(f'[{_ts()}] Joined MUC room: {room} as {nick}', flush=True)
                    self._muc_last_join[room] = datetime.now()
                    self._muc_rejoin_attempts.pop(room, None)
                    self._muc_rejoin_attempts.pop("_nick_" + room, None)
                    success = True
                    break
                except Exception as e:
                    print(f'[{_ts()}] MUC join attempt #{attempt} failed ({room}): {e}', flush=True)
                    # Anti-ban: handle 409 Conflict - coba nick alternatif
                    if '409' in str(e) or 'conflict' in str(e).lower():
                        nick_attempts = self._muc_rejoin_attempts.get("_nick_" + room, 0)
                        if nick_attempts < MUC_NICK_SUFFIX_MAX:
                            nick_attempts += 1
                            self._muc_rejoin_attempts["_nick_" + room] = nick_attempts
                            print(f'[{_ts()}] MUC [{room}] Nick conflict, switching to: {self._get_muc_nick(room)}', flush=True)
                            # Retry segera dengan nick baru (jangan wait)
                            continue
                        else:
                            # Anti-ban: semua nick alternatif habis
                            print(f'[{_ts()}] MUC [{room}] All nick variations exhausted', flush=True)
                            break
                    elif attempt < 3:
                        # Anti-ban: error biasa, wait before retry (2s, 4s)
                        retry_delay = 2.0 * attempt
                        print(f'[{_ts()}] MUC [{room}] Retrying in {retry_delay:.0f}s...', flush=True)
                        await asyncio.sleep(retry_delay)
            if not success:
                # Anti-ban: semua attempt gagal, schedule background rejoin
                print(f'[{_ts()}] MUC [{room}] All join attempts failed, scheduling background rejoin', flush=True)
                self._schedule_muc_rejoin(room)

    def _on_message(self, msg):
        if msg['type'] not in ('chat', 'normal'):
            return
        jid = msg['from'].bare
        body = msg['body'].strip()
        if not body:
            return
        print(f'[{_ts()}] DM from {jid}: {body[:60]}', flush=True)
        threading.Thread(target=self._process_dm, args=(jid, body), daemon=True).start()

    def _on_groupchat_message(self, msg):
        if msg['type'] != 'groupchat':
            return
        room = msg['from'].bare
        nick = msg['from'].resource
        if self._is_my_nick(room, nick):
            return
        room = msg['from'].bare
        if room not in self._muc_ready:
            return
        body = msg['body'].strip()
        if not body:
            return
        print(f'[{_ts()}] MUC [{room}] <{nick}>: {body[:60]}', flush=True)
        threading.Thread(target=self._process_muc, args=(room, nick, body), daemon=True).start()

    def _is_my_nick(self, room: str, nick: str) -> bool:
        """Anti-ban: cek apakah nick yang dimasukan sesuai dengan nick bot di room."""
        expected = self._get_muc_nick(room)
        # Bandingkan dengan nick yang diharapkan, plus base nick tanpa suffix
        base = config.XMPP_NICKNAME.strip() or self._muc_nick
        return nick == expected or nick == base

    def _on_muc_presence(self, presence):
        room = presence['from'].bare
        nick = presence['from'].resource
        ptype = presence['type']
        if self._is_my_nick(room, nick) and ptype not in ('unavailable', 'error'):
            self._muc_ready.add(room)
            # Reset rejoin counter on successful join (anti-ban: avoid accumulating backoff)
            self._muc_rejoin_attempts.pop(room, None)
            self._muc_rejoin_attempts.pop("_nick_" + room, None)
        if ptype == 'unavailable':
            print(f'[{_ts()}] MUC [{room}] <{nick}> left', flush=True)
            # Anti-ban: remove from ready set on unavailable to keep state consistent
            self._muc_ready.discard(room)
            # Anti-ban: trigger auto-rejoin with exponential backoff
            if self._is_my_nick(room, nick):
                self._schedule_muc_rejoin(room)
        elif ptype == 'error':
            print(f'[{_ts()}] MUC [{room}] error: {presence}', flush=True)
            # Anti-ban: also rejoin on error (e.g. temporary failure)
            if self._is_my_nick(room, nick):
                self._muc_ready.discard(room)
                self._schedule_muc_rejoin(room)
        else:
            print(f'[{_ts()}] MUC [{room}] <{nick}> joined (type={ptype})', flush=True)

    def _process_dm(self, jid, body):
        session = self._session_mgr.get_or_create(
            jid, self._build_system_prompt(
                tools_definition=self._tools_def,
                character=config.AGENT_CHARACTER or None,
                skills=config.AGENT_SKILLS.split(",") if config.AGENT_SKILLS else None,
            )
        )
        session.cancel_timer()

        self.send_presence_subscription(pto=jid, ptype='subscribed')

        if body == ':new':
            self._session_mgr.reset(jid)
            print(f'[{_ts()}] Session reset for {jid}', flush=True)
            self._schedule_send(jid, 'Memulai sesi baru. Ada yang bisa di bantu?')
            return

        if 'roleplayer' in self._skill:
            try:
                results = ragroleplay.user_load(config.ragroleplay_db_path, unique_id=jid, character=personality.PERSONALITY.name)
                if results:
                    u = results[0]
                    context = (
                        f'[User Context]\n'
                        f'ID: {u["id"]}\n'
                        f'Nama: {u["fullname"]} ({u["nickname"]})\n'
                        f'Family: {u.get("familyname", "-") or "-"}\n'
                        f'Character: {u.get("character", "-") or "-"}\n'
                        f'Alias: {u.get("alias", "-") or "-"}\n'
                        f'Salutation: {u.get("salutation", "-") or "-"}\n'
                        f'Persona: {u.get("persona", "-") or "-"}\n'
                        f'Telegram: {u.get("telegram_id", "-") or "-"} / @{u.get("telegram_username", "-") or "-"}\n'
                        f'XMPP: {u.get("xmpp_username", "-") or "-"}\n'
                        f'[/User Context]\n'
                        f'[PENTING: Kamu SUDAH mengenal user ini. WAJIB panggil memories_latest(character="{personality.PERSONALITY.name}", user_id="{u["id"]}", limit=10) untuk mengambil riwayat percakapan terbaru. GUNAKAN data kondisi emotional/physical dari memori terbaru untuk melanjutkan state character secara natural sebelum merespon.]'
                    )
                    session.add_message('system', context)
                else:
                    session.add_message('system',
                        f'[User Context: Pengguna baru — belum ada di database]\n'
                        f'[PENTING: Kamu BELUM mengenal user ini. WAJIB tanya nama sebagai pembuka. '
                        f'Setelah nama diketahui, simpan via users_store. '
                        f'Lanjutkan percakapan natural dan proaktif melengkapi data user lainnya di pesan berikutnya.]\n'
                        f'[Platform: XMPP ({jid})]')
            except Exception:
                pass

        session.add_message('user', body)

        is_roleplay = 'roleplayer' in self._skill
        if not is_roleplay:
            self._schedule_send(jid, f'> {body}\nThinking...')

        # Delay 1: simulasi membaca pesan user
        if self._loop and not self._loop.is_closed():
            asyncio.run_coroutine_threadsafe(_read_delay(), self._loop)

        my_name = personality.PERSONALITY.name
        quote = body

        def on_tool_calls(content):
            if content and content.strip():
                self._schedule_send(jid, content, 'chat')

        tool_reminder = None

        final_content, should_close = run_agent_loop(
            session, self._llm, self._TOOLS, self._TOOL_HANDLERS,
            self._max_iterations, on_tool_calls=on_tool_calls, tool_reminder=tool_reminder
        )

        if final_content is not None:
            if is_roleplay:
                if config.XMPP_SELECTIVE_RESPONSE:
                    recent_msgs = []
                    for msg in session.messages[-6:]:
                        if msg.get('role') == 'user':
                            recent_msgs.append(f"User: {msg.get('content', '')}")
                        elif msg.get('role') == 'assistant' and msg.get('content'):
                            recent_msgs.append(f"{my_name}: {msg.get('content', '')}")
                    recent_history = "\n".join(recent_msgs)

                    if should_respond(
                        message=quote,
                        sender_nickname=jid,
                        recent_history=recent_history,
                        my_name=my_name,
                    ):
                        print(f'[{_ts()}] need_response=True → sending response', flush=True)
                        self._schedule_send(jid, final_content, 'chat')
                    else:
                        print(f'[{_ts()}] need_response=False → staying silent', flush=True)
                else:
                    from tools.roleplayer import _name_mentioned
                    if _name_mentioned(my_name, quote):
                        print(f'[{_ts()}] Name mentioned → sending response', flush=True)
                        self._schedule_send(jid, final_content, 'chat')
                    else:
                        print(f'[{_ts()}] Name not mentioned → staying silent', flush=True)
            else:
                self._schedule_send(jid, f'> {quote}\n{final_content}', 'chat')
        else:
            msg = 'Max iterations reached without final answer.'
            if is_roleplay:
                self._schedule_send(jid, msg, 'chat')
            else:
                self._schedule_send(jid, f'> {quote}\n{msg}', 'chat')

        # Natural close: DM only, roleplayer only
        if should_close and 'roleplayer' in self._skill:
            print(f'[{_ts()}] Natural close triggered for {jid}', flush=True)
            self._session_mgr.reset(jid)
        else:
            # DM: timeout 24 jam (efektif tidak auto-close), MUC tetap 5 menit
            session.start_timer(86400, self._timeout_session, jid, 'chat')

    def _process_muc(self, room, nick, body):
        session = self._session_mgr.get_or_create(
            room, self._build_system_prompt(
                tools_definition=self._tools_def,
                character=config.AGENT_CHARACTER or None,
                skills=config.AGENT_SKILLS.split(",") if config.AGENT_SKILLS else None,
            )
        )
        session.cancel_timer()

        if body == ':new':
            self._session_mgr.reset(room)
            print(f'[{_ts()}] Session reset for MUC room {room}', flush=True)
            self._schedule_send(room, 'Memulai sesi baru. Ada yang bisa di bantu?', mtype='groupchat')
            return

        prefixed = f'[{nick}] {body}'
        session.add_message('user', prefixed)

        if self._skill != 'roleplayer':
            self._schedule_send(room, f'> [{nick}] {body}\nThinking...', mtype='groupchat')

        # Delay 1: simulasi membaca pesan user
        if self._loop and not self._loop.is_closed():
            asyncio.run_coroutine_threadsafe(_read_delay(), self._loop)

        my_name = personality.PERSONALITY.name
        quote = f'[{nick}] {body}'

        _is_roleplay = 'roleplayer' in self._skill

        def on_tool_calls(content):
            if content and content.strip():
                self._schedule_send(room, content, 'groupchat')

        tool_reminder = None

        final_content, _should_close = run_agent_loop(
            session, self._llm, self._TOOLS, self._TOOL_HANDLERS,
            self._max_iterations, on_tool_calls=on_tool_calls, tool_reminder=tool_reminder
        )

        if final_content is not None:
            if _is_roleplay:
                if config.XMPP_SELECTIVE_RESPONSE:
                    recent_msgs = []
                    for msg in session.messages[-6:]:
                        if msg.get('role') == 'user':
                            recent_msgs.append(f"User: {msg.get('content', '')}")
                        elif msg.get('role') == 'assistant' and msg.get('content'):
                            recent_msgs.append(f"{my_name}: {msg.get('content', '')}")
                    recent_history = "\n".join(recent_msgs)

                    if should_respond(
                        message=quote,
                        sender_nickname=nick,
                        recent_history=recent_history,
                        my_name=my_name,
                    ):
                        print(f'[{_ts()}] need_response=True → sending response', flush=True)
                        self._schedule_send(room, final_content, 'groupchat')
                    else:
                        print(f'[{_ts()}] need_response=False → staying silent', flush=True)
                else:
                    from tools.roleplayer import _name_mentioned
                    if _name_mentioned(my_name, quote):
                        print(f'[{_ts()}] Name mentioned → sending response', flush=True)
                        self._schedule_send(room, final_content, 'groupchat')
                    else:
                        print(f'[{_ts()}] Name not mentioned → staying silent', flush=True)
            else:
                self._schedule_send(room, f'> {quote}\n{final_content}', 'groupchat')
        else:
            msg = 'Max iterations reached without final answer.'
            if _is_roleplay:
                self._schedule_send(room, msg, 'groupchat')
            else:
                self._schedule_send(room, f'> {quote}\n{msg}', 'groupchat')

        session.start_timer(300, self._timeout_session, room, 'groupchat')

    def _execute_tool(self, tool_call):
        from lib.agent_loop import execute_tool
        return execute_tool(tool_call, self._TOOL_HANDLERS)

    def _schedule_send(self, to, body, mtype='chat'):
        if self._loop and not self._loop.is_closed():
            asyncio.run_coroutine_threadsafe(
                self._send_coro(to, body, mtype), self._loop
            )
        else:
            print(f'[{_ts()}] WARNING: cannot send to {to} — loop unavailable', flush=True)

    async def _send_coro(self, to, body, mtype):
        try:
            # Delay 2: simulasi mengetik (proporsional dengan panjang pesan)
            delay = _typing_delay(body)
            print(f'[{_ts()}] Typing delay: {delay:.1f}s ({len(body)} chars)', flush=True)
            await asyncio.sleep(delay)

            msg = self.make_message(mto=to, mbody=body, mtype=mtype)
            msg.send()
        except Exception as e:
            print(f'[{_ts()}] SEND ERROR: {e}', flush=True)

    def _timeout_session(self, session_id, mtype):
        print(f'[{_ts()}] Session timeout: {session_id}', flush=True)
        self._schedule_send(session_id, 'Sesi ditutup. Sampai jumpa', mtype)
        self._session_mgr.reset(session_id)

    def start(self):
        print(f'[{_ts()}] Starting XMPP service...', flush=True)
        asyncio.run(self._run())

    async def _run(self):
        self._stopped = asyncio.Event()
        self._loop = asyncio.get_running_loop()

        # Hanya tangani SIGTERM untuk shutdown.
        # SENGATKAN SIGHUP: nohup kirim SIGHUP saat terminal close,
        # dan kita tidak mau proses mati karena itu.
        try:
            self._loop.add_signal_handler(signal.SIGTERM, self._stopped.set)
        except (NotImplementedError, RuntimeError):
            pass

        await self.connect()
        try:
            await self._stopped.wait()
        except (asyncio.CancelledError, KeyboardInterrupt):
            pass
        print(f'[{_ts()}] Shutting down...', flush=True)
        await self.disconnect()

    def stop(self):
        if self._loop and not self._loop.is_closed():
            asyncio.run_coroutine_threadsafe(self._async_stop(), self._loop)

    async def _async_stop(self):
        self._stopped.set()