src/chat_capture.py

#!/usr/bin/env python3
"""
TikTok Live Chat Capture - Battle Score Tracker
FIXED: Extracts actual scores from LinkMicArmiesEvent
NEW: Battle events can be toggled ON/OFF
"""

import sys
import os
import json
import site
import time
import textwrap
from datetime import datetime

try:
    site.addsitedir(site.getusersitepackages())
except:
    pass

sys.stdout.reconfigure(encoding="utf-8")

# Import ALL available events
try:
    from TikTokLive import TikTokLiveClient
    from TikTokLive.events import *

    EVENTS_IMPORTED = True
except ImportError:
    EVENTS_IMPORTED = False
    print("❌ Could not import TikTokLive events")

# --- CONFIG (Overlay Settings) ---
FONT_SIZE = 16
FONT_NAME = "Arial"
WRAP_WIDTH = 40
SIDEBAR_X = 1560
BOTTOM_Y = 1040
LINE_HEIGHT = 38
MAX_LINES_ON_SCREEN = 28
VIDEO_WIDTH = 1920
VIDEO_HEIGHT = 1080

# Network settings
MAX_RETRIES = 3
RETRY_DELAY = 10

# ⚔️ BATTLE SETTINGS ⚔️
TRACK_BATTLES = False  # ← Set to False to disable battle tracking
BATTLE_DEBUG_MODE = False  # ← Set to False to hide debug info (participant list)

# 😊 EMOJI MAPPING - TikTok codes to Unicode emoji
EMOJI_MAP = {
    "[laugh]": "😆",
    "[thanks]": "🙏",
    "[thumb]": "👍",
    "[hi]": "👋",
    "[congrat]": "🎉",
    "[rockyserious]": "🗿",
    "[rockyloveit]": "🥰",
    "[rockycool]": "😎",
    "[rosiedislike]": "👎",
    "[rosieawkward]": "😅",
    "[rosiekisskiss]": "😘",
    "[rosiecute]": "🥺",
    "[jolliekissingface]": "😚",
    "[jolliewow]": "😮",
    "[jolliespeechless]": "😶",
    "[jolliesatisfied]": "😌",
    "[sagethink]": "🤔",
    "[sagefulfilled]": "😇",
    "[sageclever]": "🤓",
    "[sagemoney]": "🤑",
}
# -------------------------------------


