"""Stream Telegram media to disk with resume, disk-space checks and hashing."""
from __future__ import annotations

import asyncio
import errno
import hashlib
import logging
import os
import re
import shutil
from pathlib import Path
from typing import AsyncIterator, Callable, Protocol

from errors import DiskSpaceError, Interrupted, TransientError
from models import LocalFile, State, VideoRecord

log = logging.getLogger(__name__)

# Partial downloads are resumed from a 1 MiB boundary: Telegram's upload.getFile
# requires offsets aligned to the requested chunk size (512 KiB).
RESUME_ALIGNMENT = 1024 * 1024
HASH_BLOCK = 4 * 1024 * 1024
# Only files with exactly this shape are ever deleted by cleanup routines.
TEMP_NAME = re.compile(r"^(\d+)_(\d+)\.(part|video)$")


class MediaSource(Protocol):
    def iter_media_chunks(
        self, message_id: int, *, offset: int, expected_size: int | None, channel_id: int | None = None
    ) -> AsyncIterator[bytes]: ...


class _Hashes:
    def __init__(self, sha256: bool):
        self.md5 = hashlib.md5()
        self.sha256 = hashlib.sha256() if sha256 else None

    def update(self, data: bytes) -> None:
        self.md5.update(data)
        if self.sha256:
            self.sha256.update(data)

    def digest(self) -> tuple[str, str | None]:
        return self.md5.hexdigest(), self.sha256.hexdigest() if self.sha256 else None


def hash_file(path: Path, *, sha256: bool, limit: int | None = None, hashes: _Hashes | None = None) -> _Hashes:
    hashes = hashes or _Hashes(sha256)
    remaining = limit
    with open(path, "rb") as fh:
        while remaining is None or remaining > 0:
            block = fh.read(HASH_BLOCK if remaining is None else min(HASH_BLOCK, remaining))
            if not block:
                break
            hashes.update(block)
            if remaining is not None:
                remaining -= len(block)
    return hashes


