"""Parallel (multi-connection) transfers over MTProto.

The part-scheduling logic (``download_parts`` / ``upload_parts``) only talks to a
tiny ``Transport`` interface, so it is tested offline with fakes.  The Telethon
specific glue (``open_transports`` and friends) uses a few private Telethon
helpers; if anything in the fast path fails the wrappers ``download`` /
``upload`` fall back to Telethon's stock implementation, so the bot keeps
working (just slower).
"""
from __future__ import annotations

import asyncio
import itertools
import logging
import math
import os
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Awaitable, Callable, List, Optional, Protocol, Sequence, Tuple

from .errors import Cancelled

log = logging.getLogger(__name__)

DOWNLOAD_PART = 1024 * 1024  # upload.getFile allows up to 1 MiB per request
UPLOAD_PART = 512 * 1024  # upload.saveBigFilePart parts are at most 512 KiB
BIG_FILE_THRESHOLD = 10 * 1024 * 1024  # Telegram wants saveBigFilePart above 10 MB
RELAY_BUFFER_MIN = 4  # downloaded parts a streamed rename may hold back at least (2 per upload connection otherwise)
PART_TIMEOUT = 120.0
MAX_RETRIES = 5
RETRY_BASE_DELAY = 1.0

# Errors that retrying cannot fix: fail fast and let the caller fall back.
_FATAL_ERROR_NAMES = {
    "FileReferenceExpiredError", "FileReferenceInvalidError", "FileIdInvalidError", "LocationInvalidError",
    "OffsetInvalidError", "LimitInvalidError", "FilePartsInvalidError", "FilePartSizeInvalidError",
    "FileMigrateError", "AuthKeyUnregisteredError", "AuthKeyInvalidError", "SessionRevokedError",
    "UserDeactivatedError", "UserDeactivatedBanError",
}

ProgressCb = Callable[[int, int], None]


@dataclass
class TransferStats:
    """Which way the transfers since the start went: over several connections, or by Telethon itself (shown by /status)."""

    fast_down: int = 0  # downloads over several connections
    standard_down: int = 0  # downloads by Telethon itself: parallel transfer is off, or it failed
    fast_up: int = 0  # uploads over several connections
    small_up: int = 0  # uploads of files up to BIG_FILE_THRESHOLD, which Telethon's own upload handles on purpose
    standard_up: int = 0  # uploads by Telethon itself: parallel transfer is off, or it failed
    fallbacks: int = 0  # how often the parallel path failed and the standard one took over
    last_fallback: str = ""  # what failed last, e.g. "download: RuntimeError: ..."
    streamed: int = 0  # renames that went straight from the download into the upload, without touching the disk
    stream_fallbacks: int = 0  # streamed renames that failed, so the file was saved to disk first
    last_stream_fallback: str = ""  # what failed last, e.g. "RuntimeError: ..."

    def fell_back(self, what: str, exc: BaseException) -> None:
        self.fallbacks += 1
        self.last_fallback = f"{what}: {type(exc).__name__}: {exc}"

    def stream_failed(self, exc: BaseException) -> None:
        self.stream_fallbacks += 1
        self.last_stream_fallback = f"{type(exc).__name__}: {exc}"


stats = TransferStats()


class Transport(Protocol):
    async def get_file(self, location: Any, offset: int, limit: int) -> bytes: ...

    async def save_part(self, file_id: int, part: int, total_parts: int, data: bytes) -> None: ...

    async def close(self) -> None: ...


# ---------------------------------------------------------------------------
# positional file I/O that also works where os.pread/os.pwrite are missing
# ---------------------------------------------------------------------------


class _File:
    def __init__(self, path: Path | str, mode: str):
        flags = {"r": os.O_RDONLY, "w": os.O_RDWR | os.O_CREAT | os.O_TRUNC}[mode]
        flags |= getattr(os, "O_BINARY", 0)
        self.fd = os.open(str(path), flags, 0o644)
        self._lock = threading.Lock()
        self._positional = hasattr(os, "pwrite") and hasattr(os, "pread")

    def pwrite(self, data: bytes, offset: int) -> None:
        view = memoryview(data)
        if self._positional:
            written = 0
            while written < len(view):
                written += os.pwrite(self.fd, view[written:], offset + written)
            return
        with self._lock:
            os.lseek(self.fd, offset, os.SEEK_SET)
            written = 0
            while written < len(view):
                written += os.write(self.fd, view[written:])

    def pread(self, size: int, offset: int) -> bytes:
        if self._positional:
            return os.pread(self.fd, size, offset)
        with self._lock:
            os.lseek(self.fd, offset, os.SEEK_SET)
            return os.read(self.fd, size)

    def close(self) -> None:
        os.close(self.fd)


