#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
=============================================================================
FastTelethon - موتور دانلود موازی چنداتصاله MTProto بر پایه صف ناهمگام (Async Queue)
تست شده و کاملاً سازگار با تلگرام، بدون هنگ و بدون خطا در بخش‌های انتهایی فایل
=============================================================================
"""

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

# ماژول‌های داخلی پروژه در پوشه 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}")

# ابعاد چانک استاندارد تلگرام: ۵۱۲ کیلوبایت (۵۲۴۲۸۸ بایت)
CHUNK_SIZE = 512 * 1024


async def _create_single_sender(client: TelegramClient, dc_id: int) -> 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:
        for attempt in range(3):
            try:
                auth = await client(ExportAuthorizationRequest(dc_id))
                query = ImportAuthorizationRequest(id=auth.id, bytes=auth.bytes)
                break
            except Exception as ex:
                if attempt == 2:
                    raise ex
                await asyncio.sleep(0.25)
    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 = 12
) -> int:
    """
    دانلود موازی و چند کانکشنه بر پایه Worker Queue و ذخیره مستقیم روی دیسک با os.pwrite
    """
    # استخراج صحیح شیء مدیا و حجم فایل
    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, max(1, part_count))

    # باز کردن فایل روی دیسک
    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 باشد.")

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

    senders: List[MTProtoSender] = []
    try:
        # ایجاد کانکشن‌های موازی پایدار با دیتاسنتر
        for i in range(connections):
            try:
                s = await _create_single_sender(client, dc_id)
                senders.append(s)
            except Exception as conn_err:
                # اگر حداقل ۲ اتصال برقرار شده، با همان‌ها موازی دانلود کن
                if len(senders) >= 2:
                    break
                if not senders and i == connections - 1:
                    raise conn_err
                await asyncio.sleep(0.2)

        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")
        file_lock = asyncio.Lock()
        cancel_event = asyncio.Event()
        stalled_flag = {"v": False}
        last_progress = {"t": time.time()}

        async def write_chunk(part_idx, chunk):
            """نوشتن چانک روی دیسک و ثبت آن دقیقاً یک‌بار (بدون شمارش دوباره)"""
            nonlocal downloaded_bytes
            offset = part_idx * CHUNK_SIZE
            if offset + len(chunk) > total_size:
                chunk = chunk[:total_size - offset]
            if has_pwrite:
                os.pwrite(fd, chunk, offset)
            else:
                async with file_lock:
                    os.lseek(fd, offset, os.SEEK_SET)
                    os.write(fd, 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()

        async def worker(sender_obj):
            current_sender = sender_obj
            fails = 0
            # ورکر تا دانلود کامل «همه» چانک‌ها زنده می‌ماند (نه فقط تا خالی شدن صف)؛
            # چون ممکن است ورکر دیگری چانک شکست‌خورده را دوباره به صف برگرداند.
            while not cancel_event.is_set() and len(downloaded_chunks) < part_count:
                try:
                    part_idx = await asyncio.wait_for(queue.get(), timeout=0.5)
                except asyncio.TimeoutError:
                    continue

                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

                try:
                    req = functions.upload.GetFileRequest(
                        location=input_location,
                        offset=offset,
                        limit=CHUNK_SIZE
                    )
                    res = await asyncio.wait_for(client._call(current_sender, req), timeout=15.0)
                    if not hasattr(res, "bytes"):
                        # مثلاً FileCdnRedirect؛ این مسیر پشتیبانی نمی‌شود
                        raise RuntimeError("پاسخ نامعتبر (بدون bytes) از تلگرام")
                    chunk = res.bytes
                    if len(chunk) == 0 and offset < total_size:
                        raise RuntimeError("چانک خالی دریافت شد")
                    await write_chunk(part_idx, chunk)
                    fails = 0
                except errors.FloodWaitError as fwe:
                    await queue.put(part_idx)
                    await asyncio.sleep(min(getattr(fwe, "seconds", 5), 10))
                except asyncio.CancelledError:
                    raise
                except Exception:
                    # چانک را به صف برمی‌گردانیم و ورکر را «زنده نگه می‌داریم»
                    await queue.put(part_idx)
                    fails += 1
                    await asyncio.sleep(min(0.3 * fails, 2.0))
                    if fails >= 3:
                        # اتصال این ورکر خراب شده؛ یک اتصال تازه می‌سازیم
                        try:
                            new_sender = await _create_single_sender(client, dc_id)
                            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)

        # گزارش‌دهنده پیشرفت مستقل + نگهبان قفل‌شدن (watchdog)
        reporter_stop = asyncio.Event()

        async def reporter():
            while not reporter_stop.is_set():
                try:
                    await asyncio.sleep(0.6)
                    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

        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()
                await asyncio.gather(*worker_tasks, return_exceptions=True)
                raise

            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]
            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 asyncio.CancelledError:
                        raise
                    except Exception:
                        await asyncio.sleep(0.5)

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

            # فایل از قبل با ftruncate «پیش‌تخصیص» شده؛ پس بدون این بررسی فایل ناقص هم هم‌اندازهٔ کامل دیده می‌شود
            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:
            reporter_stop.set()
            reporter_task.cancel()
            try:
                await reporter_task
            except (asyncio.CancelledError, Exception):
                pass

        # گزارش قطعی ۱۰۰٪ در پایان دانلود
        if progress_callback:
            try:
                cb_res = progress_callback(total_size, total_size)
                if inspect.isawaitable(cb_res):
                    await cb_res
            except Exception:
                pass

        try:
            os.fsync(fd)
        except Exception:
            pass

        return downloaded_bytes

    finally:
        # بستن سریع فایل روی دیسک
        if is_own_fd and fd is not None:
            try:
                os.close(fd)
            except Exception:
                pass
            fd = None

        # بستن کانکشن‌های تلگرام در پس‌زمینه بدون ایجاد تاخیر در مرحله اتمام ۱۰۰٪
        for s in senders:
            try:
                asyncio.create_task(s.disconnect())
            except Exception:
                pass