class Downloader:
    def __init__(
        self,
        source: MediaSource,
        download_dir: Path,
        *,
        min_free_bytes: int = 0,
        max_temp_bytes: int = 0,
        compute_sha256: bool = True,
        disk_usage: Callable[[str], tuple] = shutil.disk_usage,
    ):
        self.source = source
        self.download_dir = download_dir
        self.min_free_bytes = min_free_bytes
        self.max_temp_bytes = max_temp_bytes
        self.compute_sha256 = compute_sha256
        self._disk_usage = disk_usage
        # Optional callback that frees space by removing abandoned temp files;
        # returns the number of files deleted. Set by the queue manager.
        self.reclaim: Callable[[], int] | None = None
        download_dir.mkdir(parents=True, exist_ok=True)
        self._root = download_dir.resolve()

    # Local files are named by message identity, never by caption text, so
    # captions cannot cause path problems.
    def complete_path(self, rec: VideoRecord) -> Path:
        return self.download_dir / f"{abs(rec.channel_id)}_{rec.message_id}.video"

    def partial_path(self, rec: VideoRecord) -> Path:
        return self.download_dir / f"{abs(rec.channel_id)}_{rec.message_id}.part"

    def has_complete_local(self, rec: VideoRecord) -> bool:
        path = self.complete_path(rec)
        if not path.is_file():
            return False
        expected = rec.expected_size or rec.local_size
        return expected is not None and path.stat().st_size == expected

    async def local_file(self, rec: VideoRecord) -> LocalFile:
        """Describe an already complete download, reusing stored hashes if valid."""
        path = self.complete_path(rec)
        size = path.stat().st_size
        if rec.md5 and rec.local_size == size:
            return LocalFile(path, size, rec.md5, rec.sha256)
        hashes = await asyncio.to_thread(hash_file, path, sha256=self.compute_sha256)
        md5, sha = hashes.digest()
        return LocalFile(path, size, md5, sha)

    def check_space(self, needed: int, *, exclude: Path | None = None) -> None:
        free = self._disk_usage(str(self.download_dir))[2]
        if free - needed < self.min_free_bytes:
            raise DiskSpaceError(
                f"not enough disk space: need {needed} bytes plus {self.min_free_bytes} reserve, {free} free",
                public="Not enough free disk space for the next video; free space and /resume.",
            )
        if self.max_temp_bytes:
            used = sum(
                p.stat().st_size for p in self.download_dir.iterdir() if p.is_file() and p != exclude
            )
            if used + needed > self.max_temp_bytes:
                raise DiskSpaceError(
                    f"MAX_TEMP_STORAGE exceeded: {used} used + {needed} needed > {self.max_temp_bytes}",
                    public="Temporary storage limit (MAX_TEMP_STORAGE) reached; /retry or clean up downloads.",
                )

    async def download(
        self,
        rec: VideoRecord,
        *,
        on_progress: Callable[[int, int | None], None],
        should_stop: Callable[[], bool],
    ) -> LocalFile:
        if self.has_complete_local(rec):
            return await self.local_file(rec)
        part = self.partial_path(rec)
        final = self.complete_path(rec)
        expected = rec.expected_size
        offset = part.stat().st_size if part.exists() else 0
        offset -= offset % RESUME_ALIGNMENT
        if expected is not None and offset > expected:
            offset = 0
        needed = (expected - offset) if expected else 0
        try:
            self.check_space(needed, exclude=part)
        except DiskSpaceError:
            if not (self.reclaim and self.reclaim()):
                raise
            self.check_space(needed, exclude=part)

        hashes = _Hashes(self.compute_sha256)
        if offset:
            log.info("Resuming download of message %s at byte %d", rec.message_id, offset)
            await asyncio.to_thread(hash_file, part, sha256=self.compute_sha256, limit=offset, hashes=hashes)
        written = offset
        chunks = self.source.iter_media_chunks(
            rec.message_id, offset=offset, expected_size=expected, channel_id=rec.channel_id
        )
        try:
            with open(part, "r+b" if part.exists() else "wb") as fh:
                fh.truncate(offset)
                fh.seek(offset)
                on_progress(written, expected)
                async for chunk in chunks:
                    fh.write(chunk)
                    hashes.update(chunk)
                    written += len(chunk)
                    on_progress(written, expected)
                    if expected is not None and written > expected:
                        break
                    if should_stop():
                        fh.flush()
                        raise Interrupted(f"download paused at byte {written}")
                fh.flush()
                os.fsync(fh.fileno())
        except OSError as exc:
            if exc.errno == errno.ENOSPC:
                raise DiskSpaceError("disk full while downloading", public="Local disk is full; free space and /resume.") from exc
            raise TransientError(f"download failed: {exc}", public="Download failed (I/O or network error)") from exc
        finally:
            await chunks.aclose()

        size = part.stat().st_size
        if expected is not None and size != expected:
            if size > expected:
                part.unlink(missing_ok=True)
            raise TransientError(
                f"incomplete download: {size} of {expected} bytes", public="Download incomplete or corrupted; will retry"
            )
        if size == 0:
            part.unlink(missing_ok=True)
            raise TransientError("downloaded file is empty", public="Downloaded file is empty; will retry")
        os.replace(part, final)
        md5, sha = hashes.digest()
        return LocalFile(final, size, md5, sha)

    def cleanup(self, rec: VideoRecord) -> int:
        """Delete this video's temporary files. Returns how many files were removed."""
        return sum(self.safe_unlink(path) for path in (self.complete_path(rec), self.partial_path(rec)))

    def safe_unlink(self, path: Path) -> bool:
        """Delete ``path`` only if it is a regular temp file directly inside the download dir."""
        try:
            if path.is_symlink() or path.parent.resolve() != self._root or not TEMP_NAME.match(path.name):
                log.error("Refusing to delete %s: not an application temp file", path.name)
                return False
            if not path.exists():
                return False
            path.unlink()
            log.info("Deleted temporary file %s", path.name)
            return True
        except OSError as exc:
            log.warning("Could not delete temporary file %s: %s", path.name, exc)
            return False

    def cleanup_orphans(self, db, *, keep_failed: bool = False) -> int:
        """Remove temp files that no unfinished item needs any more.

        A file is deleted only when it matches the application's temp-name pattern,
        lives directly in the download directory (symlinks are never followed), and
        its database record is missing or finished (COMPLETED / SKIPPED, or FAILED
        unless ``keep_failed``). Leftover ``.part`` files of fully downloaded items
        are removed too. Anything else in the directory is left untouched.
        """
        deleted = 0
        try:
            entries = list(os.scandir(self._root))
        except FileNotFoundError:
            return 0
        for entry in entries:
            match = TEMP_NAME.match(entry.name)
            if not match or entry.is_symlink() or not entry.is_file(follow_symlinks=False):
                continue
            channel_id, message_id, kind = -int(match.group(1)), int(match.group(2)), match.group(3)
            rec = db.get((channel_id, message_id))
            if rec is None or rec.state in (State.COMPLETED, State.SKIPPED):
                reason = "record finished" if rec else "no database record"
            elif rec.state == State.FAILED and not keep_failed:
                reason = "item failed (can be downloaded again)"
            elif kind == "part" and rec.state in (State.DOWNLOADED, State.UPLOADING):
                reason = "stale partial of a completed download"
            else:
                continue
            if self.safe_unlink(Path(entry.path)):
                deleted += 1
                log.info("Orphan cleanup: removed %s (%s)", entry.name, reason)
        return deleted