utils.py

import base64
import hashlib
import html
import html.parser
import json
import logging
import os
import random
import re
import secrets
import string
import subprocess
import time
from collections import OrderedDict
from logging.handlers import RotatingFileHandler
from typing import Annotated, Literal

import msgspec
from requests import Session
from requests import get as requests_get

from enums import MediaType
from logging import Logger


def is_post_media_file_video(post):
    return post.get('ext', '').endswith(('webm', 'mp4', 'gif'))


def is_post_media_file_image(post):
    return post.get('ext', '').endswith(('jpg', 'png', 'jpeg', 'webp', 'bmp'))


def convert_to_asagi_capcode(a):
    if a:
        if a == "mod": return "M"
        if a == "admin": return "A"
        if a == "admin_highlight": return "A"
        if a == "developer": return "D"
        if a == "verified": return "V"
        if a == "founder": return "F"
        if a == "manager": return "G"

        return "M"

    return "N"


def convert_to_asagi_comment(a):
    if not a:
        return a

    # literal tags
    if "[" in a:
        a = re.sub(
            "\\[(/?(spoiler|code|math|eqn|sub|sup|b|i|o|s|u|banned|info|fortune|shiftjis|sjis|qstcolor))\\]",
            "[\\1:lit]",
            a
        )

    # abbr, exif, oekaki
    if "\"abbr" in a: a = re.sub("((<br>){0-2})?<span class=\"abbr\">(.*?)</span>", "", a)
    if "\"exif" in a: a = re.sub("((<br>)+)?<table class=\"exif\"(.*?)</table>", "", a)
    if ">Oek" in a: a = re.sub("((<br>)+)?<small><b>Oekaki(.*?)</small>", "", a)

    # banned
    if "<stro" in a:
        a = re.sub("<strong style=\"color: ?red;?\">(.*?)</strong>", "[banned]\\1[/banned]", a)

    # fortune
    if "\"fortu" in a:
        a = re.sub(
            "<span class=\"fortune\" style=\"color:(.+?)\"><br><br><b>(.*?)</b></span>",
            "\n\n[fortune color=\"\\1\"]\\2[/fortune]",
            a
        )

    # dice roll
    if "<b>" in a:
        a = re.sub(
            "<b>(Roll(.*?))</b>",
            "[b]\\1[/b]",
            a
        )

    # code tags
    if "<pre" in a:
        a = re.sub("<pre[^>]*>", "[code]", a)
        a = a.replace("</pre>", "[/code]")

    # math tags
    if "\"math" in a:
        a = re.sub("<span class=\"math\">(.*?)</span>", "[math]\\1[/math]", a)
        a = re.sub("<div class=\"math\">(.*?)</div>", "[eqn]\\1[/eqn]", a)

    # sjis tags
    if "\"sjis" in a:
        a = re.sub("<span class=\"sjis\">(.*?)</span>", "[shiftjis]\\1[/shiftjis]", a) # use [sjis] maybe?

    # quotes & deadlinks
    if "<span" in a:
        a = re.sub("<span class=\"quote\">(.*?)</span>", "\\1", a)

        # hacky fix for deadlinks inside quotes
        for idx in range(3):
            if not "deadli" in a: break
            a = re.sub("<span class=\"(?:[^\"]*)?deadlink\">(.*?)</span>", "\\1", a)

    # other links
    if "<a" in a:
        a = re.sub("<a(?:[^>]*)>(.*?)</a>", "\\1", a)

    # spoilers
    a = a.replace("<s>", "[spoiler]")
    a = a.replace("</s>", "[/spoiler]")

    # newlines
    a = a.replace("<br>", "\n")
    a = a.replace("<br/>", "\n")
    a = a.replace("<wbr>", "")

    a = html.unescape(a)

    return a


def get_asagi_value_media(post: dict) -> str | None:
    if post_has_file(post):
        return f"{post.get('tim')}{post.get('ext')}"


