# live_monitor.py
# SRS 直播监控 + 通知模块
#   群聊关注：群推群（条目按 channel + 群号）
#   私聊关注：挂在统一身份上（条目按 uid），通知时按 TG + 已绑定 QQ 扇出镜像推送
# 主站(site=1): 远程 aliyun SRS → /remote_srs/on_publish
# 副站(site=2): 本地 SRS → /remote_srs/local/on_publish
import hmac
import json
import os
import queue
import threading
from flask import Flask, request

app = Flask(__name__)
notification_queue = queue.Queue()

RECORD_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "live_monitor")
RECORD_FILE = os.path.join(RECORD_DIR, "record.json")

websocket = None

SITE_LABEL = {"1": "主站", "2": "副站"}

# Configure SRS http_hooks with ?token=<secret>, or send X-SRS-Token. The route
# and proxy/frp setup stay unchanged. Missing configuration fails closed.
SRS_CALLBACK_SECRET = os.environ.get(
    "BOT_SRS_CALLBACK_SECRET", os.environ.get("SRS_CALLBACK_SECRET", ""))
_live_sessions = set()       # (site, room_id) currently publishing
_live_sessions_lock = threading.Lock()
_auth_warning = False

HELP_TEXT = (
    "直播监控\n"
    "━━━━━━━━━━━━\n"
    "• .live follow {房间号0-100} {1/2/空}\n"
    "• .live unfollow {房间号} {1/2/空}\n"
    "• .live status  查看当前监控的房间\n"
    "  1=主站 2=副站 空=全部\n"
    "示例：.live follow 42 1\n"
    "━━━━━━━━━━━━\n"
    "直播站 https://live.rrrosmontis.icu"
)

def _ensure_record():
    os.makedirs(os.path.dirname(RECORD_FILE), exist_ok=True)
    if not os.path.exists(RECORD_FILE):
        with open(RECORD_FILE, "w", encoding="utf-8") as f:
            json.dump({"follows": {}}, f)

def _load_records():
    _ensure_record()
    with open(RECORD_FILE, "r", encoding="utf-8") as f:
        return json.load(f)

def _save_records(data):
    """原子写：同目录临时文件 + flush/fsync + os.replace，避免写一半留下截断的
    record.json（订阅全丢）。失败时清理临时文件，原文件保持不动。
    签名与输出内容与旧实现一致（ensure_ascii=False, indent=2，无尾随换行）。
    """
    _ensure_record()
    tmp = RECORD_FILE + ".tmp"
    try:
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
            f.flush()
            os.fsync(f.fileno())
        os.replace(tmp, RECORD_FILE)
    except BaseException:
        try:
            if os.path.exists(tmp):
                os.unlink(tmp)
        except OSError:
            pass
        raise


def _private_uid(value, channel=None):
    """把历史/当前私聊标识归一为统一身份 uid。

    兼容三种历史写法:
      - uid（[UID_BASE, 1e10)，历史 record 里 chat_id 直接存的就是 uid）
      - TG ns（>= 1e10）-> identity.resolve("tg", ns)
      - 原始 QQ 号 / TG 裸 id -> identity.resolve(channel|qq, value)
    解析不出来返回 None（调用方跳过该条目）。
    """
    try:
        number = int(value)
    except (TypeError, ValueError):
        return None
    import identity
    if identity.UID_BASE <= number < 10 ** 10:
        return number
    uid = identity.resolve("tg", number)
    if uid is not None:
        return uid
    return identity.resolve(channel or "qq", number)


def _entry_uid(entry):
    """关注条目的 uid：新记录直接存 uid，旧记录回退解析 chat_id"""
    if entry.get("uid") is not None:
        try:
            return int(entry["uid"])
        except (TypeError, ValueError):
            return None
    return _private_uid(entry.get("chat_id"), entry.get("channel"))


def _follow_key(entry):
    """关注条目的去重/匹配键：私聊按 uid（跨通道合并），群聊按 channel+群号"""
    if entry.get("chat_type") == "private":
        uid = _entry_uid(entry)
        return ("private", uid) if uid is not None else None
    return ("group", entry.get("channel", "tg"), str(entry.get("chat_id")))


def _matches_follow(entry, chat_type, chat_id, channel):
    """条目是否属于 (chat_type, chat_id, channel)。私聊参数 chat_id 传 uid。"""
    if entry.get("chat_type") != chat_type:
        return False
    if chat_type == "private":
        target = _private_uid(chat_id, channel)
        return target is not None and _entry_uid(entry) == target
    return entry.get("channel", "tg") == channel and str(entry.get("chat_id")) == str(chat_id)