# ---------------------------------------------------------------------------
# generic engines
# ---------------------------------------------------------------------------


async def _retry(op: Callable[[], Awaitable[Any]], *, cancel: Optional[asyncio.Event], what: str) -> Any:
    delay = RETRY_BASE_DELAY
    for attempt in range(MAX_RETRIES + 1):
        if cancel is not None and cancel.is_set():
            raise Cancelled()
        try:
            return await asyncio.wait_for(op(), PART_TIMEOUT)
        except (asyncio.CancelledError, Cancelled):
            raise
        except Exception as exc:  # noqa: BLE001 - classified below
            if type(exc).__name__ in _FATAL_ERROR_NAMES or attempt >= MAX_RETRIES:
                raise
            log.warning("%s failed (%s: %s) - retry %d/%d", what, type(exc).__name__, exc, attempt + 1, MAX_RETRIES)
            await asyncio.sleep(delay)
            delay = min(delay * 2, 8.0)
    raise RuntimeError("unreachable")  # pragma: no cover


async def _run_workers(workers: Sequence[Callable[[], Awaitable[None]]]) -> None:
    tasks = [asyncio.ensure_future(w()) for w in workers]
    try:
        pending = set(tasks)
        while pending:
            done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_EXCEPTION)
            for task in done:
                if not task.cancelled() and task.exception() is not None:
                    raise task.exception()  # type: ignore[misc]
    finally:
        for task in tasks:
            task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)


async def download_parts(
    transports: Sequence[Transport],
    location: Any,
    size: int,
    dest: Path | str,
    *,
    part_size: int = DOWNLOAD_PART,
    progress: Optional[ProgressCb] = None,
    cancel: Optional[asyncio.Event] = None,
    max_bytes: Optional[int] = None,
) -> int:
    """Download ``size`` bytes (or the first ``max_bytes``) into ``dest``; returns bytes written."""
    total = size if max_bytes is None else min(size, max_bytes)
    out = _File(dest, "w")
    try:
        if total <= 0:
            return 0
        n_parts = math.ceil(total / part_size)
        counter = itertools.count()
        done = 0

        async def worker(transport: Transport) -> None:
            nonlocal done
            while True:
                if cancel is not None and cancel.is_set():
                    raise Cancelled()
                idx = next(counter)
                if idx >= n_parts:
                    return
                offset = idx * part_size
                want = min(part_size, total - offset)

                async def fetch() -> bytes:
                    data = await transport.get_file(location, offset, part_size)
                    if len(data) < want:
                        raise IOError(f"short read: got {len(data)} of {want} bytes at offset {offset}")
                    return data

                data = await _retry(fetch, cancel=cancel, what=f"download part {idx + 1}/{n_parts}")
                await asyncio.to_thread(out.pwrite, data[:want], offset)
                done += want
                if progress is not None:
                    progress(done, total)

        await _run_workers([lambda t=t: worker(t) for t in transports])
        if done != total:
            raise IOError(f"download incomplete: {done} of {total} bytes")
        return total
    finally:
        out.close()


async def upload_parts(
    transports: Sequence[Transport],
    path: Path | str,
    size: int,
    file_id: int,
    *,
    part_size: int = UPLOAD_PART,
    progress: Optional[ProgressCb] = None,
    cancel: Optional[asyncio.Event] = None,
) -> int:
    """Upload ``path`` with ``saveBigFilePart``; returns the number of parts."""
    n_parts = max(1, math.ceil(size / part_size))
    src = _File(path, "r")
    try:
        counter = itertools.count()
        done = 0

        async def worker(transport: Transport) -> None:
            nonlocal done
            while True:
                if cancel is not None and cancel.is_set():
                    raise Cancelled()
                idx = next(counter)
                if idx >= n_parts:
                    return
                offset = idx * part_size
                want = min(part_size, size - offset)
                data = await asyncio.to_thread(src.pread, want, offset)
                if len(data) != want:
                    raise IOError(f"read {len(data)} of {want} bytes at offset {offset}; was the file modified?")
                await _retry(
                    lambda d=data, i=idx: transport.save_part(file_id, i, n_parts, d),
                    cancel=cancel,
                    what=f"upload part {idx + 1}/{n_parts}",
                )
                done += want
                if progress is not None:
                    progress(done, size)

        await _run_workers([lambda t=t: worker(t) for t in transports])
        if done != size:
            raise IOError(f"upload incomplete: {done} of {size} bytes")
        return n_parts
    finally:
        src.close()


