"""
Downloads and uploads: one transfer lease each, then the chunks.

**The unit is the lease** (docs/MESHBAY_DESIGN.md §5.5): the node's record that
this connection is transferring something, opened before the first chunk and
closed after the last, on every path out. A lease may be granted at once or
queued behind the member's other transfers; the grant then arrives as a push.

Downloads are 1 MiB chunks, each encrypted under its own key derived from the
group key and the file's hash, and each independently decryptable — which is
what makes resuming a download cost nothing in the protocol. Uploads are 48 KB
chunks sealed under the group key, name and folder inside the seal; the node
keeps its own position, and a probe chunk asks for it, which is what resumes
an upload.
"""

import asyncio
import contextlib
import json
import logging
import math
import os
import secrets
import sys
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path

from meshbay_common import MNP_VERSION
from meshbay_common.groupbox import PURPOSE_UPLOAD, unseal
from meshbay_common.paths import long_path
from meshbay_common.protocol import UPLOAD_PROBE_INDEX, file_upload_wire
from meshbay_common.webcrypto import chunk_key_aes, decrypt_chunk_aes

from ..errors import NODE, CliError
from .session import GroupSession, NodeRefusal

log = logging.getLogger(__name__)

DOWNLOAD_CHUNK = 1024 * 1024     # the node serves files in 1 MiB chunks
UPLOAD_CHUNK = 48 * 1024         # what fits one DataChannel message after sealing
DOWNLOAD_WINDOW = 8              # chunk requests in flight per file
UPLOAD_WINDOW = 32               # chunks sent ahead of the node's acks
CHUNK_TIMEOUT = 60
# A queued lease waits behind the member's own transfers and other members';
# a queue that has said nothing for this long is not moving.
QUEUE_SILENCE = 600

PART_SUFFIX = ".meshbay-part"

Progress = Callable[[int], None]          # bytes just completed


class TransferError(CliError):
    """One file failed; the run goes on unless `fatal`."""

    def __init__(self, message: str, *, fatal: bool = False, code: str = ""):
        super().__init__(NODE, message)
        self.fatal = fatal
        self.refusal_code = code


# ── Leases ───────────────────────────────────────────────────────────────────

class Lease:
    def __init__(self, session: GroupSession, kind: str):
        self.session = session
        self.kind = kind
        self.tr = secrets.token_hex(16)
        self.closed_by_node = ""
        self._state: asyncio.Queue = asyncio.Queue()

    def _on_state(self, msg: dict) -> None:
        if msg.get("tr") != self.tr:
            return
        if msg.get("state") == "closed":
            self.closed_by_node = str(msg.get("reason") or "closed by the node")
        self._state.put_nowait(msg)

    async def open(self, size: int, chunks: int,
                   on_queued: Callable[[int], None] | None = None) -> None:
        ch = self.session.ch
        ch.on_push("transfer_state", self._on_state)
        reply = await ch.request({"type": "transfer_open", "v": MNP_VERSION, "tr": self.tr,
                                  "kind": self.kind, "bytes": size, "chunks": chunks},
                                 "transfer_state")
        state = reply.get("state")
        while state == "queued":
            if on_queued:
                on_queued(int(reply.get("ahead") or 0))
            try:
                reply = await asyncio.wait_for(self._state.get(), QUEUE_SILENCE)
            except TimeoutError:
                raise TransferError("The node's transfer queue did not move for "
                                    f"{QUEUE_SILENCE // 60} minutes.") from None
            state = reply.get("state")
        if state != "granted":
            raise TransferError(f"The node did not grant the transfer "
                                f"({reply.get('reason') or state}).")

    def check(self) -> None:
        if self.closed_by_node:
            raise TransferError(f"The node ended the transfer ({self.closed_by_node}).")

    async def close(self, reason: str) -> None:
        ch = self.session.ch
        ch.off_push("transfer_state", self._on_state)
        if self.closed_by_node or ch.closed.is_set():
            return
        with contextlib.suppress(CliError):
            await ch.request({"type": "transfer_close", "v": MNP_VERSION, "tr": self.tr,
                              "reason": reason}, "transfer_state", timeout=10)


# ── Download ─────────────────────────────────────────────────────────────────

@dataclass
class Resume:
    """What a `.meshbay-part` beside the target already holds, if it is this file."""

    part: Path
    sidecar: Path
    chunks_done: int

    @classmethod
    def find(cls, target: Path, entry: dict) -> "Resume":
        part = target.with_name(target.name + PART_SUFFIX)
        sidecar = target.with_name(target.name + PART_SUFFIX + ".json")
        done = 0
        try:
            meta = json.loads(sidecar.read_text(encoding="utf-8"))
            if meta.get("id") == entry["id"] and meta.get("size") == entry["size"]:
                # Every chunk but the last is whole, so the complete ones are
                # what the part's length says.
                done = os.path.getsize(long_path(part)) // DOWNLOAD_CHUNK
        except (OSError, ValueError):
            done = 0
        return cls(part, sidecar, done)


