import asyncio
import os
import re
import time
import traceback
from collections import defaultdict

from dotenv import load_dotenv
from pyrogram import Client, filters
from pyrogram.errors import FloodWait
from pyrogram.types import Message


# =========================================================
# تنظیمات
# =========================================================

load_dotenv()

API_ID = int(os.getenv("API_ID"))
API_HASH = os.getenv("API_HASH")

BASE_DOWNLOAD_DIR = os.getenv("DOWNLOAD_DIR", "./downloads")

TEST_GROUP_ID = -1001155301071

# حداکثر تلاش برای خطاهای شبکه و خطاهای غیر از FloodWait
MAX_RETRIES = 5

# فاصله بین دانلودهای موفق
DOWNLOAD_DELAY = 5

# اندازه ثابت هر chunk برای محاسبه آفست
CHUNK_SIZE = 1024 * 1024  # 1 MiB


# =========================================================
# اتصال Pyrogram
# =========================================================

app = Client(
    "my_account",
    api_id=API_ID,
    api_hash=API_HASH,
    sleep_threshold=0
)


# =========================================================
# ابزارها
# =========================================================

def sanitize_filename(name: str) -> str:
    clean_name = re.sub(
        r'[\\/*?:"<>|\n\r]',
        "_",
        name
    ).strip()

    return clean_name[:100]


def fa_to_en_digits(text: str) -> str:
    fa_digits = "۰۱۲۳۴۵۶۷۸۹"
    ar_digits = "٠١٢٣٤٥٦٧٨٩"
    en_digits = "0123456789"

    for i in range(10):
        text = (
            text
            .replace(fa_digits[i], en_digits[i])
            .replace(ar_digits[i], en_digits[i])
        )

    return text


# =========================================================
# پیدا کردن چت
# =========================================================

async def force_find_chat(target_id: int):

    async for dialog in app.get_dialogs():

        if dialog.chat.id == target_id:
            return dialog.chat

    raise ValueError(
        f"چت با آیدی {target_id} در گفتگوهای اکانت یافت نشد."
    )


async def resolve_target(target_raw: str):

    target_raw = fa_to_en_digits(
        target_raw.strip()
    )

    # -----------------------------------------------------
    # لینک خصوصی
    # -----------------------------------------------------

    private_link_match = re.search(
        r"t\.me/c/(\d+)",
        target_raw
    )

    if private_link_match:

        channel_id = int(
            f"-100{private_link_match.group(1)}"
        )

        try:

            return await app.get_chat(
                channel_id
            )

        except Exception:

            return await force_find_chat(
                channel_id
            )

    # -----------------------------------------------------
    # لینک عمومی
    # -----------------------------------------------------

    public_link_match = re.search(
        r"t\.me/([^/]+)",
        target_raw
    )

    if public_link_match:

        return await app.get_chat(
            public_link_match.group(1)
        )

    # -----------------------------------------------------
    # ID یا username
    # -----------------------------------------------------

    try:

        if target_raw.lstrip("-").isdigit():

            val = int(target_raw)

            try:

                return await app.get_chat(
                    val
                )

            except Exception:

                pass

            if not str(val).startswith("-100"):

                target_id = int(
                    f"-100{val}"
                )

                try:

                    return await app.get_chat(
                        target_id
                    )

                except Exception:

                    return await force_find_chat(
                        target_id
                    )

            return await force_find_chat(
                val
            )

        return await app.get_chat(
            target_raw
        )

    except Exception as e:

        raise ValueError(
            f"امکان شناسایی چت وجود ندارد: {e}"
        )


# =========================================================
# دانلود یک فایل با Resume واقعی (اصلاح شده و سازگار با Pyrogram)
# =========================================================