async def relay_parts(
    down: Sequence[Transport],
    up: Sequence[Transport],
    location: Any,
    size: int,
    file_id: int,
    *,
    part_size: int = DOWNLOAD_PART,
    send_size: int = UPLOAD_PART,
    progress: Optional[ProgressCb] = None,
    cancel: Optional[asyncio.Event] = None,
) -> int:
    """Copy ``size`` bytes of ``location`` into an upload without saving them anywhere; returns the number of upload parts.

    The ``down`` connections fetch ``part_size`` pieces and hand them to the ``up`` connections through a small
    buffer, which cut each piece into ``send_size`` parts for ``saveBigFilePart``.  Parts can be uploaded in any
    order, so the two sides only wait for each other when the buffer is full or empty: the file takes about as long as
    the slower direction, not as long as both together.  Progress counts what has been uploaded.
    """
    if size <= 0:
        raise ValueError("there is nothing to copy")
    if part_size % send_size:
        raise ValueError("a download part must hold a whole number of upload parts")
    if not down or not up:
        raise ValueError("a streamed copy needs connections both ways")
    n_fetch = math.ceil(size / part_size)
    n_parts = math.ceil(size / send_size)
    per_piece = part_size // send_size
    buffer: asyncio.Queue = asyncio.Queue(maxsize=max(RELAY_BUFFER_MIN, 2 * len(up)))
    fetch_counter, take_counter = itertools.count(), itertools.count()
    sent = 0

    async def fetch_worker(transport: Transport) -> None:
        while True:
            if cancel is not None and cancel.is_set():
                raise Cancelled()
            idx = next(fetch_counter)
            if idx >= n_fetch:
                return
            offset = idx * part_size
            want = min(part_size, size - offset)

            async def fetch() -> bytes:
                data = await transport.get_file(location, offset, part_size)
                if len(data) < want:
                    raise IOError(f"short read: got {len(data)} of {want} bytes at offset {offset}")
                return data

            data = await _retry(fetch, cancel=cancel, what=f"relay download part {idx + 1}/{n_fetch}")
            await buffer.put((idx, data[:want]))

    async def send_worker(transport: Transport) -> None:
        nonlocal sent
        while True:
            if cancel is not None and cancel.is_set():
                raise Cancelled()
            if next(take_counter) >= n_fetch:  # every downloaded part already has a sender
                return
            idx, data = await buffer.get()
            for start in range(0, len(data), send_size):
                part = idx * per_piece + start // send_size
                piece = data[start:start + send_size]
                await _retry(
                    lambda p=piece, i=part: transport.save_part(file_id, i, n_parts, p),
                    cancel=cancel,
                    what=f"relay upload part {part + 1}/{n_parts}",
                )
                sent += len(piece)
                if progress is not None:
                    progress(sent, size)

    await _run_workers([lambda t=t: fetch_worker(t) for t in down] + [lambda t=t: send_worker(t) for t in up])
    if sent != size:
        raise IOError(f"relay incomplete: {sent} of {size} bytes uploaded")
    return n_parts


async def close_all(transports: Sequence[Transport]) -> None:
    await asyncio.gather(*(t.close() for t in transports), return_exceptions=True)


# ---------------------------------------------------------------------------
# Telethon glue
# ---------------------------------------------------------------------------