def add_follow(room_id, chat_type, chat_id, site, channel="tg"):
    """关注直播间。site: '1'=主站, '2'=副站, ''=全部; channel: 群聊通知通道 (qq/tg)

    私聊关注统一挂在统一身份上（chat_id 传 uid），通知时按 TG + QQ 扇出；
    群聊关注仍是群推群（channel + 群号）。
    """
    data = _load_records()
    room_id = str(room_id)
    follows = data.setdefault("follows", {})
    room_follows = follows.setdefault(room_id, [])
    if chat_type == "private":
        uid = _private_uid(chat_id, channel)
        if uid is None:
            return
        entry = {"chat_type": "private", "uid": uid, "site": site}
    else:
        entry = {"channel": channel, "chat_type": chat_type,
                 "chat_id": chat_id, "site": site}
    key = _follow_key(entry)
    if key is None or not any(_follow_key(f) == key and f.get("site", "") == site
                              for f in room_follows):
        room_follows.append(entry)
    _save_records(data)


def remove_follow(room_id, chat_type, chat_id, site, channel="tg"):
    """取消关注"""
    data = _load_records()
    room_id = str(room_id)
    follows = data.get("follows", {})
    if room_id in follows:
        follows[room_id] = [
            f for f in follows[room_id]
            if not (_matches_follow(f, chat_type, chat_id, channel)
                    and f.get("site", "") == site)
        ]
        if not follows[room_id]:
            del follows[room_id]
    _save_records(data)


def get_follows(chat_type, chat_id, channel="tg"):
    """获取某聊天关注的所有房间号和站点（私聊按 uid，群聊按 channel+群号）"""
    data = _load_records()
    follows = data.get("follows", {})
    result = []
    for room_id, flist in follows.items():
        for f in flist:
            if _matches_follow(f, chat_type, chat_id, channel):
                result.append((room_id, f.get("site", "")))
                break
    return result


def status_text(chat_type=None, chat_id=None, channel=None):
    """当前已监控的房间列表；标注会推送到当前会话的房间（不暴露其它会话的 id）"""
    data = _load_records()
    follows = data.get("follows") or {}
    if not follows:
        return "当前没有监控任何直播间。\n用法：.live follow {房间号0-100} {1/2/空}"
    lines = ["直播关注（共 %d 个房间）" % len(follows), "━━━━━━━━━━━━"]
    for room in sorted(follows, key=lambda x: int(x) if str(x).isdigit() else 0):
        items = follows[room] or []
        sites = []
        for item in items:
            label = SITE_LABEL.get(item.get("site", ""), "全部")
            if label not in sites:
                sites.append(label)
        marks = []
        if chat_id is not None and any(
                _matches_follow(f, chat_type, chat_id, channel) for f in items):
            marks.append("本群" if chat_type == "group" else "本会话")
        if any(f.get("chat_type") == "private" for f in items):
            marks.append("私聊")
        suffix = "（%s）" % "、".join(marks) if marks else ""
        lines.append("• 房间 %-3s %s %s" % (room, "/".join(sites), suffix))
    lines.append("━━━━━━━━━━━━")
    lines.append("关注：.live follow <房间号> [1/2]　取消：.live unfollow <房间号> [1/2]")
    return "\n".join(lines)


def _notification(room_id, site, f):
    """把关注条目转成一条通知。

    私聊条目归一为 {"chat_type":"private","uid":..}（消费者按 TG + QQ 扇出）；
    群聊条目保持 {"channel":.., "chat_type":"group", "chat_id":群号}（群推群）；
    私聊条目若 uid 解析不出来，退回旧写法（按原 channel/target 发送），保证兼容。
    """
    if f.get("chat_type") == "private":
        uid = _entry_uid(f)
        if uid is not None:
            return {"room_id": room_id, "site": site,
                    "chat_type": "private", "uid": uid}
    return {"room_id": room_id, "site": site, "chat_type": f.get("chat_type"),
            "chat_id": f.get("chat_id"), "channel": f.get("channel", "tg")}


def _srs_callback_authorized():
    global _auth_warning
    if not SRS_CALLBACK_SECRET:
        if not _auth_warning:
            print("[live_monitor] 拒绝 SRS 回调: 未配置 BOT_SRS_CALLBACK_SECRET")
            _auth_warning = True
        return False
    authorization = request.headers.get("Authorization", "")
    supplied = (request.headers.get("X-SRS-Token", "")
                or request.args.get("token", "")
                or (authorization[7:] if authorization.startswith("Bearer ") else ""))
    return hmac.compare_digest(str(supplied), str(SRS_CALLBACK_SECRET))