async def download_single_file(
    message: Message,
    file_path: str,
    temp_file_path: str,
    file_name: str
):
    media = (
        message.audio
        or message.document
        or message.video
        or message.voice
    )

    expected_size = getattr(
        media,
        "file_size",
        0
    ) or 0

    network_retry = 0

    while True:

        try:

            current_size = 0

            if os.path.exists(temp_file_path):
                current_size = os.path.getsize(
                    temp_file_path
                )

            # اگر فایل قبلاً کامل شده
            if (
                expected_size > 0
                and current_size >= expected_size
            ):
                print(
                    f"📦 فایل موقت از قبل کامل است: "
                    f"{file_name}"
                )

                if os.path.exists(file_path):
                    os.remove(file_path)

                os.rename(
                    temp_file_path,
                    file_path
                )

                return True

            completed_chunks = (
                current_size // CHUNK_SIZE
            )

            resume_offset = (
                completed_chunks * CHUNK_SIZE
            )

            if current_size != resume_offset:
                truncated_bytes = (
                    current_size
                    - resume_offset
                )

                print(
                    f"🧹 حذف {truncated_bytes} بایت "
                    f"از انتهای ناقص فایل "
                    f"{file_name}"
                )

                with open(
                    temp_file_path,
                    "r+b"
                ) as f:
                    f.truncate(
                        resume_offset
                    )

            if (
                expected_size > 0
                and resume_offset >= expected_size
            ):
                if os.path.exists(file_path):
                    os.remove(file_path)

                os.rename(
                    temp_file_path,
                    file_path
                )

                return True

            print(
                f"📥 دانلود/ادامه فایل: {file_name}\n"
                f"   ├─ حجم فعلی: "
                f"{resume_offset / 1024 / 1024:.2f} MB\n"
                f"   └─ شروع از chunk: "
                f"{completed_chunks}"
            )

            # استفاده از stream_media بدون ارسال chunk_size (چون در برخی نسخه‌ها پشتیبانی نمی‌شود)
            with open(
                temp_file_path,
                "ab"
            ) as f:

                async for chunk in app.stream_media(
                    message,
                    offset=completed_chunks
                ):

                    if not chunk:
                        continue

                    f.write(chunk)

            if not os.path.exists(
                temp_file_path
            ):
                raise Exception(
                    "فایل موقت بعد از دانلود پیدا نشد."
                )

            final_size = os.path.getsize(
                temp_file_path
            )

            if final_size <= 0:
                raise Exception(
                    "فایل دانلود شده خالی است."
                )

            if (
                expected_size > 0
                and final_size != expected_size
            ):
                raise Exception(
                    f"دانلود کامل نشده است. "
                    f"دریافت شده: {final_size} "
                    f"از {expected_size} بایت."
                )

            print(
                f"✅ دانلود کامل شد: "
                f"{file_name} | "
                f"{final_size / 1024 / 1024:.2f} MB"
            )

            if os.path.exists(
                file_path
            ):
                os.remove(
                    file_path
                )

            os.rename(
                temp_file_path,
                file_path
            )

            print(
                f"📦 فایل نهایی ذخیره شد: "
                f"{file_path}"
            )

            return True

        except FloodWait as e:

            wait_time = int(
                e.value
            )

            print(
                f"⏳ FloodWait برای فایل "
                f"{file_name}: "
                f"{wait_time} ثانیه"
            )

            await asyncio.sleep(
                wait_time + 1
            )

            continue

        except Exception as e:

            network_retry += 1

            print(
                f"❌ خطا در دانلود "
                f"{file_name}: {e}"
            )

            traceback.print_exc()

            if network_retry >= MAX_RETRIES:

                print(
                    f"🚫 حداکثر تلاش‌ها برای "
                    f"{file_name} تمام شد."
                )

                if os.path.exists(
                    temp_file_path
                ):
                    try:
                        os.remove(
                            temp_file_path
                        )
                    except Exception:
                        pass

                return False

            retry_delay = (
                10 * network_retry
            )

            print(
                f"🔄 تلاش مجدد شماره "
                f"{network_retry} "
                f"بعد از {retry_delay} ثانیه..."
            )

            await asyncio.sleep(
                retry_delay
            )


# =========================================================
# دانلود کانال
# =========================================================