class TelethonTransport:
    def __init__(self, client: Any, sender: Any):
        self._client = client
        self._sender = sender

    async def get_file(self, location: Any, offset: int, limit: int) -> bytes:
        from telethon.tl.functions.upload import GetFileRequest

        result = await self._client._call(self._sender, GetFileRequest(location=location, offset=offset, limit=limit))
        if not hasattr(result, "bytes"):
            raise IOError(f"unexpected upload.getFile result: {type(result).__name__}")
        return result.bytes

    async def save_part(self, file_id: int, part: int, total_parts: int, data: bytes) -> None:
        from telethon.tl.functions.upload import SaveBigFilePartRequest

        ok = await self._client._call(self._sender, SaveBigFilePartRequest(file_id, part, total_parts, data))
        if not ok:
            raise IOError("upload.saveBigFilePart returned False")

    async def close(self) -> None:
        await self._sender.disconnect()


async def _open_sender(client: Any, dc_id: int, auth_key: Any) -> Any:
    from telethon.network import MTProtoSender

    dc = await client._get_dc(dc_id)
    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
        )
    )
    return sender


async def open_transports(client: Any, dc_id: int, connections: int) -> List[TelethonTransport]:
    """Open up to ``connections`` extra MTProto connections to ``dc_id``."""
    from telethon.tl.alltlobjects import LAYER
    from telethon.tl.functions import InvokeWithLayerRequest
    from telethon.tl.functions.auth import ExportAuthorizationRequest, ImportAuthorizationRequest

    if dc_id == client.session.dc_id:
        auth_key = client.session.auth_key
        first = await _open_sender(client, dc_id, auth_key)
    else:
        # a different DC needs its own key, plus our authorization exported/imported once
        first = await _open_sender(client, dc_id, None)
        try:
            exported = await client(ExportAuthorizationRequest(dc_id))
            client._init_request.query = ImportAuthorizationRequest(id=exported.id, bytes=exported.bytes)
            await first.send(InvokeWithLayerRequest(LAYER, client._init_request))
        except BaseException:
            await first.disconnect()
            raise
        auth_key = first.auth_key
    senders = [first]
    if connections > 1:
        results = await asyncio.gather(
            *(_open_sender(client, dc_id, auth_key) for _ in range(connections - 1)), return_exceptions=True
        )
        for res in results:
            if isinstance(res, BaseException):
                log.warning("could not open an extra connection to DC %s: %s", dc_id, res)
            else:
                senders.append(res)
    log.info("opened %d connection(s) to DC %s", len(senders), dc_id)
    return [TelethonTransport(client, s) for s in senders]


async def fast_download(
    client: Any,
    document: Any,
    dest: Path | str,
    *,
    connections: int = 8,
    progress: Optional[ProgressCb] = None,
    cancel: Optional[asyncio.Event] = None,
    max_bytes: Optional[int] = None,
) -> int:
    from telethon import utils

    dc_id, location = utils.get_input_location(document)
    transports = await open_transports(client, dc_id, connections)
    try:
        written = await download_parts(
            transports, location, int(document.size), dest, progress=progress, cancel=cancel, max_bytes=max_bytes
        )
    finally:
        await close_all(transports)
    if max_bytes is None:  # reading the start of a file is not a transfer
        stats.fast_down += 1
    return written


async def fast_upload(
    client: Any,
    path: Path | str,
    file_name: str,
    *,
    connections: int = 8,
    progress: Optional[ProgressCb] = None,
    cancel: Optional[asyncio.Event] = None,
) -> Any:
    from telethon.tl.types import InputFileBig

    size = os.path.getsize(path)
    if size <= BIG_FILE_THRESHOLD:
        small = await client.upload_file(str(path), file_name=file_name, progress_callback=_checked(progress, cancel))
        stats.small_up += 1
        return small
    file_id = int.from_bytes(os.urandom(8), "big", signed=True)
    transports = await open_transports(client, client.session.dc_id, connections)
    try:
        parts = await upload_parts(transports, path, size, file_id, progress=progress, cancel=cancel)
    finally:
        await close_all(transports)
    stats.fast_up += 1
    return InputFileBig(id=file_id, parts=parts, name=file_name)


async def open_relay_transports(
    client: Any, dc_id: int, connections: int
) -> Tuple[List[TelethonTransport], List[TelethonTransport]]:
    """Connections to fetch from ``dc_id`` and connections to upload to our own DC, opened together."""
    results = await asyncio.gather(
        open_transports(client, dc_id, connections),
        open_transports(client, client.session.dc_id, connections),
        return_exceptions=True,
    )
    failed = next((r for r in results if isinstance(r, BaseException)), None)
    if failed is not None:
        for r in results:
            if not isinstance(r, BaseException):
                await close_all(r)
        raise failed
    return results[0], results[1]  # type: ignore[return-value]


