import asyncio
import os
import sys
import json
import re
import requests
from telethon import TelegramClient, events
from telethon.sessions import StringSession
from telethon.tl.types import MessageService
from telethon.errors import FloodWaitError

sys.stdout.reconfigure(line_buffering=True)

API_ID = int(os.environ.get("API_ID", "0"))
API_HASH = os.environ.get("API_HASH")
SESSION_STRING = os.environ.get("TG_SESSION_STRING")
SOURCES = [s.strip() for s in os.environ.get("SOURCES", "").split(",") if s.strip()]
DESTINATION = os.environ.get("DESTINATION")
UNLOCK_API_URL = os.environ.get("UNLOCK_API_URL", "http://localhost:7860")
UNLOCK_PASSWORD = os.environ.get("UNLOCK_PASSWORD", "")

if not all([API_ID, API_HASH, SESSION_STRING, SOURCES, DESTINATION]):
    print("缺少必要的环境变量")
    sys.exit(1)

slug_cache = {}

def extract_slugs_from_text(text):
    if not text:
        return []
    hdhive_pattern = r"https?://hdhive\.com/resource/(?:[^/]+/)?([a-f0-9]{32})"
    matches = re.findall(hdhive_pattern, text, re.IGNORECASE)
    if matches:
        return matches
    return []

def get_first_slug_from_message(message):
    if not message.text:
        return None
    for entity in message.entities or []:
        url = None
        if hasattr(entity, 'url'):
            url = entity.url
        if url:
            slugs = extract_slugs_from_text(url)
            if slugs:
                return slugs[0]
    slugs = extract_slugs_from_text(message.text)
    return slugs[0] if slugs else None

def get_direct_link(slug):
    if slug in slug_cache:
        return slug_cache[slug]
    url = f"{UNLOCK_API_URL}/api/cache/unlock"
    payload = {"slug": slug, "allow_points": True}
    if UNLOCK_PASSWORD:
        payload["password"] = UNLOCK_PASSWORD
    try:
        resp = requests.post(url, json=payload, timeout=10)
        if resp.status_code == 200:
            data = resp.json()
            if data.get("success") and data.get("link"):
                direct_link = data["link"]
                slug_cache[slug] = direct_link
                return direct_link
    except Exception as e:
        print(f"解锁错误: {e}")
    return None

def replace_links_in_text(text, slug, direct_link):
    if not text:
        return text
    pattern = rf"(https?://hdhive\.com/resource/(?:[^/]+/)?{slug})"
    return re.sub(pattern, direct_link, text, flags=re.IGNORECASE)

async def send_modified_message(client, dest_entity, original_msg):
    if isinstance(original_msg, MessageService):
        return

    try:
        slug = get_first_slug_from_message(original_msg)
        if not slug:
            await client.forward_messages(dest_entity, original_msg)
            print(f"直接转发 {original_msg.id}（无HDHive链接）")
            return

        direct_link = get_direct_link(slug)
        if not direct_link:
            await client.forward_messages(dest_entity, original_msg)
            print(f"直接转发 {original_msg.id}（解锁失败）")
            return

        if original_msg.text:
            new_text = replace_links_in_text(original_msg.text, slug, direct_link)
            try:
                await client.send_message(dest_entity, new_text)
                print(f"✅ 已发送替换后消息 {original_msg.id}")
            except FloodWaitError as e:
                print(f"⏳ 速率限制，等待 {e.seconds} 秒后重试...")
                await asyncio.sleep(e.seconds)
                await client.send_message(dest_entity, new_text)
                print(f"✅ 重试成功 {original_msg.id}")
        else:
            await client.forward_messages(dest_entity, original_msg)
    except Exception as e:
        print(f"❌ 处理失败 {original_msg.id}: {e}")

async def main():
    print("等待 60 秒，确保旧容器已退出...")
    await asyncio.sleep(60)

    while True:
        client = None
        try:
            client = TelegramClient(StringSession(SESSION_STRING), API_ID, API_HASH)
            await client.start()
            print("✅ 连接成功")

            dest = await client.get_entity(DESTINATION)
            print(f"📢 目标频道: {DESTINATION}")

            # 添加官方通知(777000)的实时监听
            @client.on(events.NewMessage(from_users=777000))
            async def verification_handler(event):
                msg = event.message
                text = msg.text or ""
                print(f"\n📩 收到新官方通知（实时）")
                print(f"   消息ID: {msg.id}")
                print(f"   时间: {msg.date}")
                print(f"   内容: {text[:500]}")
                match = re.search(r"\b(\d{5,6})\b", text)
                if match:
                    print(f"   🔑 验证码: {match.group(1)}")
                else:
                    print("   ℹ️ 未找到验证码")
                print()

            # 启动时获取最近 5 条官方通知（仅打印，不转发）
            try:
                print("📜 获取最近 5 条官方通知...")
                async for msg in client.iter_messages(777000, limit=5, reverse=False):
                    text = msg.text or ""
                    print(f"--- 消息ID: {msg.id} | 时间: {msg.date} ---")
                    print(f"内容: {text[:300]}")
                    match = re.search(r"\b(\d{5,6})\b", text)
                    if match:
                        print(f"验证码: {match.group(1)}")
                    print()
            except Exception as e:
                print(f"获取官方通知失败: {e}")

            # 为每个源频道添加实时监听
            for src in SOURCES:
                try:
                    entity = await client.get_entity(src)
                    @client.on(events.NewMessage(chats=entity))
                    async def handler(event, dest=dest, src_name=src):
                        await send_modified_message(client, dest, event.message)
                    print(f"✅ 已监听: {src}")
                except Exception as e:
                    print(f"❌ 监听失败 {src}: {e}")

            print(f"🔄 实时监听已启动，监控 {len(SOURCES)} 个源")
            print("📢 官方通知监听已启动，等待新消息...")

            # 心跳检测
            heartbeat_count = 0
            while True:
                try:
                    await asyncio.wait_for(client.get_me(), timeout=30)
                    heartbeat_count += 1
                    print(f"💓 心跳正常 [{heartbeat_count}]")
                    await asyncio.sleep(60)
                except asyncio.TimeoutError:
                    print("⚠️ 心跳超时，可能连接断开")
                    break
                except Exception as e:
                    print(f"⚠️ 心跳失败: {e}")
                    break

            await client.disconnect()

        except FloodWaitError as e:
            print(f"⏳ 主循环触发速率限制，等待 {e.seconds} 秒")
            await asyncio.sleep(e.seconds)
            continue
        except asyncio.TimeoutError:
            print("⏰ 连接超时，10秒后重连...")
            await asyncio.sleep(10)
            continue
        except (ConnectionError, OSError) as e:
            print(f"🔌 连接错误: {e}，10秒后重连...")
            await asyncio.sleep(10)
            continue
        except Exception as e:
            print(f"❌ 未知错误: {e}，10秒后重连...")
            await asyncio.sleep(10)
            continue

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("👋 用户停止")
    except Exception as e:
        print(f"💥 主程序异常: {e}")
        import traceback
        traceback.print_exc()