async def download(session: GroupSession, entry: dict, target: Path | None, *,
                   overwrite: bool = False, progress: Progress | None = None,
                   on_queued: Callable[[int], None] | None = None) -> int:
    """
    Fetch `entry` into `target` (None: stdout). Returns the bytes written.

    Written to `<name>.meshbay-part` and renamed at the end, so an interrupted run
    never leaves a truncated file under the real name; the part is resumed from
    its last whole chunk when the sidecar says it is this file.
    """
    size = int(entry["size"])
    chunks = max(1, math.ceil(size / DOWNLOAD_CHUNK))
    file_hash = bytes.fromhex(entry["id"])

    resume = None
    if target is not None:
        resume = Resume.find(target, entry)
        start = min(resume.chunks_done, chunks)
        await asyncio.to_thread(_prepare_part, resume, entry, start)
        out = open(long_path(resume.part), "r+b")
        out.seek(start * DOWNLOAD_CHUNK)
    else:
        start = 0
        out = sys.stdout.buffer
    if progress and start:
        progress(min(size, start * DOWNLOAD_CHUNK))

    lease = Lease(session, "download")
    reason = "error"
    try:
        await lease.open(size, chunks, on_queued)
        written = await _fetch_chunks(session, lease, entry["id"], file_hash, start, chunks,
                                      out, progress, streaming=target is None)
        if target is not None:
            await asyncio.to_thread(out.flush)
            await asyncio.to_thread(os.fsync, out.fileno())
            total = start * DOWNLOAD_CHUNK + written
            if total != size:
                raise TransferError(f"The node sent {total} bytes of a {size}-byte file.")
        reason = "done"
    except asyncio.CancelledError:
        reason = "cancelled"
        raise
    finally:
        if target is not None:
            out.close()
        await lease.close(reason)

    if target is not None:
        await asyncio.to_thread(_publish, resume, target, overwrite)
    return size if target is not None else written


def _prepare_part(resume: Resume, entry: dict, start: int) -> None:
    part = long_path(resume.part)
    if start == 0:
        with open(part, "wb"):
            pass
    else:
        os.truncate(part, start * DOWNLOAD_CHUNK)
    resume.sidecar.write_text(json.dumps({"id": entry["id"], "size": entry["size"]}),
                              encoding="utf-8")


def _publish(resume: Resume, target: Path, overwrite: bool) -> None:
    part, dest = long_path(resume.part), long_path(target)
    if overwrite:
        os.replace(part, dest)
    else:
        # A file that appeared meanwhile is not replaced: a hard link refuses
        # an existing target, where a rename would silently take its place.
        try:
            os.link(part, dest)
            os.unlink(part)
        except FileExistsError:
            raise TransferError(f"{target} appeared while downloading; the download "
                                f"is kept as {resume.part.name}.") from None
        except OSError:
            if os.path.exists(dest):
                raise TransferError(f"{target} already exists.") from None
            os.replace(part, dest)
    with contextlib.suppress(OSError):
        resume.sidecar.unlink()
    if sys.platform == "win32":
        mark_of_the_web(target)


async def _fetch_chunks(session: GroupSession, lease: Lease, file_id: str, file_hash: bytes,
                        start: int, chunks: int, out, progress: Progress | None,
                        *, streaming: bool) -> int:
    """Chunk requests `DOWNLOAD_WINDOW` ahead, written in order as they land."""
    ch, gek = session.ch, session.gek
    inflight: dict[int, asyncio.Task] = {}
    nxt, written = start, 0

    def ask(i: int) -> asyncio.Task:
        return asyncio.ensure_future(ch.request(
            {"type": "file_req", "v": MNP_VERSION, "file_id": file_id, "chunk_index": i,
             "tr": lease.tr}, "file_chunk", timeout=CHUNK_TIMEOUT))

    try:
        for i in range(start, chunks):
            while nxt < chunks and nxt - i < DOWNLOAD_WINDOW:
                inflight[nxt] = ask(nxt)
                nxt += 1
            try:
                msg = await inflight.pop(i)
            except NodeRefusal as e:
                raise TransferError(f"chunk {i}: {e.detail}", code=e.refusal_code) from e
            except CliError as e:
                raise TransferError(f"chunk {i}: {e.message}") from e
            lease.check()
            data = decrypt_chunk_aes(chunk_key_aes(gek, file_hash, i), msg["nonce"], msg["ct"])
            if streaming:
                out.write(data)
                out.flush()
            else:
                await asyncio.to_thread(out.write, data)
            written += len(data)
            if progress:
                progress(len(data))
    finally:
        for task in inflight.values():
            task.cancel()
        for task in inflight.values():
            with contextlib.suppress(BaseException):
                await task
    return written