def get_asagi_value_preview(post: dict) -> str | None:
    if post_has_file(post):
        return f"{post.get('tim')}s.jpg"


post_has_file_keys = ('tim', 'ext', 'md5')
def post_has_file(post: dict) -> bool:
    return all(post.get(k) for k in post_has_file_keys)


def create_thumbnail(post: dict, full_path: str, thumb_path: str, logger=None):
    if is_post_media_file_video(post):
        create_thumbnail_from_video(full_path, thumb_path, logger=logger)
        return

    if is_post_media_file_image(post):
        create_thumbnail_from_image(full_path, thumb_path, logger=logger)
        return


def get_media_filename(post: dict, unescape_data_b4_db_write: bool) -> str | None:
    if post.get('ext') == 'deleted':
        return

    if post.get('filename') and post.get('ext'):
        if unescape_data_b4_db_write:
            return html.unescape(f"{post.get('filename')}{post.get('ext')}")
        return f"{post.get('filename')}{post.get('ext')}"


def get_d_board(post: dict, media_id: int | None = None, unescape_data_b4_db_write: bool=True):
    return {
        # 'doc_id': post.get('doc_id'), # autoincremented
        'media_id': media_id or 0, # inserted/updated by triggers
        'poster_ip': post.get('poster_ip', '0'),
        'num': post.get('no', 0),
        'subnum': post.get('subnum', 0),
        'thread_num': post.get('no') if post.get('resto') == 0 else post.get('resto'),
        'op': 1 if post.get('resto') == 0 else 0,
        'timestamp': post.get('time', 0),
        'timestamp_expired': post.get('archived_on', 0),
        'preview_orig': get_asagi_value_preview(post),
        'preview_w': post.get('tn_w', 0),
        'preview_h': post.get('tn_h', 0),
        'media_filename': get_media_filename(post, unescape_data_b4_db_write),
        'media_w': post.get('w', 0),
        'media_h': post.get('h', 0),
        'media_size': post.get('fsize', 0),
        'media_hash': post.get('md5'),
        'media_orig': get_asagi_value_media(post),
        'spoiler': post.get('spoiler', 0),
        'deleted': post.get('filedeleted', 0),
        'capcode': convert_to_asagi_capcode(post.get('capcode')),
        'email': post.get('email'),
        'name': html.unescape(post.get('name')) if post.get('name') and unescape_data_b4_db_write else None,
        'trip': post.get('trip'),
        'title': html.unescape(post.get('sub')) if post.get('sub') and unescape_data_b4_db_write else None,
        'comment': convert_to_asagi_comment(post.get('com')) if unescape_data_b4_db_write else post.get('com'),
        'delpass': post.get('delpass'),
        'sticky': post.get('sticky', 0),
        'locked': post.get('closed', 0),
        'poster_hash': post.get('id'),
        'poster_country': post.get('country_name'),
        'exif': json.dumps({'uniqueIps': int(post.get('unique_ips'))}) if post.get('unique_ips') else None,
    }


def get_thread_id_2_last_replies(catalog):
    thread_id_2_last_replies = {}
    for page in catalog:
        for thread in page['threads']:
            if thread.get('last_replies'):
                thread_id_2_last_replies[thread['no']] = thread.get('last_replies')
    return thread_id_2_last_replies


def get_d_image(post: dict, is_op: bool):
    return {
        # 'media_id': post.get('media_id'), # autoincremented
        'media_hash': post.get('md5'),
        'media': get_asagi_value_media(post),
        'preview_op': get_asagi_value_preview(post) if is_op else None,
        'preview_reply': get_asagi_value_preview(post) if not is_op else None,
        'total': 0,
        'banned': 0,
    }


PositiveInt = Annotated[int, msgspec.Meta(gt=0)]
NonNegativeInt = Annotated[int, msgspec.Meta(ge=0)]
ZeroOrOne = Annotated[int, msgspec.Meta(ge=0, le=1)]