async def relay(
    client: Any,
    document: Any,
    file_name: str,
    *,
    connections: int = 8,
    progress: Optional[ProgressCb] = None,
    cancel: Optional[asyncio.Event] = None,
) -> Any:
    """Upload a copy of ``document`` named ``file_name`` while it downloads, without saving it; returns an ``InputFileBig``.

    Telethon has no equivalent to fall back to here: when this fails the caller downloads and uploads the usual way.
    """
    from telethon import utils
    from telethon.tl.types import InputFileBig

    dc_id, location = utils.get_input_location(document)
    size = int(document.size)
    file_id = int.from_bytes(os.urandom(8), "big", signed=True)
    down, up = await open_relay_transports(client, dc_id, connections)
    try:
        parts = await relay_parts(down, up, location, size, file_id, progress=progress, cancel=cancel)
    finally:
        await close_all(down + up)
    return InputFileBig(id=file_id, parts=parts, name=file_name)


def is_file_part_error(exc: BaseException) -> bool:
    """Is this Telegram refusing an upload because of its parts (one missing, or the wrong size or number of them)?"""
    return type(exc).__name__.startswith("FilePart")


# ---------------------------------------------------------------------------
# public wrappers with fallback
# ---------------------------------------------------------------------------


def _checked(progress: Optional[ProgressCb], cancel: Optional[asyncio.Event]) -> ProgressCb:
    """Progress callback for Telethon's own transfers; raising aborts them."""

    def cb(current: int, total: int) -> None:
        if cancel is not None and cancel.is_set():
            raise Cancelled()
        if progress is not None:
            progress(current, total)

    return cb


def _unlink(path: Path | str) -> None:
    try:
        os.unlink(path)
    except FileNotFoundError:
        pass


async def download(
    client: Any,
    message: Any,
    dest: Path | str,
    *,
    fast: bool = True,
    connections: int = 8,
    progress: Optional[ProgressCb] = None,
    cancel: Optional[asyncio.Event] = None,
    max_bytes: Optional[int] = None,
) -> None:
    """Download a message's media to ``dest`` (optionally only the first ``max_bytes``)."""
    if fast and getattr(message, "document", None) is not None:
        try:
            await fast_download(
                client, message.document, dest, connections=connections, progress=progress, cancel=cancel,
                max_bytes=max_bytes,
            )
            return
        except Cancelled:
            raise
        except Exception as exc:  # noqa: BLE001
            log.exception("fast download failed (%s: %s); using Telethon's standard download", type(exc).__name__, exc)
            stats.fell_back("download", exc)
            _unlink(dest)
    if max_bytes is not None:
        request = 512 * 1024
        with open(dest, "wb") as fh:
            async for chunk in client.iter_download(
                message.document, request_size=request, limit=max(1, math.ceil(max_bytes / request))
            ):
                fh.write(chunk)
                if cancel is not None and cancel.is_set():
                    raise Cancelled()
            fh.truncate(min(fh.tell(), max_bytes))
        return
    await client.download_media(message, file=str(dest), progress_callback=_checked(progress, cancel))
    stats.standard_down += 1


async def upload(
    client: Any,
    path: Path | str,
    file_name: str,
    *,
    fast: bool = True,
    connections: int = 8,
    progress: Optional[ProgressCb] = None,
    cancel: Optional[asyncio.Event] = None,
) -> Any:
    """Upload ``path``; returns an ``InputFile``/``InputFileBig`` for ``send_file``."""
    if fast:
        try:
            return await fast_upload(client, path, file_name, connections=connections, progress=progress, cancel=cancel)
        except Cancelled:
            raise
        except Exception as exc:  # noqa: BLE001
            log.exception("fast upload failed (%s: %s); using Telethon's standard upload", type(exc).__name__, exc)
            stats.fell_back("upload", exc)
    standard = await client.upload_file(str(path), file_name=file_name, progress_callback=_checked(progress, cancel))
    stats.standard_up += 1
    return standard