def main():
    if len(sys.argv) < 4:
        print("Usage: python chat_capture.py <username> <output_dir> <format>")
        sys.exit(1)

    username = sys.argv[1]
    output_dir = os.path.normpath(sys.argv[2])
    save_format = sys.argv[3]

    if not os.path.exists(output_dir):
        os.makedirs(output_dir, exist_ok=True)

    timestamp_str = datetime.now().strftime("%Y%m%d_%H%M%S")
    start_time = datetime.now()

    print(f"💬 Chat Capture started for @{username}")
    print(f"📂 ID: {timestamp_str}")
    print(f"ℹ️ Battle Tracking: {'✅ ENABLED' if TRACK_BATTLES else '❌ DISABLED'}")
    if TRACK_BATTLES and BATTLE_DEBUG_MODE:
        print(f"ℹ️ Battle Debug Mode: ✅ ON (showing participant lists)")
    elif TRACK_BATTLES:
        print(f"ℹ️ Battle Debug Mode: ❌ OFF (clean output only)")
    sys.stdout.flush()

    if not EVENTS_IMPORTED:
        print("❌ Library missing. Please install TikTokLive.")
        sys.exit(1)

    chat_messages = []
    last_battle_time = 0
    last_scores = {}
    current_battle_id = None  # Track current battle
    retry_count = 0

    client = TikTokLiveClient(unique_id=f"@{username}")

    # Suppress logs
    import logging

    logging.getLogger("TikTokLive").setLevel(logging.ERROR)

    def get_elapsed():
        delta = datetime.now() - start_time
        return delta.total_seconds()

    def replace_tiktok_emojis(text):
        """Replace TikTok emoji codes with Unicode emojis"""
        for code, emoji in EMOJI_MAP.items():
            text = text.replace(code, emoji)
        return text

    def format_srt_time(seconds):
        hours = int(seconds // 3600)
        minutes = int((seconds % 3600) // 60)
        secs = int(seconds % 60)
        millis = int((seconds % 1) * 1000)
        return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"

    def format_ass_time(seconds):
        hours = int(seconds // 3600)
        minutes = int((seconds % 3600) // 60)
        secs = int(seconds % 60)
        millis = int((seconds % 1) * 100)
        return f"{hours}:{minutes:02d}:{secs:02d}.{millis:02d}"

    def save_chat():
        base_filename = f"{username}_{timestamp_str}_chat"
        file_path_base = os.path.join(output_dir, base_filename)

        try:
            if save_format == "json":
                with open(file_path_base + ".json", "w", encoding="utf-8") as f:
                    json.dump(
                        {"username": username, "messages": chat_messages},
                        f,
                        indent=2,
                        ensure_ascii=False,
                    )

            elif save_format == "txt":
                with open(file_path_base + ".txt", "w", encoding="utf-8") as f:
                    f.write(f"TikTok Chat - @{username} ({timestamp_str})\n\n")
                    for msg in chat_messages:
                        if msg["type"] == "comment":
                            f.write(f"{msg['username']}: {msg['message']}\n")
                        elif msg["type"] == "gift":
                            f.write(
                                f"Gift: {msg['username']} sent {msg['gift_count']}x {msg['gift']}\n"
                            )
                        elif msg["type"] == "battle":
                            f.write(f"{msg['status']}\n")

            elif save_format == "srt":
                with open(file_path_base + ".srt", "w", encoding="utf-8") as f:
                    counter = 1
                    for i, msg in enumerate(chat_messages):
                        if msg["type"] not in ["comment", "gift", "battle"]:
                            continue
                        elapsed = msg.get("elapsed_seconds", 0)
                        text = ""
                        if msg["type"] == "comment":
                            text = f"{msg['username']}: {msg['message']}"
                        elif msg["type"] == "gift":
                            text = f"🎁 {msg['username']} sent {msg['gift']}"
                        elif msg["type"] == "battle":
                            text = f"{msg['status']}"
                        text = text.replace("\n", " ")

                        next_elapsed = elapsed + 3600.0
                        for j in range(i + 1, len(chat_messages)):
                            if chat_messages[j]["type"] in [
                                "comment",
                                "gift",
                                "battle",
                            ]:
                                next_elapsed = chat_messages[j].get(
                                    "elapsed_seconds", elapsed + 3600.0
                                )
                                break

                        f.write(
                            f"{counter}\n{format_srt_time(elapsed)} --> {format_srt_time(next_elapsed)}\n{text}\n\n"
                        )
                        counter += 1

            elif save_format == "ass":
                with open(file_path_base + ".ass", "w", encoding="utf-8") as f:
                    f.write(f"""[Script Info]
ScriptType: v4.00+
PlayResX: {VIDEO_WIDTH}
PlayResY: {VIDEO_HEIGHT}

[V4+ Styles]
Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding
Style: Feed,{FONT_NAME},{FONT_SIZE},&H00FFFFFF,&H000000FF,&H00000000,&H80000000,0,0,0,0,100,100,0,0,1,1,0,1,10,10,10,1

[Events]
Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text
""")
                    screen_states = []
                    current_buffer = []
                    valid_messages = [
                        m
                        for m in chat_messages
                        if m["type"] in ["comment", "gift", "battle"]
                    ]

                    for i, msg in enumerate(valid_messages):
                        raw_text = ""
                        if msg["type"] == "comment":
                            raw_text = f"{msg['username']}: {msg['message']}"
                        elif msg["type"] == "gift":
                            raw_text = f"🎁 {msg['username']} sent {msg['gift']}"
                        elif msg["type"] == "battle":
                            raw_text = f"{msg['status']}"

                        wrapped_lines = textwrap.wrap(
                            raw_text.replace("\n", " "), width=WRAP_WIDTH
                        )
                        current_buffer.extend(wrapped_lines)

                        while len(current_buffer) > MAX_LINES_ON_SCREEN:
                            current_buffer.pop(0)

                        start_t = msg.get("elapsed_seconds", 0)
                        if i + 1 < len(valid_messages):
                            end_t = valid_messages[i + 1].get(
                                "elapsed_seconds", start_t + 3600.0
                            )
                            if end_t <= start_t:
                                end_t = start_t + 0.1
                        else:
                            end_t = start_t + 3600.0

                        screen_states.append(
                            {
                                "start": start_t,
                                "end": end_t,
                                "lines": list(current_buffer),
                            }
                        )

                    for state in screen_states:
                        s_time = format_ass_time(state["start"])
                        e_time = format_ass_time(state["end"])
                        for idx, line_text in enumerate(reversed(state["lines"])):
                            y_position = BOTTOM_Y - (idx * LINE_HEIGHT)
                            safe_text = (
                                line_text.replace(",", "\\,")
                                .replace("{", "(")
                                .replace("}", ")")
                            )
                            f.write(
                                f"Dialogue: 0,{s_time},{e_time},Feed,,0,0,0,,{{\\pos({SIDEBAR_X},{y_position})}}{safe_text}\n"
                            )

        except Exception as e:
            print(f"⚠️ Save error: {e}")

    def log_battle_event(status_text):
        nonlocal last_battle_time
        current_time = time.time()

        # Prevent duplicate logs within 0.5 seconds
        if current_time - last_battle_time < 0.5:
            return

        last_battle_time = current_time

        message = {
            "timestamp": datetime.now().isoformat(),
            "type": "battle",
            "status": status_text,
            "elapsed_seconds": get_elapsed(),
        }
        chat_messages.append(message)
        print(f"⚔️ {status_text}")
        sys.stdout.flush()
        save_chat()

    def extract_battle_scores(event):
        """Extract battle scores from LinkMicArmiesEvent"""
        try:
            if not hasattr(event, "armies") or not event.armies:
                return None

            scores = []

            # armies is a dictionary: {user_id: BattleUserArmies}
            for user_id, army_data in event.armies.items():
                if hasattr(army_data, "user_armies") and army_data.user_armies:
                    # user_armies is a list of BattleUserArmy objects
                    for user_army in army_data.user_armies:
                        if hasattr(user_army, "score") and hasattr(
                            user_army, "nickname"
                        ):
                            scores.append(
                                {
                                    "nickname": user_army.nickname,
                                    "score": user_army.score,
                                    "user_id": user_army.user_id
                                    if hasattr(user_army, "user_id")
                                    else None,
                                }
                            )

            return scores if scores else None

        except Exception as e:
            if BATTLE_DEBUG_MODE:
                print(f"⚠️ Error extracting scores: {e}")
            return None

    @client.on(ConnectEvent)
    async def on_connect(event):
        nonlocal retry_count
        retry_count = 0
        print(f"✅ Connected to @{username}")
        if TRACK_BATTLES:
            print(f"⚔️ Battle tracking active...")
        sys.stdout.flush()

    @client.on(DisconnectEvent)
    async def on_disconnect(event):
        print(f"⚠️ Disconnected from @{username}")
        save_chat()
        sys.stdout.flush()

    @client.on(CommentEvent)
    async def on_comment(event):
        try:
            user_name = event.user.nickname or event.user.unique_id or "User"
            comment_text = replace_tiktok_emojis(event.comment)

            message = {
                "timestamp": datetime.now().isoformat(),
                "type": "comment",
                "username": user_name,
                "message": comment_text,
                "elapsed_seconds": get_elapsed(),
            }
            chat_messages.append(message)
            print(f"💬 {user_name}: {comment_text}")
            sys.stdout.flush()
            save_chat()
        except Exception:
            pass

    @client.on(GiftEvent)
    async def on_gift(event):
        try:
            user_name = event.user.nickname or "Viewer"
            gift_name = event.gift.name if hasattr(event.gift, "name") else "Gift"
            count = event.gift.count if hasattr(event.gift, "count") else 1

            message = {
                "timestamp": datetime.now().isoformat(),
                "type": "gift",
                "username": user_name,
                "gift": gift_name,
                "gift_count": count,
                "elapsed_seconds": get_elapsed(),
            }
            chat_messages.append(message)
            print(f"🎁 {user_name} sent {count}x {gift_name}")
            sys.stdout.flush()
            save_chat()
        except Exception:
            pass

    # Battle event handlers - ONLY if TRACK_BATTLES is True
    if TRACK_BATTLES:
        try:

            @client.on(LinkMicArmiesEvent)
            async def on_link_mic_armies(event):
                nonlocal last_scores

                # Debug: print ALL attributes to find real battle scores
                if BATTLE_DEBUG_MODE:
                    print(f"\n{'=' * 60}")
                    print(f"🔍 LinkMicArmiesEvent - Searching for battle scores...")

                    # Skip these problematic attributes
                    skip_attrs = {
                        "as_base64",
                        "bytes",
                        "SerializeToString",
                        "FromString",
                        "to_dict",
                        "to_json",
                        "from_dict",
                        "from_json",
                    }

                    # Check common score attributes
                    score_attrs = {}
                    for attr in dir(event):
                        if attr.startswith("_") or attr in skip_attrs:
                            continue

                        try:
                            value = getattr(event, attr)
                            if callable(value):
                                continue

                            attr_lower = attr.lower()
                            if any(
                                keyword in attr_lower
                                for keyword in [
                                    "score",
                                    "point",
                                    "total",
                                    "battle",
                                    "host",
                                    "guest",
                                    "left",
                                    "right",
                                ]
                            ):
                                score_attrs[attr] = value
                                print(f"   🎯 {attr}: {value}")
                        except:
                            pass

                    # Check battle_settings for scores
                    if hasattr(event, "battle_settings"):
                        try:
                            print(f"   📋 battle_settings: {event.battle_settings}")
                        except:
                            pass

                    print(f"{'=' * 60}\n")

                scores = extract_battle_scores(event)

                if scores and len(scores) >= 2:
                    # Sort by score descending to get top 2
                    scores_sorted = sorted(
                        scores, key=lambda x: x["score"], reverse=True
                    )

                    user1 = scores_sorted[0]
                    user2 = scores_sorted[1]

                    # Create score key for comparison
                    score_key = f"{user1['score']}_{user2['score']}"

                    # Only log to chat if scores changed
                    if last_scores.get("key") != score_key:
                        last_scores = {"key": score_key, "user1": user1, "user2": user2}

                        status_text = f"BATTLE: {user1['score']} ({user1['nickname']}) VS {user2['score']} ({user2['nickname']})"
                        log_battle_event(status_text)

                        if BATTLE_DEBUG_MODE:
                            print(f"🔍 Total participants in armies: {len(scores)}")
                            for idx, s in enumerate(scores_sorted[:5], 1):
                                print(
                                    f"   #{idx}: {s['nickname']} - {s['score']} points"
                                )

                elif BATTLE_DEBUG_MODE and scores:
                    print(f"⚠️ Only {len(scores)} participant(s) found in armies")
                    for s in scores:
                        print(f"   • {s['nickname']}: {s['score']}")

        except Exception as e:
            print(f"❌ Could not register LinkMicArmiesEvent handler: {e}")

        # Additional battle events (optional)
        try:

            @client.on(LinkMicBattleEvent)
            async def on_link_mic_battle(event):
                nonlocal current_battle_id, last_scores

                # Check if this is a FINISH event with results
                if hasattr(event, "action") and hasattr(event, "battle_result"):
                    action = str(event.action) if event.action else ""

                    if "FINISH" in action and event.battle_result:
                        # Extract final scores from battle_result
                        results = []
                        for user_id, battle_result in event.battle_result.items():
                            if hasattr(battle_result, "score"):
                                nickname = "Unknown"

                                # Try to find nickname from anchor_info
                                if hasattr(event, "anchor_info") and event.anchor_info:
                                    for anchor in event.anchor_info:
                                        if (
                                            hasattr(anchor, "user_id")
                                            and anchor.user_id == user_id
                                        ):
                                            if hasattr(anchor, "user_info") and hasattr(
                                                anchor.user_info, "user"
                                            ):
                                                if hasattr(
                                                    anchor.user_info.user, "nick_name"
                                                ):
                                                    nickname = (
                                                        anchor.user_info.user.nick_name
                                                    )
                                                    break

                                results.append(
                                    {
                                        "user_id": user_id,
                                        "nickname": nickname,
                                        "score": battle_result.score,
                                        "result": str(battle_result.result)
                                        if hasattr(battle_result, "result")
                                        else "UNKNOWN",
                                    }
                                )

                        if len(results) >= 2:
                            # Sort by score to get winner first
                            results_sorted = sorted(
                                results, key=lambda x: x["score"], reverse=True
                            )

                            winner = results_sorted[0]
                            loser = results_sorted[1]

                            status_text = f"BATTLE FINISHED: {winner['score']} ({winner['nickname']}) VS {loser['score']} ({loser['nickname']}) - Winner: {winner['nickname']}"
                            log_battle_event(status_text)

                            if BATTLE_DEBUG_MODE:
                                print(f"🏆 Final results:")
                                for idx, r in enumerate(results_sorted, 1):
                                    print(
                                        f"   #{idx}: {r['nickname']} - {r['score']} points ({r['result']})"
                                    )

                        # Reset for next battle
                        current_battle_id = None
                        last_scores = {}

                    elif "OPEN" in action:
                        # New battle starting - reset tracking
                        battle_id = getattr(event, "battle_id", None)
                        if battle_id:
                            current_battle_id = battle_id
                            last_scores = {}
                            if BATTLE_DEBUG_MODE:
                                print(f"🔔 New battle starting: {battle_id}")
        except:
            if BATTLE_DEBUG_MODE:
                print("ℹ️ LinkMicBattleEvent not available in this TikTokLive version")

    # --- MAIN LOOP ---
    print("✅ Event handlers registered")
    print("🎯 Monitoring chat...")

    while retry_count < MAX_RETRIES:
        try:
            print(
                f"🔌 Connecting to @{username}... (Attempt {retry_count + 1}/{MAX_RETRIES})"
            )
            sys.stdout.flush()
            client.run()
            break
        except KeyboardInterrupt:
            print("\n🛑 Exit requested...")
            save_chat()
            sys.exit(0)
        except Exception as e:
            retry_count += 1
            print(f"❌ Connection error: {e}")
            if retry_count < MAX_RETRIES:
                print(f"⏳ Retrying in {RETRY_DELAY} seconds...")
                time.sleep(RETRY_DELAY)
            else:
                print("❌ Max retries reached. Exiting.")
                save_chat()
                sys.exit(1)


if __name__ == "__main__":
    main()