MultiplyMediaTim = Annotated[str, msgspec.Meta(pattern=r'^$|^\d+(?:-\d+)?$')]

ExtLiteral = Literal['.jpg', '.jpeg', '.png', '.gif', '.pdf', '.swf', '.mp4', '.mp3', '.webm', '.webp', 'deleted']

StrLength32 = Annotated[str, msgspec.Meta(max_length=32)]
StrLength512 = Annotated[str, msgspec.Meta(min_length=0, max_length=512)]
StrLength16384 = Annotated[str, msgspec.Meta(min_length=0, max_length=16_384)]


class BasePost(msgspec.Struct, kw_only=True):
    no: PositiveInt
    resto: NonNegativeInt
    sticky: ZeroOrOne | None = None
    closed: ZeroOrOne | None = None
    now: StrLength512 | None = None
    time: PositiveInt
    name: StrLength512 | None = None
    trip: StrLength512 | None = None
    id: StrLength32 | None = None
    capcode: StrLength32 | None = None
    country: Annotated[str, msgspec.Meta(min_length=2, max_length=2)] | None = None
    country_name: StrLength512 | None = None
    sub: StrLength512 | None = None
    com: StrLength16384 | None = None
    tim: PositiveInt | MultiplyMediaTim | None = None
    filename: StrLength512 | None = None
    ext: ExtLiteral | None = None
    fsize: PositiveInt | None = None
    md5: Annotated[str, msgspec.Meta(min_length=24, max_length=24)] | None = None
    w: PositiveInt | None = None
    h: PositiveInt | None = None
    tn_w: PositiveInt | None = None
    tn_h: PositiveInt | None = None
    filedeleted: ZeroOrOne | None = None
    spoiler: ZeroOrOne | None = None
    custom_spoiler: Annotated[int, msgspec.Meta(ge=1, le=10)] | None = None
    m_img: ZeroOrOne | None = None

    replies: NonNegativeInt | None = None
    images: NonNegativeInt | None = None
    bumplimit: ZeroOrOne | None = None
    imagelimit: ZeroOrOne | None = None
    tag: StrLength512 | None = None
    semantic_url: StrLength512 | None = None
    since4pass: Annotated[int, msgspec.Meta(ge=2000, le=2099)] | None = None
    unique_ips: PositiveInt | None = None


class ChanPost(BasePost):
    '''https://github.com/4chan/4chan-API/blob/master/pages/Threads.md'''
    board_flag: StrLength512 | None = None
    flag_name: StrLength512 | None = None
    archived: ZeroOrOne | None = None
    archived_on: PositiveInt | None = None


class ChanThread(BasePost):
    '''https://github.com/4chan/4chan-API/blob/master/pages/Catalog.md'''
    last_modified: PositiveInt | None = None
    omitted_posts: NonNegativeInt | None = None
    omitted_images: NonNegativeInt | None = None
    last_replies: list[ChanPost] | None = None


def assert_thumbnail_deps(logger: Logger):
    ffmpeg_path = subprocess.run(['which', 'ffmpeg'], capture_output=True, text=True).stdout.strip()
    convert_path = subprocess.run(['which', 'convert'], capture_output=True, text=True).stdout.strip()
    logger.info(f'FFmpeg Path: {ffmpeg_path}')
    logger.info(f'Convert Path: {convert_path}')
    if not ffmpeg_path or not convert_path:
        raise ValueError(ffmpeg_path, convert_path)


def make_path(*filepaths):
    return os.path.join(os.path.abspath(os.path.dirname(__file__)), *filepaths)


def setup_logger(logger_name, log_file=False, stdout=True, file_rotate_size=1 * 1024 * 1024, max_files=3, log_level=logging.INFO):
    logger = logging.getLogger(logger_name)
    logger.setLevel(log_level)

    formatter = logging.Formatter('%(message)s')

    if stdout:
        console_handler = logging.StreamHandler()
        console_handler.setFormatter(formatter)
        logger.addHandler(console_handler)

    if log_file:
        file_handler = RotatingFileHandler(log_file, maxBytes=file_rotate_size, backupCount=max_files)
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)

    return logger


