max2telegram/app/media_transfer.py

from pathlib import Path
from uuid import uuid4

import aiohttp
import structlog

from app.models.tasks import MediaItem

logger = structlog.get_logger(__name__)

_MAX_CDN_HEADERS = {
  "User-Agent": (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
    "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
  ),
}


def tmp_dir(data_dir: str) -> Path:
  path = Path(data_dir) / "tmp"
  path.mkdir(parents=True, exist_ok=True)
  return path


def _item_field(item: MediaItem | dict, name: str, default=None):
  if isinstance(item, MediaItem):
    return getattr(item, name, default)
  return item.get(name, default)


def _default_file_name(kind: str) -> str:
  if "photo" in kind:
    return f"{uuid4().hex}.jpg"
  if "video" in kind:
    return f"{uuid4().hex}.mp4"
  return f"{uuid4().hex}.bin"


def resolve_file_name(item: MediaItem | dict) -> str:
  kind = _item_field(item, "kind", "") or ""
  file_name = _item_field(item, "file_name")
  if file_name:
    return file_name
  return _default_file_name(kind)


async def _download_url(
  session: aiohttp.ClientSession,
  url: str,
  dest: Path,
) -> bool:
  try:
    async with session.get(url, headers=_MAX_CDN_HEADERS) as resp:
      if resp.status != 200:
        logger.warning(
          "media_download_http_error",
          url=url,
          status=resp.status,
        )
        return False
      dest.write_bytes(await resp.read())
    return True
  except aiohttp.ClientError as exc:
    logger.warning("media_download_failed", url=url, error=str(exc))
    return False


async def _resolve_max_url(client, item: MediaItem | dict) -> str | None:
  chat_id = _item_field(item, "max_chat_id")
  message_id = _item_field(item, "max_message_id")
  if chat_id is None or message_id is None:
    return _item_field(item, "url")

  file_id = _item_field(item, "max_file_id")
  if file_id is not None:
    file_req = await client.get_file_by_id(chat_id, message_id, file_id)
    return file_req.url if file_req else None

  video_id = _item_field(item, "max_video_id")
  if video_id is not None:
    video = await client.get_video_by_id(chat_id, message_id, video_id)
    return video.url if video else None

  return _item_field(item, "url")


async def download_max_media(
  client,
  items: list[dict],
  dest_dir: Path,
) -> list[dict]:
  if not items:
    return []

  downloaded: list[dict] = []
  async with aiohttp.ClientSession() as session:
    for item in items:
      url = await _resolve_max_url(client, item)
      if not url:
        logger.warning(
          "max_media_url_unresolved",
          kind=item.get("kind"),
          max_file_id=item.get("max_file_id"),
          max_video_id=item.get("max_video_id"),
        )
        continue

      file_name = resolve_file_name(item)
      dest = dest_dir / f"{uuid4().hex}_{file_name}"
      if not await _download_url(session, url, dest):
        logger.warning(
          "max_media_download_skipped",
          kind=item.get("kind"),
          max_file_id=item.get("max_file_id"),
          max_video_id=item.get("max_video_id"),
        )
        continue

      stored = dict(item)
      stored["local_path"] = str(dest)
      stored.pop("url", None)
      downloaded.append(stored)

  return downloaded


async def download_media_item(
  bot,
  session: aiohttp.ClientSession,
  item: MediaItem | dict,
  dest_dir: Path,
  *,
  max_client=None,
) -> Path | None:
  local_path = _item_field(item, "local_path")
  if local_path:
    path = Path(local_path)
    if path.is_file():
      return path
    logger.warning("media_local_path_missing", path=local_path)

  file_id = _item_field(item, "file_id")
  file_name = resolve_file_name(item)

  if file_id:
    tg_file = await bot.get_file(file_id)
    if tg_file.file_path is None:
      return None
    dest = dest_dir / file_name
    await bot.download_file(tg_file.file_path, destination=dest)
    return dest

  url = _item_field(item, "url")
  if url is None and max_client is not None:
    url = await _resolve_max_url(max_client, item)

  if url:
    dest = dest_dir / file_name
    if await _download_url(session, url, dest):
      return dest
    return None

  return None


def cleanup_paths(paths: list[Path]) -> None:
  for path in paths:
    try:
      path.unlink(missing_ok=True)
    except OSError:
      logger.warning("media_temp_cleanup_failed", path=str(path))