async def download_from_channel(
    chat_obj,
    max_files: int,
    status_msg: Message
):

    print(
        f"🔄 شروع عملیات دانلود برای: "
        f"{chat_obj.title} "
        f"(ID: {chat_obj.id})"
    )

    channel_name = sanitize_filename(
        chat_obj.title
        or str(chat_obj.id)
    )

    channel_dir = os.path.join(
        BASE_DOWNLOAD_DIR,
        channel_name
    )

    os.makedirs(
        channel_dir,
        exist_ok=True
    )

    await status_msg.edit_text(
        f"🔍 در حال بررسی آرشیو "
        f"«{chat_obj.title}»..."
    )

    try:

        raw_messages = []

        async for msg in app.get_chat_history(
            chat_obj.id
        ):

            if (
                msg.audio
                or msg.document
                or msg.video
                or msg.voice
            ):

                raw_messages.append(
                    msg
                )

        raw_messages.reverse()

        print(
            f"📊 تعداد کل فایل‌های یافت شده: "
            f"{len(raw_messages)}"
        )

        media_group_captions = {}

        for msg in raw_messages:

            if (
                msg.media_group_id
                and msg.caption
                and msg.caption.strip()
            ):

                if (
                    msg.media_group_id
                    not in media_group_captions
                ):

                    media_group_captions[
                        msg.media_group_id
                    ] = msg.caption.strip()

        media_group_counters = defaultdict(
            int
        )

        downloaded_count = 0
        already_downloaded_count = 0
        failed_count = 0

        last_update_time = time.time()

        for message in raw_messages:

            if (
                max_files
                and downloaded_count >= max_files
            ):
                break

            media = (
                message.audio
                or message.document
                or message.video
                or message.voice
            )

            if not media:
                continue

            file_ext = ""

            media_file_name = getattr(
                media,
                "file_name",
                None
            )

            if (
                media_file_name
                and "." in media_file_name
            ):

                file_ext = (
                    "."
                    + media_file_name.split(".")[-1]
                )

            elif message.audio:

                file_ext = ".mp3"

            elif message.video:

                file_ext = ".mp4"

            caption_to_use = None
            gallery_suffix = ""

            if message.media_group_id:

                media_group_counters[
                    message.media_group_id
                ] += 1

                gallery_suffix = (
                    f" _ gallery"
                    f"{media_group_counters[message.media_group_id]}"
                )

                caption_to_use = (
                    message.caption
                    or media_group_captions.get(
                        message.media_group_id
                    )
                )

            else:

                caption_to_use = (
                    message.caption
                )

            if (
                caption_to_use
                and caption_to_use.strip()
            ):

                file_name = (
                    f"{sanitize_filename(caption_to_use)}"
                    f"{gallery_suffix}"
                    f"{file_ext}"
                )

            else:

                raw_name = (
                    media_file_name
                    or f"file_{message.id}"
                )

                file_name = sanitize_filename(
                    raw_name
                )

                if (
                    file_ext
                    and not file_name.endswith(
                        file_ext
                    )
                ):

                    file_name += file_ext

            file_path = os.path.join(
                channel_dir,
                file_name
            )

            temp_file_path = (
                file_path
                + ".downloading"
            )

            if os.path.exists(
                file_path
            ):

                already_downloaded_count += 1

                print(
                    f"⏩ فایل قبلاً وجود دارد: "
                    f"{file_name}"
                )

                continue

            if (
                time.time()
                - last_update_time
                > 4
            ):

                await status_msg.edit_text(
                    f"🚀 در حال دانلود از "
                    f"«{chat_obj.title}»...\n\n"
                    f"📥 موفق: {downloaded_count}\n"
                    f"⏩ موجود: "
                    f"{already_downloaded_count}\n"
                    f"❌ ناموفق: "
                    f"{failed_count}\n\n"
                    f"📄 فایل جاری:\n"
                    f"`{file_name}`"
                )

                last_update_time = time.time()

            success = await download_single_file(
                message,
                file_path,
                temp_file_path,
                file_name
            )

            if success:

                downloaded_count += 1

                print(
                    f"🎉 فایل شماره "
                    f"{downloaded_count} "
                    f"با موفقیت دانلود شد."
                )

                await asyncio.sleep(
                    DOWNLOAD_DELAY
                )

            else:

                failed_count += 1

                print(
                    f"⚠️ فایل ناموفق بود: "
                    f"{file_name}"
                )

        await status_msg.edit_text(
            f"🎉 **دانلود کانال "
            f"«{chat_obj.title}» تکمیل شد!**\n\n"
            f"✅ جدید: {downloaded_count}\n"
            f"⏩ قبلی: {already_downloaded_count}\n"
            f"❌ ناموفق: {failed_count}\n\n"
            f"📂 پوشه ذخیره:\n"
            f"`{channel_dir}`"
        )

    except Exception as e:

        print(
            f"💥 خطای غیرمنتظره: {e}"
        )

        traceback.print_exc()

        try:

            await status_msg.edit_text(
                f"❌ خطای غیرمنتظره:\n"
                f"`{e}`"
            )

        except Exception:

            pass


# =========================================================
# دریافت دستور
# =========================================================

@app.on_message(
    filters.chat(TEST_GROUP_ID)
    & (filters.me | ~filters.me)
)
async def handle_commands(
    client: Client,
    message: Message
):

    raw_text = (
        message.text
        or message.caption
        or ""
    ).strip()

    text = fa_to_en_digits(
        raw_text
    )

    print(
        f"📩 دستور دریافتی در گروه تست: "
        f"{text}"
    )

    target_raw = None
    max_files = None

    if text.lower().startswith("dl "):

        parts = text.split()

        if len(parts) >= 2:

            target_raw = parts[1]

            if (
                len(parts) > 2
                and parts[2].isdigit()
            ):

                max_files = int(
                    parts[2]
                )

    elif "t.me/" in text:

        target_raw = text.split()[0]

    if target_raw:

        status_msg = await message.reply_text(
            f"⏳ در حال شناسایی هدف:\n"
            f"`{target_raw}`..."
        )

        try:

            chat_obj = await resolve_target(
                target_raw
            )

            print(
                f"✅ هدف پیدا شد: "
                f"{chat_obj.title} "
                f"(ID: {chat_obj.id})"
            )

            asyncio.create_task(
                download_from_channel(
                    chat_obj,
                    max_files,
                    status_msg
                )
            )

        except Exception as e:

            await status_msg.edit_text(
                f"❌ خطا در شناسایی "
                f"کانال/چت:\n"
                f"`{e}`"
            )


# =========================================================
# Start
# =========================================================

if __name__ == "__main__":

    print(
        "🚀 در حال استارت ربات..."
    )

    app.start()

    print(
        "👀 ربات آماده شنود دستورات است..."
    )

    from pyrogram.sync import idle

    idle()

    app.stop()