def write_json_obj_to_file(filepath: str, obj):
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, mode='w', encoding='utf-8') as f:
        json.dump(obj, f)


def read_json(fpath) -> dict:
    if not os.path.isfile(fpath):
        return None

    with open(fpath, mode='r', encoding='utf-8') as f:
        return json.load(f)


def sleep(t: int, add_random: bool=False):
    if add_random:
        t += random.uniform(0.0, 1.0)

    time.sleep(t)


def log_util(logger: Logger, message: str):
    if logger:
        logger.warning(message)
    else:
        print(message)


class TextExtractor(html.parser.HTMLParser):
    def __init__(self):
        super().__init__()
        self.text = []

    def handle_data(self, data: str):
        self.text.append(data)

    def get_text(self) -> str:
        return ' '.join(self.text)


def extract_text_from_html(html_str: str) -> str:
    if not html_str:
        return ''
    parser = TextExtractor()
    parser.feed(html_str)
    return html.unescape(parser.get_text())


def fullmatch_sub_and_com(post: dict, pattern: str) -> bool:
    """Compares a post's raw api data to patterns."""
    sub = post.get('sub')
    com = post.get('com')

    if sub:
        sub_text = html.unescape(sub)
        if re.fullmatch(pattern, sub_text, re.IGNORECASE):
            return True

    if com:
        com_text = extract_text_from_html(com)
        if re.fullmatch(pattern, com_text, re.IGNORECASE):
            return True

    return False


def get_n_random_chars(n: int) -> str:
    return ''.join(secrets.choice(string.ascii_letters) for _ in range(n))


def get_random_querystring() -> str:
    return f'{get_n_random_chars(5)}={get_n_random_chars(5)}'


def get_media_url(url_format: str, board: str, post: dict, media_type: MediaType) -> str:
    if media_type == MediaType.thumbnail:
        url = url_format.format(board=board, image_id=post['tim']) # ext is always .jpg
    elif media_type == MediaType.full_media:
        url = url_format.format(board=board, image_id=post['tim'], ext=post['ext'])
    else:
        raise ValueError(media_type)

    # try to avoid cloudflare's recompressed media
    url = f'{url}?{get_random_querystring()}'
    return url


def get_md5_b64_hash(content: bytes) -> str:
    return base64.b64encode(hashlib.md5(content).digest()).decode('ascii')


_b64_fs_trans = str.maketrans({
    '+': '-',
    '/': '_',
})


def get_fs_safe_b64(b64: str) -> str:
    return b64.translate(_b64_fs_trans)


def download_to_memory(url: str, session: Session, headers: dict | None) -> bytes | None:
    resp = session.get(url, headers=headers)
    if resp.status_code >= 400:
        return
    return resp.content


def fetch_media_bytes(
    url: str,
    ext: str,
    video_cooldown_sec: float=3.2,
    image_cooldown_sec: float=2.2,
    headers: dict | None=None,
    logger: Logger | None=None,
    session: Session | None=None,
    max_bytes: int | None=None,
) -> bytes | None:
    """Handles sleeping after requests"""

    resp = (session.get if session else requests_get)(url, headers=headers, stream=True)

    try:
        if resp.status_code != 200:
            log_util(logger, f'{url=} {resp.status_code=}')
            return

        length = resp.headers.get('content-length')
        if max_bytes is not None and length and int(length) > max_bytes:
            log_util(logger, f'Download stopped: {url=} bytes={length} > {max_bytes=}')
            sleep(2.0)
            return

        data = bytearray()

        for chunk in resp.iter_content(65_536):
            if not chunk:
                continue

            data.extend(chunk)

            if max_bytes is not None and len(data) > max_bytes:
                log_util(logger, f'Download stopped: {url=} bytes={len(data)} > {max_bytes=}')
                return

        if not data:
            return

    finally:
        # always return connection to session pool
        resp.close()

    # We only sleep if we decide to download a file
    time_to_sleep = video_cooldown_sec if is_video_path(ext) else image_cooldown_sec if is_image_path(ext) else 2.0
    sleep(time_to_sleep)

    return bytes(data)