def _callback_room_id(data):
    raw_stream = str(data.get("stream", ""))
    return (raw_stream.replace("room_", "", 1)
            if raw_stream.startswith("room_") else raw_stream)


def _begin_live_session(site, room_id):
    """Atomically deduplicate repeated on_publish for one active session."""
    key = (str(site), str(room_id))
    with _live_sessions_lock:
        if key in _live_sessions:
            return False
        _live_sessions.add(key)
        return True


def _end_live_session(site, room_id):
    with _live_sessions_lock:
        _live_sessions.discard((str(site), str(room_id)))


def _handle_on_publish(site):
    """处理 on_publish 回调，site='1' 主站 '2' 副站"""
    if not _srs_callback_authorized():
        return '{"code": 1, "msg": "forbidden"}', 403
    data = request.get_json(force=True, silent=True) or {}
    print(f"[live_monitor] on_publish(site={site}) received: {json.dumps(data, ensure_ascii=False)}")
    raw_stream = str(data.get("stream", ""))
    if not raw_stream:
        print(f"[live_monitor] on_publish(site={site}): no stream field")
        return '{"code": 0}'

    room_id = _callback_room_id(data)
    if not _begin_live_session(site, room_id):
        print(f"[live_monitor] on_publish(site={site}): duplicate active session room={room_id}")
        return '{"code": 0}'

    records = _load_records()
    room_follows = records.get("follows", {}).get(room_id, [])
    matched = [f for f in room_follows if f.get("site", "") in ("", site)]
    print(f"[live_monitor] on_publish(site={site}): raw_stream={raw_stream}, room_id={room_id}, matched={matched}")

    for f in matched:
        notif = _notification(room_id, site, f)
        notification_queue.put(notif)
        print("[live_monitor] on_publish(site=%s): queued for %s"
              % (site, notif.get("uid") if notif.get("uid") is not None
                 else "%s/%s/%s" % (notif.get("channel", "qq"), notif.get("chat_type"),
                                    notif.get("chat_id"))))

    return '{"code": 0}'


@app.route("/internal/notify", methods=["POST"])
def internal_notify():
    """live-api 内部开播通知（仅本机 127.0.0.1 可达）：按 room+site 直接入队。"""
    # Flask 绑在 0.0.0.0（SRS 公网回调需要），本端点必须自验来源，否则外部可伪造开播通知
    if request.remote_addr not in ("127.0.0.1", "::1"):
        print(f"[live_monitor] internal_notify 拒绝非本机来源: {request.remote_addr}")
        return '{"code": 1, "msg": "forbidden"}', 403
    data = request.get_json(force=True, silent=True) or {}
    room_id = str(data.get("room_id", "")).strip()
    site = str(data.get("site", "")).strip()
    title = str(data.get("title", ""))
    nickname = str(data.get("nickname", ""))
    if not room_id or site not in ("1", "2"):
        return '{"code": 1, "msg": "bad params"}'
    if not _begin_live_session(site, room_id):
        return '{"code": 0}'
    records = _load_records()
    room_follows = records.get("follows", {}).get(room_id, [])
    matched = [f for f in room_follows if f.get("site", "") in ("", site)]
    for f in matched:
        notification_queue.put(_notification(room_id, site, f))
    print(f"[live_monitor] internal_notify room={room_id} site={site} matched={len(matched)} title={title!r} anchor={nickname!r}")
    return '{"code": 0}'


@app.route("/remote_srs/on_publish", methods=["POST"])
def on_publish_main():
    """主站(aliyun 远程) on_publish"""
    return _handle_on_publish("1")


def _handle_on_unpublish(site):
    if not _srs_callback_authorized():
        return '{"code": 1, "msg": "forbidden"}', 403
    data = request.get_json(force=True, silent=True) or {}
    raw_stream = str(data.get("stream", ""))
    if raw_stream:
        _end_live_session(site, _callback_room_id(data))
    return '{"code": 0}'


@app.route("/remote_srs/on_unpublish", methods=["POST"])
def on_unpublish_main():
    return _handle_on_unpublish("1")


@app.route("/remote_srs/local/on_publish", methods=["POST"])
def on_publish_local():
    """副站(本地) on_publish"""
    return _handle_on_publish("2")


@app.route("/remote_srs/local/on_unpublish", methods=["POST"])
def on_unpublish_local():
    return _handle_on_unpublish("2")


def run_flask():
    from waitress import serve
    serve(app, host=os.environ.get("BOT_LIVE_HOST", "127.0.0.1"), port=1986)