#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
=============================================================================
FastTelethon - موتور دانلود موازی فوق‌سریع MTProto تلگرام
ارتقایافته برای سرعت ۸+ مگابایت بر ثانیه با چانک‌های ۱ مگابایتی، پایپ‌لاین موازی
و ایزوله‌سازی کامل I/O دیسک از Event Loop تلگرام
=============================================================================
"""

import os
import sys
import math
import time
import inspect
import asyncio
from pathlib import Path
from typing import Optional, List, Union

# ماژول‌های داخلی پروژه در پوشه libs
CURRENT_DIR = Path(__file__).resolve().parent
LIBS_DIR = CURRENT_DIR / "libs"
if LIBS_DIR.exists() and str(LIBS_DIR) not in sys.path:
    sys.path.insert(0, str(LIBS_DIR))

try:
    from telethon import utils, helpers, functions, errors, types, TelegramClient
    from telethon.network import MTProtoSender
    from telethon.tl.alltlobjects import LAYER
    from telethon.tl.functions import InvokeWithLayerRequest
    from telethon.tl.functions.auth import ExportAuthorizationRequest, ImportAuthorizationRequest
    from telethon.tl.functions.upload import GetFileRequest
except ImportError as err:
    raise ImportError(f"خطا در بارگذاری وابستگی‌های تلتون برای FastTelethon: {err}")

# خطای ریت‌لیمیت پرمیوم
FloodPremiumWait = getattr(errors, 'FloodPremiumWaitError', errors.FloodWaitError)

# ابعاد چانک استاندارد برای حداکثر پهنای باند: ۱ مگابایت (۱۰۴۸۵۷۶ بایت)
# تلگرام تا سقف ۱ مگابایت را در GetFileRequest می‌پذیرد و آفست باید مضرب ۴ کیلوبایت باشد.
CHUNK_SIZE = 1024 * 1024


async def _create_single_sender(client: TelegramClient, dc_id: int, exported_auth=None) -> MTProtoSender:
    """ایجاد و مقداردهی اولیه یک کانکشن امن MTProto با دیتاسنتر هدف"""
    dc = await client._get_dc(dc_id)
    is_current_dc = (dc_id == client.session.dc_id)
    auth_key = client.session.auth_key if is_current_dc else None

    sender = MTProtoSender(auth_key, loggers=client._log)
    await sender.connect(client._connection(
        dc.ip_address,
        dc.port,
        dc.id,
        loggers=client._log,
        proxy=client._proxy,
        local_addr=client._local_addr
    ))

    # تلگرام برای شناسایی پروتکل روی هر اتصال جدید نیاز به InvokeWithLayerRequest دارد
    if not is_current_dc:
        if exported_auth is None:
            exported_auth = await client(ExportAuthorizationRequest(dc_id))
        query = ImportAuthorizationRequest(id=exported_auth.id, bytes=exported_auth.bytes)
    else:
        query = functions.help.GetNearestDcRequest()

    init_conn = functions.InitConnectionRequest(
        api_id=client.api_id,
        device_model=client._init_request.device_model,
        system_version=client._init_request.system_version,
        app_version=client._init_request.app_version,
        system_lang_code=client._init_request.system_lang_code,
        lang_pack=client._init_request.lang_pack,
        lang_code=client._init_request.lang_code,
        query=query
    )
    await sender.send(InvokeWithLayerRequest(LAYER, init_conn))
    return sender


async def download_file(
    client: TelegramClient,
    location,
    out: Union[str, Path, int],
    file_size: Optional[int] = None,
    progress_callback=None,
    connection_count: int = 8
) -> int:
    """
    دانلود موازی و چند کانکشنه فوق‌سریع بر پایه پایپ‌لاین ناهمگام و چانک‌های ۱ مگابایتی
    با خروج کامل I/O و fsync از EventLoop تلگرام
    """
    loc_for_dc = location
    if hasattr(location, "media"):
        loc_for_dc = location.media
    if hasattr(loc_for_dc, "document"):
        size = getattr(loc_for_dc.document, "size", None)
        loc_for_dc = loc_for_dc.document
    elif hasattr(loc_for_dc, "size"):
        size = getattr(loc_for_dc.size, None)
    else:
        size = file_size

    total_size = size or file_size or 0
    if not total_size or total_size <= 0:
        raise ValueError("حجم فایل برای دانلود چنداتصاله نامشخص است.")

    dc_id, input_location = utils.get_input_location(loc_for_dc)

    # تعداد قطعات ۱ مگابایتی
    part_count = math.ceil(total_size / CHUNK_SIZE)
    if part_count <= 0:
        return 0

    # تعداد اتصال‌های فعال موازی (حداکثر ۸ اتصال برای پایداری و حداکثر پهنای باند)
    connections = min(connection_count, 8)
    if total_size < 10 * 1024 * 1024:
        connections = min(connections, 4)
    else:
        connections = max(connections, 6)

    # باز کردن فایل روی دیسک
    is_own_fd = False
    fd = None
    if isinstance(out, (str, Path)):
        path_obj = Path(out).resolve()
        path_obj.parent.mkdir(parents=True, exist_ok=True)
        fd = os.open(str(path_obj), os.O_RDWR | os.O_CREAT | os.O_TRUNC, 0o666)
        is_own_fd = True
    elif hasattr(out, "fileno"):
        fd = out.fileno()
    elif isinstance(out, int):
        fd = out
    else:
        raise TypeError("خروجی دانلود باید مسیر فایل یا آبجکت دارای fileno باشد.")

    loop = asyncio.get_running_loop()

    # پیش‌تخصیص فضای فایل روی دیسک خارج از ایونت‌لوپ
    try:
        await loop.run_in_executor(None, os.ftruncate, fd, total_size)
    except Exception:
        pass

    senders: List[MTProtoSender] = []
    t_start_all = time.time()

    try:
        # ۱) ساخت موازی اتصالات + ExportAuthorization فقط یک‌بار
        exported_auth = None
        if dc_id != client.session.dc_id:
            try:
                exported_auth = await client(ExportAuthorizationRequest(dc_id))
            except Exception as ex:
                pass

        async def init_sender_safe():
            for att in range(3):
                try:
                    return await _create_single_sender(client, dc_id, exported_auth=exported_auth)
                except (errors.FloodWaitError, FloodPremiumWait) as fwe:
                    await asyncio.sleep(min(getattr(fwe, "seconds", 3), 6))
                except Exception:
                    await asyncio.sleep(0.25 * (att + 1))
            return None

        # اجرای موازی اتصال اولیه
        raw_res = await asyncio.gather(*(init_sender_safe() for _ in range(connections)), return_exceptions=True)
        for s in raw_res:
            if isinstance(s, MTProtoSender):
                senders.append(s)

        # تلاش مجدد برای اتصال‌های ناموفق تا اتصال‌ها به حداقل مطلوب برسند
        retry_rounds = 0
        min_desired = min(connections, max(4, connections - 2))
        while len(senders) < min_desired and retry_rounds < 2:
            retry_rounds += 1
            needed = connections - len(senders)
            retry_res = await asyncio.gather(*(init_sender_safe() for _ in range(needed)), return_exceptions=True)
            new_added = 0
            for s in retry_res:
                if isinstance(s, MTProtoSender):
                    senders.append(s)
                    new_added += 1
            if new_added == 0:
                break

        if not senders:
            raise RuntimeError("امکان برقراری اتصال موازی با دیتاسنتر تلگرام وجود ندارد.")

        # پر کردن صف چانک‌ها
        queue = asyncio.Queue()
        for i in range(part_count):
            queue.put_nowait(i)

        downloaded_chunks = set()
        chunks_lock = asyncio.Lock()
        downloaded_bytes = 0
        has_pwrite = hasattr(os, "pwrite")
        cancel_event = asyncio.Event()
        stalled_flag = {"v": False}
        last_progress = {"t": time.time()}

        write_queue = asyncio.Queue(maxsize=32)
        file_lock_sync = None
        if not has_pwrite:
            import threading
            file_lock_sync = threading.Lock()

        def _do_write_sync(offset, data):
            if has_pwrite:
                os.pwrite(fd, data, offset)
            else:
                with file_lock_sync:
                    os.lseek(fd, offset, os.SEEK_SET)
                    os.write(fd, data)

        # ورکر اختصاصی نوشتن روی دیسک به صورت کاملاً غیرمسدودکننده در ThreadPoolExecutor
        async def disk_writer():
            nonlocal downloaded_bytes
            while not cancel_event.is_set():
                try:
                    item = await asyncio.wait_for(write_queue.get(), timeout=0.5)
                except asyncio.TimeoutError:
                    if len(downloaded_chunks) >= part_count:
                        break
                    continue

                if item is None:
                    write_queue.task_done()
                    break

                part_idx, chunk = item
                offset = part_idx * CHUNK_SIZE
                if offset + len(chunk) > total_size:
                    chunk = chunk[:total_size - offset]

                # نوشتن در پس‌زمینه بدون کوچک‌ترین وقفه در ایونت‌لوپ تلگرام
                await loop.run_in_executor(None, _do_write_sync, offset, chunk)

                async with chunks_lock:
                    if part_idx not in downloaded_chunks:
                        downloaded_chunks.add(part_idx)
                        downloaded_bytes += len(chunk)
                        last_progress["t"] = time.time()

                write_queue.task_done()
                if len(downloaded_chunks) >= part_count:
                    break

        async def write_chunk(part_idx, chunk):
            """ارسال چانک به صف نوشتن دیسک بدون توقف ایونت‌لوپ شبکه"""
            await write_queue.put((part_idx, chunk))

        # ۲ و ۳) پایپ‌لاین: ۲ درخواست همزمان در پرواز برای هر اتصال
        MAX_IN_FLIGHT = 2

        async def worker(sender_obj):
            current_sender = sender_obj
            fails = 0
            in_flight = {}  # task -> part_idx

            while not cancel_event.is_set() and len(downloaded_chunks) < part_count:
                # پر کردن پنجره پایپ‌لاین
                while len(in_flight) < MAX_IN_FLIGHT and not cancel_event.is_set() and (len(downloaded_chunks) + len(in_flight) < part_count):
                    try:
                        part_idx = queue.get_nowait()
                    except asyncio.QueueEmpty:
                        break

                    if part_idx in downloaded_chunks:
                        continue

                    offset = part_idx * CHUNK_SIZE
                    if offset >= total_size:
                        async with chunks_lock:
                            downloaded_chunks.add(part_idx)
                        continue

                    req = functions.upload.GetFileRequest(
                        location=input_location,
                        offset=offset,
                        limit=CHUNK_SIZE
                    )
                    t = asyncio.create_task(client._call(current_sender, req))
                    in_flight[t] = part_idx

                if not in_flight:
                    if len(downloaded_chunks) >= part_count or cancel_event.is_set():
                        break
                    await asyncio.sleep(0.04)
                    continue

                # انتظار برای اتمام حداقل یکی از درخواست‌های موازی این اتصال
                done, _ = await asyncio.wait(in_flight.keys(), return_when=asyncio.FIRST_COMPLETED)
                for t in done:
                    p_idx = in_flight.pop(t)
                    try:
                        res = await t
                        if not hasattr(res, "bytes") or len(res.bytes) == 0:
                            raise RuntimeError("پاسخ خالی از تلگرام")
                        await write_chunk(p_idx, res.bytes)
                        fails = 0
                    except (errors.FloodWaitError, FloodPremiumWait) as fwe:
                        # ۲) گرفتن جداگانه FloodPremiumWaitError و FloodWaitError
                        await queue.put(p_idx)
                        sec = getattr(fwe, "seconds", 3)
                        await asyncio.sleep(min(sec, 8))
                    except asyncio.CancelledError:
                        await queue.put(p_idx)
                        raise
                    except Exception as ex:
                        await queue.put(p_idx)
                        fails += 1
                        await asyncio.sleep(min(0.2 * fails, 1.5))
                        if fails >= 4:
                            try:
                                new_sender = await _create_single_sender(client, dc_id, exported_auth=exported_auth)
                                old_sender = current_sender
                                current_sender = new_sender
                                senders.append(new_sender)
                                try:
                                    asyncio.create_task(old_sender.disconnect())
                                except Exception:
                                    pass
                                fails = 0
                            except Exception:
                                await asyncio.sleep(1.0)

            # لغو کارهای ناتمام در صورت خروج
            for t in list(in_flight.keys()):
                if not t.done():
                    t.cancel()

        # گزارش‌دهنده پیشرفت سبک
        reporter_stop = asyncio.Event()

        async def reporter():
            while not reporter_stop.is_set():
                try:
                    await asyncio.sleep(0.8)
                    if progress_callback:
                        cb_res = progress_callback(downloaded_bytes, total_size)
                        if inspect.isawaitable(cb_res):
                            await cb_res
                    if time.time() - last_progress["t"] > 45 and len(downloaded_chunks) < part_count:
                        stalled_flag["v"] = True
                        cancel_event.set()
                        break
                except (asyncio.CancelledError, KeyboardInterrupt):
                    cancel_event.set()
                    break
                except Exception as cb_err:
                    if "cancel" in str(cb_err).lower():
                        cancel_event.set()
                        break

        writer_task = asyncio.create_task(disk_writer())
        reporter_task = asyncio.create_task(reporter())
        worker_tasks = [asyncio.create_task(worker(s)) for s in list(senders)]

        try:
            try:
                await asyncio.gather(*worker_tasks)
            except (Exception, asyncio.CancelledError, BaseException):
                cancel_event.set()
                for t in worker_tasks:
                    if not t.done():
                        t.cancel()
                if not writer_task.done():
                    writer_task.cancel()
                await asyncio.gather(*worker_tasks, return_exceptions=True)
                raise

            # اطمینان از تخلیه و نوشتن تمامی چانک‌های بافر شده
            if not writer_task.done():
                await write_queue.join()
                await write_queue.put(None)
                await writer_task

            if cancel_event.is_set():
                if stalled_flag["v"]:
                    raise RuntimeError("دانلود موازی بیش از ۴۵ ثانیه بدون پیشرفت ماند")
                raise asyncio.CancelledError("Download cancelled")

            # بازبینی چانک‌های باقی‌مانده
            missing_parts = [i for i in range(part_count) if i not in downloaded_chunks]
            if missing_parts and not cancel_event.is_set():
                writer_task_missing = asyncio.create_task(disk_writer())
                for p_idx in missing_parts:
                    if cancel_event.is_set():
                        break
                    offset = p_idx * CHUNK_SIZE
                    for attempt in range(4):
                        try:
                            req = functions.upload.GetFileRequest(
                                location=input_location,
                                offset=offset,
                                limit=CHUNK_SIZE
                            )
                            if senders:
                                res = await asyncio.wait_for(client._call(senders[0], req), timeout=15.0)
                            else:
                                res = await asyncio.wait_for(client(req), timeout=15.0)
                            if hasattr(res, "bytes") and len(res.bytes) > 0:
                                await write_chunk(p_idx, res.bytes)
                                break
                        except (errors.FloodWaitError, FloodPremiumWait) as fe:
                            await asyncio.sleep(min(getattr(fe, "seconds", 3), 6))
                        except asyncio.CancelledError:
                            raise
                        except Exception:
                            await asyncio.sleep(0.5)

                if not writer_task_missing.done():
                    await write_queue.join()
                    await write_queue.put(None)
                    await writer_task_missing

            if cancel_event.is_set():
                if stalled_flag["v"]:
                    raise RuntimeError("دانلود موازی بیش از ۴۵ ثانیه بدون پیشرفت ماند")
                raise asyncio.CancelledError("Download cancelled")

            # ۴) خارج کردن fsync از EventLoop تلگرام
            if hasattr(os, "fsync"):
                try:
                    await loop.run_in_executor(None, os.fsync, fd)
                except Exception:
                    pass

            still_missing = [i for i in range(part_count) if i not in downloaded_chunks]
            if still_missing:
                raise RuntimeError(f"{len(still_missing)} قطعه از فایل دانلود نشد؛ رد کردن فایل ناقص")
        finally:
            cancel_event.set()
            reporter_stop.set()
            reporter_task.cancel()
            try:
                await reporter_task
            except (asyncio.CancelledError, Exception):
                pass
            if 'writer_task' in locals() and not writer_task.done():
                writer_task.cancel()
            if 'writer_task_missing' in locals() and not writer_task_missing.done():
                writer_task_missing.cancel()

        # گزارش نهایی به کالبک
        if progress_callback:
            try:
                cb_res = progress_callback(total_size, total_size)
                if inspect.isawaitable(cb_res):
                    await cb_res
            except Exception:
                pass

        # ۶) لاگ سرعت واقعی و تعداد اتصالات فعال در پایان
        t_elapsed = max(0.001, time.time() - t_start_all)
        real_mb_s = (downloaded_bytes / (1024 * 1024)) / t_elapsed
        active_cnt = len(senders)
        print(f"🏁 دانلود FastTelethon با موفقیت به پایان رسید | سرعت میانگین واقعی: {real_mb_s:.2f} MB/s | اتصالات فعال: {active_cnt} | زمان کل: {t_elapsed:.1f} ثانیه")

        return downloaded_bytes
    finally:
        # بستن تمیز همه اتصال‌های موازی
        for s in senders:
            try:
                asyncio.create_task(s.disconnect())
            except Exception:
                pass
        if is_own_fd and fd is not None:
            try:
                os.close(fd)
            except Exception:
                pass