def mark_of_the_web(path: Path) -> None:
    """ZoneId 3, as a browser marks a download, so SmartScreen and Protected View
    apply when it is opened. FAT and exFAT have no alternate streams: nothing to do."""
    with contextlib.suppress(OSError):
        with open(long_path(path) + ":Zone.Identifier", "w", encoding="ascii") as f:
            f.write("[ZoneTransfer]\r\nZoneId=3\r\n")


# ── Upload ───────────────────────────────────────────────────────────────────

async def create_dir(session: GroupSession, parent: str, name: str) -> str:
    """A folder `name` inside `parent` (root/folders). Returns its path."""
    try:
        ack = await session.ch.request({"type": "dir_create", "v": MNP_VERSION,
                                        "dir": parent, "name": name}, "dir_create_ack")
    except NodeRefusal as e:
        raise TransferError(f"{parent}/{name}: {e.detail}", code=e.refusal_code,
                            fatal=e.refusal_code in FATAL_UPLOAD_CODES) from e
    return str(ack.get("dir") or f"{parent}/{name}")


# Refusals that will not get better for the next file: the run stops.
FATAL_UPLOAD_CODES = frozenset({"disk_full", "root_unavailable", "root_read_only",
                                "no_writable_root", "no_roots", "no_group_key"})


async def upload(session: GroupSession, read: Callable[[int, int], bytes], size: int,
                 filename: str, folder: str, *, progress: Progress | None = None,
                 on_queued: Callable[[int], None] | None = None) -> str:
    """
    Send `size` bytes, read by `read(offset, length)`, as `filename` in `folder`.
    Returns the name the node stored it under — it never replaces a file, and
    takes a free name when this one is taken by the time the upload ends.
    """
    gek, gid, ch = session.gek, session.group_id, session.ch
    chunks = max(1, math.ceil(size / UPLOAD_CHUNK))
    upload_id = secrets.token_hex(16)

    def chunk(i: int, data: bytes) -> dict:
        return {**file_upload_wire(gek, gid, upload_id=upload_id, chunk_index=i,
                                   total_chunks=chunks, filename=filename, data=data,
                                   dir=folder), "tr": lease.tr}

    def refused(e: NodeRefusal) -> TransferError:
        return TransferError(f"{filename}: {e.detail}", code=e.refusal_code,
                             fatal=e.refusal_code in FATAL_UPLOAD_CODES)

    lease = Lease(session, "upload")
    reason = "error"
    try:
        await lease.open(size, chunks, on_queued)
        # Where am I? Only the node knows how much of this file it already holds.
        try:
            probe = await ch.request(chunk(UPLOAD_PROBE_INDEX, b""), "file_upload_ack",
                                     timeout=10)
            start = int(_open_ack(gek, gid, probe).get("resume_from") or 0)
        except NodeRefusal as e:
            raise refused(e) from e
        except CliError:
            start = 0
        if not 0 < start < chunks:
            start = 0
        if progress and start:
            progress(min(size, start * UPLOAD_CHUNK))

        pending: list[tuple[int, asyncio.Task]] = []
        last: dict = {}
        try:
            for i in range(start, chunks):
                while len(pending) >= UPLOAD_WINDOW:
                    last = await _ack(pending.pop(0), refused, progress, size)
                lease.check()
                data = await asyncio.to_thread(read, i * UPLOAD_CHUNK, UPLOAD_CHUNK)
                pending.append((i, asyncio.ensure_future(
                    ch.request(chunk(i, data), "file_upload_ack", timeout=CHUNK_TIMEOUT))))
            while pending:
                last = await _ack(pending.pop(0), refused, progress, size)
        finally:
            for _, task in pending:
                task.cancel()
            for _, task in pending:
                with contextlib.suppress(BaseException):
                    await task
        reason = "done"
    except asyncio.CancelledError:
        reason = "cancelled"
        raise
    finally:
        await lease.close(reason)
    return str(_open_ack(gek, gid, last).get("stored_as") or filename)


async def _ack(item: tuple[int, asyncio.Task], refused, progress: Progress | None,
               size: int) -> dict:
    i, task = item
    try:
        msg = await task
    except NodeRefusal as e:
        raise refused(e) from e
    except CliError as e:
        raise TransferError(f"chunk {i}: {e.message}") from e
    if progress:
        progress(min(UPLOAD_CHUNK, size - i * UPLOAD_CHUNK))
    return msg


def _open_ack(gek: bytes, group_id: str, msg: dict) -> dict:
    try:
        return unseal(gek, PURPOSE_UPLOAD, "file_upload_ack", group_id, msg)
    except Exception as e:
        raise TransferError("The node's acknowledgement does not open under the group "
                            "key.", fatal=True) from e