video_exts = ('webm', 'mp4', 'gif')
def is_video_path(path: str) -> bool:
    return path.endswith(video_exts)


image_exts = ('jpg', 'jpeg', 'png', 'webp', 'bmp')
def is_image_path(path: str) -> bool:
    return path.endswith(image_exts)


def makedir_p(dir: str):
    if not os.path.isdir(dir):
        os.makedirs(dir, mode=0o775, exist_ok=True)


class MaxQueue:
    def __init__(self, boards, max_items_per_board=150*151*2):
        """
        threads_per_catalog = 150
        images_per_thread = 151
        """

        # use OrderedDict for lookup speed, rather than a list
        self.items = {b: OrderedDict() for b in boards}

        self.max_items_per_board = max_items_per_board

    def add(self, board: str, filepath: str):
        board_items = self.items[board]

        if filepath in board_items:
            return

        while len(board_items) >= self.max_items_per_board:
            board_items.popitem(last=False)

        board_items[filepath] = 1

    def __contains__(self, filepath: str) -> bool:
        return any(filepath in board_items for board_items in self.items.values())

    def __getitem__(self, board: str) -> OrderedDict:
        return self.items[board]


def create_thumbnail_from_video(video_path: str, out_path: str, width: int=400, height: int=400, quality: int=25, logger=None):
    """width and height form the max box boundary for the resulting image"""

    if not is_video_path(video_path):
        raise ValueError(video_path)

    command = f"""ffmpeg -hide_banner -loglevel error -ss 0 -i "{video_path}" -pix_fmt yuvj420p -q:v 2 -frames:v 1 -f image2pipe - | convert - -resize {width}x{height} -quality {quality} "{out_path}" """

    try:
        subprocess.run(command, shell=True, check=True, stdout=subprocess.DEVNULL)
        if logger:
            logger.info(f'    Created thumb {os.path.getsize(video_path) / 1024:.1f}kb -> {os.path.getsize(out_path) / 1024:.1f}kb')
    except Exception as e:
        if logger:
            logger.error(f'Error creating thumbnail from {video_path}\n{str(e)}')


def create_thumbnail_from_image(image_path: str, out_path: str, width: int=400, height: int=400, quality: int=25, logger=None):
    """width and height form the max box boundary for the resulting image"""

    if not is_image_path(image_path):
        raise ValueError(image_path)

    command = f"""convert "{image_path}" -resize {width}x{height} -quality {quality} "{out_path}" """

    try:
        subprocess.run(command, shell=True, check=True, stdout=subprocess.DEVNULL)
        if logger:
            logger.info(f'    Created thumb {os.path.getsize(image_path) / 1024:.1f}kb -> {os.path.getsize(out_path) / 1024:.1f}kb')
    except Exception as e:
        if logger:
            logger.error(f'    Error creating thumbnail from {image_path}\n{str(e)}')


def fetch_and_save_boards_json(filepath: str, url_boards: str, logger: Logger) -> dict:
    if not url_boards:
        return {}

    logger.info(f'Fetching {url_boards}...')
    resp = requests_get(url_boards, timeout=10)
    resp.raise_for_status()
    data = resp.json()

    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, 'w', encoding='utf-8') as f:
        json.dump(data, f)

    logger.info(f'Saved boards.json to {filepath}')
    return data


def load_boards_with_archive(boards_json: dict) -> set[str]:
    return {b['board'] for b in boards_json.get('boards', []) if b.get('is_archived')}