pyrogram-bridge/api_server.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-


# flake8: noqa
# pylint: disable=broad-exception-raised, raise-missing-from, too-many-arguments, redefined-outer-name
# pylint: disable=multiple-statements, logging-fstring-interpolation, trailing-whitespace, line-too-long
# pylint: disable=broad-exception-caught, missing-function-docstring, missing-class-docstring
# pylint: disable=f-string-without-interpolation
# pylance: disable=reportMissingImports, reportMissingModuleSource

import logging
import os
import mimetypes
from typing import List, Union, Any
from urllib.parse import quote

import json
from datetime import datetime
import time
from contextlib import asynccontextmanager
import random
import asyncio
import uvloop
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.background import BackgroundTask
import sys

import magic
from pyrogram import errors
from pyrogram.types import Message
from pyrogram.enums import MessageMediaType
from fastapi import FastAPI, HTTPException, Response, Request
from fastapi.responses import HTMLResponse, FileResponse, StreamingResponse
from telegram_client import TelegramClient
from config import get_settings, setup_logging
from rss_generator import generate_channel_rss, generate_channel_html
from post_parser import PostParser
from url_signer import verify_media_digest, generate_media_digest
from file_io import (DB_PATH, init_db_sync, get_all_media_file_ids_sync,
                     update_media_file_access_sync, remove_media_file_ids_sync,
                     get_mime_type_sync, set_mime_type_sync)

# Global python-magic instance for MIME type detection
magic_mime = magic.Magic(mime=True)

# Define custom exception for zero-size files
class ZeroSizeFileError(Exception):
    """Custom exception for zero-size files found or downloaded."""

class RequestLoggingMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        # Log only method and URL at debug level to avoid flooding logs on active RSS polling
        logger.debug(f"Request: {request.method} {request.url}")
        try:
            response = await call_next(request)
            logger.debug(f"Response status: {response.status_code}")
            return response
        except Exception as e:
            logger.error(f"Request processing error: {str(e)}")
            raise

logger = logging.getLogger(__name__)
if not logger.handlers: pass  

client = TelegramClient()
Config = get_settings()
HTTP_DOWNLOAD_SEMAPHORE = asyncio.Semaphore(3)  # semaphore for live HTTP media requests
BACKGROUND_DOWNLOAD_SEMAPHORE = asyncio.Semaphore(2)  # semaphore for background cache worker
download_queue = asyncio.Queue(maxsize=100)

@asynccontextmanager
async def lifespan(_: FastAPI):
    setup_logging(Config["log_level"])

    
    base_cache_dir = os.path.abspath("./data/cache")
    os.makedirs(base_cache_dir, exist_ok=True) # Create cache directory

    # Initialize SQLite database (creates table if not present)
    await asyncio.to_thread(init_db_sync, DB_PATH)

    await client.start()
    background_task = asyncio.create_task(cache_media_files()) # Start background task
    worker_task = asyncio.create_task(background_download_worker()) # Start download worker
    yield
    background_task.cancel() # Cleanup
    worker_task.cancel()
    try:
        await background_task
    except asyncio.CancelledError:
        pass
    try:
        await worker_task
    except asyncio.CancelledError:
        pass
    await client.stop()

app = FastAPI(title="Pyrogram Bridge", lifespan=lifespan)
app.add_middleware(RequestLoggingMiddleware)

def mask_sensitive_value(input_str: str) -> str:
    """Mask middle part of sensitive value, showing only first and last 4 characters"""
    if not input_str or len(input_str) < 8:
        return '***'
    visible_chars = 4
    return f"{input_str[:visible_chars]}{'*' * (len(input_str) - visible_chars * 2)}{input_str[-visible_chars:]}"


if __name__ == "__main__":
    import uvicorn
    
    setup_logging(Config["log_level"])
    
    logger.info("Starting server with configuration:")
    for key, value in Config.items():
        if any(sensitive in key.lower() for sensitive in ['token', 'tg_api_id', 'tg_api_hash']):
            logger.info(f"    {key}: {mask_sensitive_value(str(value))}")
        else:
            logger.info(f"    {key}: {value}")
    
    # Log uvloop status
    logger.info("    uvloop: enabled (asyncio speedup active)")
    
    try:
        uvicorn.run(
            "api_server:app",
            host=Config["api_host"],
            port=Config["api_port"],
            loop="uvloop"
        )
    except OSError as e:
        if "[Errno 98] Address already in use" in str(e):
            logger.critical(f"Port {Config['api_port']} is already in use. Exiting with code 1 to trigger Docker restart.")
            sys.exit(1)
        else:
            logger.critical(f"Failed to start server: {str(e)}")
            sys.exit(1)

async def find_file_id_in_message(message: Message, file_unique_id: str) -> Union[str, None]:
    """Find file_id by checking all possible media types in message"""
    if message.media == MessageMediaType.POLL:
        logger.debug(f"Message {message.id} is a poll, skipping media search")
        return None
        
    media_found = []
    if message.photo:
        media_found.append(f"photo ({message.photo.file_unique_id})")
        if message.photo.file_unique_id == file_unique_id:
            return message.photo.file_id
    if message.video:
        media_found.append(f"video ({message.video.file_unique_id})")
        if message.video.file_unique_id == file_unique_id:
            return message.video.file_id
    if message.animation:
        media_found.append(f"animation ({message.animation.file_unique_id})")
        if message.animation.file_unique_id == file_unique_id:
            return message.animation.file_id
    if message.video_note:
        media_found.append(f"video_note ({message.video_note.file_unique_id})")
        if message.video_note.file_unique_id == file_unique_id:
            return message.video_note.file_id
    if message.audio:
        media_found.append(f"audio ({message.audio.file_unique_id})")
        if message.audio.file_unique_id == file_unique_id:
            return message.audio.file_id
    if message.voice:
        media_found.append(f"voice ({message.voice.file_unique_id})")
        if message.voice.file_unique_id == file_unique_id:
            return message.voice.file_id
    if message.sticker:
        media_found.append(f"sticker ({message.sticker.file_unique_id})")
        if message.sticker.file_unique_id == file_unique_id:
            return message.sticker.file_id
    if message.web_page and message.web_page.photo:
        media_found.append(f"web_page.photo ({message.web_page.photo.file_unique_id})")
        if message.web_page.photo.file_unique_id == file_unique_id:
            return message.web_page.photo.file_id
    if message.document:
        media_found.append(f"document ({message.document.file_unique_id})")
        if message.document.file_unique_id == file_unique_id:
            return message.document.file_id
            
    # If we reached here, the file_unique_id was not found
    channel_id_log = message.chat.id if message.chat else 'unknown_chat'
    logger.warning(f"Could not find media with file_unique_id '{file_unique_id}' in message {message.id} (channel: {channel_id_log}). Found media: {', '.join(media_found) or 'None'}")
    return None


def delayed_delete_file(file_path: str, delay: int = 300) -> None:
    """
    Delete temporary file after a delay to ensure complete file delivery.
    Delay is set to 5 minutes by default.
    NOTE: Runs in a background thread via Starlette BackgroundTask.
    """
    time.sleep(delay)
    try:
        os.remove(file_path)
        logger.info(f"Deleted temporary file {file_path} after delay of {delay} seconds")
    except Exception as e:
        logger.error(f"Failed to delete temporary file {file_path}: {str(e)}")


async def prepare_file_response(file_path: str, request: Request, delete_after: bool = False,
                                media_key: tuple[str, int, str] | None = None) -> StreamingResponse:
    """Prepare a streaming file response with HTTP Range request support."""
    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="File not found")

    media_type: str | None = None

    if media_key is not None:
        # Try to load the cached MIME type from the database (avoids repeated python-magic I/O)
        channel_key, post_id_key, file_unique_id_key = media_key
        media_type = await asyncio.to_thread(get_mime_type_sync, DB_PATH, channel_key, post_id_key, file_unique_id_key)

    if not media_type:
        # Cache miss or no media_key — detect with python-magic in a thread to avoid blocking the event loop
        try:
            media_type = await asyncio.to_thread(magic_mime.from_file, file_path)
        except Exception as e:
            logger.warning(f"Failed to determine MIME type using python-magic: {str(e)}")
            media_type = None

        # Persist the detected MIME type so the next request can skip python-magic
        if media_type and media_key is not None:
            await asyncio.to_thread(set_mime_type_sync, DB_PATH, channel_key, post_id_key, file_unique_id_key, media_type)

    if not media_type: media_type, _ = mimetypes.guess_type(file_path)  # Fallback to mimetypes if python-magic failed
    if not media_type: media_type = "application/octet-stream"  # Final fallback to octet-stream

    logger.debug(f"Determined media type for {os.path.basename(file_path)}: {media_type}")

    try:
        total_size = os.path.getsize(file_path)
    except FileNotFoundError:
        raise HTTPException(status_code=404, detail="File not found")
    range_header = request.headers.get("range")

    if range_header:
        # Parse Range header according to RFC 7233
        try:
            range_value = range_header.strip()
            if not range_value.startswith("bytes="):
                raise ValueError("Only bytes ranges are supported")
            range_spec = range_value[len("bytes="):]
            if range_spec.startswith("-"):
                # Suffix range: bytes=-N (last N bytes)
                suffix_length = int(range_spec[1:])
                start = max(0, total_size - suffix_length)
                end = total_size - 1
            elif range_spec.endswith("-"):
                # Open-ended range: bytes=START-
                start = int(range_spec[:-1])
                end = total_size - 1
            else:
                # Full range: bytes=START-END
                start_str, end_str = range_spec.split("-", 1)
                start = int(start_str)
                end = int(end_str)
        except (ValueError, IndexError) as e:
            logger.debug(f"Invalid Range header '{range_header}': {e}")
            return Response(
                status_code=416,
                headers={"Content-Range": f"bytes */{total_size}"}
            )

        # Clamp end to file size - 1 (RFC 7233 allows end >= total_size)
        end = min(end, total_size - 1)

        # If start is beyond file size, return 416
        if start >= total_size:
            return Response(
                status_code=416,
                headers={"Content-Range": f"bytes */{total_size}"}
            )

        content_length = end - start + 1
        status_code = 206
        headers = {
            "Content-Disposition": (
                f"inline; filename=\"{os.path.basename(file_path)}\"; "
                f"filename*=UTF-8''{quote(os.path.basename(file_path))}"
            ),
            # Files are addressed by file_unique_id which is immutable in Telegram,
            # so it is safe to cache them aggressively on the client side.
            "Cache-Control": "public, max-age=86400, immutable",
            "Accept-Ranges": "bytes",
            "Content-Range": f"bytes {start}-{end}/{total_size}",
            "Content-Length": str(content_length),
        }
    else:
        # No Range header — serve full file with status 200
        start = 0
        end = total_size - 1
        status_code = 200
        headers = {
            "Content-Disposition": (
                f"inline; filename=\"{os.path.basename(file_path)}\"; "
                f"filename*=UTF-8''{quote(os.path.basename(file_path))}"
            ),
            # Files are addressed by file_unique_id which is immutable in Telegram,
            # so it is safe to cache them aggressively on the client side.
            "Cache-Control": "public, max-age=86400, immutable",
            "Accept-Ranges": "bytes",
            "Content-Length": str(total_size),
        }

    chunk_size = 64 * 1024  # 64 KB chunks

    async def file_chunk_generator():
        """Async generator that reads file in chunks; each chunk is a separate open/seek/read/close."""
        bytes_remaining = end - start + 1
        offset = start
        while bytes_remaining > 0:
            to_read = min(chunk_size, bytes_remaining)
            def read_at(path, off, size):
                # Open, seek, read, and close within a single thread call to avoid fd leaks
                with open(path, "rb") as f:
                    f.seek(off)
                    return f.read(size)
            data = await asyncio.to_thread(read_at, file_path, offset, to_read)
            if not data:
                break
            bytes_remaining -= len(data)
            offset += len(data)
            yield data

    background = BackgroundTask(delayed_delete_file, file_path) if delete_after else None
    return StreamingResponse(
        content=file_chunk_generator(),
        status_code=status_code,
        media_type=media_type,
        headers=headers,
        background=background,
    )

async def download_media_file(channel: Union[str, int], post_id: int, file_unique_id: str) -> tuple[Union[str, None], bool]:
    """
    Download media file from Telegram and save to cache
    Returns tuple of (file path, delete_after)
    """
    import time as _time
    _fn_start = _time.monotonic()
    base_cache_dir = os.path.abspath("./data/cache")
    
    # Create nested cache structure
    channel_dir = os.path.join(base_cache_dir, str(channel))
    post_dir = os.path.join(channel_dir, str(post_id))
    os.makedirs(post_dir, exist_ok=True)
    
    # Convert numeric channel ID to int if needed
    channel_id: Union[str, int] = channel
    if isinstance(channel, str) and channel.startswith('-100'):
        channel_id = int(channel)
    
    try:
        message = await client.safe_get_messages(channel_id, post_id)
    except asyncio.TimeoutError:
        logger.error(f"Timeout getting messages for {channel}/{post_id}")
        raise HTTPException(status_code=504, detail="Request timeout")
    
    if message.media == MessageMediaType.POLL:
        return None, False
    
    # Check if it is a video and if its size exceeds 100 MB
    is_large_video = False
    if message.video:
        try:
            if message.video.file_size > 100 * 1024 * 1024:
                is_large_video = True
        except Exception as e:
            logger.error(f"Failed to get video file size for message {post_id}: {str(e)}")
    
    if is_large_video:
        # For large video, download without permanent caching; use a temporary file
        temp_file_path = os.path.join(post_dir, f"temp_{file_unique_id}")
        if os.path.exists(temp_file_path):
            logger.info(f"Temporary file {temp_file_path} already exists, serving cached large video")
            return temp_file_path, False
        file_id = await find_file_id_in_message(message, file_unique_id)
        if not file_id:
            logger.error(f"Media with file_unique_id {file_unique_id} not found in message {post_id}")
            raise HTTPException(status_code=404, detail="File not found in message")
        logger.info(f"Downloading large video file {file_unique_id} to temporary path {temp_file_path}")
        try:
            file_path = await client.safe_download_media(file_id, temp_file_path)
        except asyncio.TimeoutError:
            logger.error(f"Timeout downloading large video {file_unique_id}")
            raise HTTPException(status_code=504, detail="Download timeout")
        logger.info(f"Downloaded large video file {file_unique_id} to temporary path {temp_file_path}")
        return file_path, False

    # Normal caching flow
    cache_path = os.path.join(post_dir, file_unique_id)
    if os.path.exists(cache_path):
        # Check if the cached file is zero size
        if os.path.getsize(cache_path) == 0:
            logger.warning(f"zero_size_cache_found: Found zero-size cached file: {cache_path}. Deleting and attempting redownload.")
            try:
                os.remove(cache_path)
                logger.info(f"Removed zero-size cached file: {cache_path}")
            except OSError as e:
                logger.error(f"cleanup_error: Failed to remove zero-size cached file {cache_path}: {e}")
            # Do not raise error here, proceed to download below
        else:
            # File exists and is not zero size, update access timestamp and return
            logger.info(f"Found cached media file: {cache_path}")
            try:
                await asyncio.to_thread(
                    update_media_file_access_sync,
                    DB_PATH, str(channel), post_id, file_unique_id,
                    datetime.now().timestamp()
                )
            except Exception as e:
                logger.error(f"Failed to update timestamp for {channel}/{post_id}/{file_unique_id}: {str(e)}")
            return cache_path, False

    file_id = await find_file_id_in_message(message, file_unique_id)
    if not file_id:
        error_message = f"Media with file_unique_id {file_unique_id} not found in message {post_id} for channel {channel}"
        logger.error(error_message)

        # Remove the invalid entry from the SQLite database
        try:
            await asyncio.to_thread(
                remove_media_file_ids_sync,
                DB_PATH, [(str(channel), post_id, file_unique_id)]
            )
            logger.info(f"Removed invalid entry for {channel}/{post_id}/{file_unique_id} from SQLite")
        except Exception as e:
            logger.error(f"Failed to remove entry for {channel}/{post_id}/{file_unique_id} from SQLite: {str(e)}")

        raise HTTPException(status_code=404, detail="File not found in message")
        
    try:
        file_path = await client.safe_download_media(file_id, cache_path)
    except asyncio.TimeoutError:
        logger.error(f"Timeout downloading media {file_unique_id}")
        raise HTTPException(status_code=504, detail="Download timeout")
    
    # Check if the downloaded file exists and has a size greater than zero
    if not file_path or not os.path.exists(file_path) or os.path.getsize(file_path) == 0:
        logger.error(f"download_failed_zero_size: Downloaded file {file_unique_id} for {channel}/{post_id} is zero size or missing.")
        # Attempt to clean up the invalid file
        if file_path and os.path.exists(file_path):
            try:
                os.remove(file_path)
                logger.info(f"Removed zero-size file: {file_path}")
            except OSError as e:
                logger.error(f"cleanup_error: Failed to remove zero-size file {file_path}: {e}")
        # Raise an error to indicate download failure
        # Raise specific error
        raise ZeroSizeFileError(f"Downloaded file {file_unique_id} for {channel}/{post_id} is zero size or missing after download attempt.")

    logger.info(f"Downloaded media file {file_unique_id} to {cache_path}")
    return file_path, False


def remove_old_cached_files_sync(media_files: list, cache_dir: str) -> tuple[list, int]:
    """
    Remove files that haven't been accessed for more than 20 days
    Returns tuple of (updated media files list, number of removed files)
    """
    current_time = datetime.now().timestamp()
    updated_media_files = []
    files_removed = 0

    for file_data in media_files:
        try:
            channel = file_data.get('channel')
            post_id = file_data.get('post_id')
            file_unique_id = file_data.get('file_unique_id')
            
            if not all([channel, post_id, file_unique_id]):
                continue

            last_access_time = file_data.get('added', 0)
            days_since_access = (current_time - last_access_time) / (24 * 3600)

            if days_since_access > 20:
                cache_path = os.path.join(cache_dir, str(channel), str(post_id), file_unique_id)
                
                if os.path.exists(cache_path):
                    try:
                        os.remove(cache_path)
                        # try to remove empty parent directories
                        post_dir = os.path.dirname(cache_path)
                        channel_dir = os.path.dirname(post_dir)
                        if not os.listdir(post_dir):
                            os.rmdir(post_dir)
                        if not os.listdir(channel_dir):
                            os.rmdir(channel_dir)
                            
                        files_removed += 1
                        logger.info(f"Removed old cached file: {cache_path}, last access {days_since_access:.1f} days ago")
                    except Exception as e:
                        logger.error(f"Failed to remove cached file {cache_path}: {str(e)}")
                        updated_media_files.append(file_data)
                continue

            updated_media_files.append(file_data)

        except Exception as e:
            logger.error(f"Error processing cache entry {file_data}: {str(e)}")
            continue

    # Additionally, remove temporary files with prefix "temp_" if they are older than a threshold (1 hour)
    temp_threshold = 3600  # 1 hour in seconds
    for root, _, files in os.walk(cache_dir):
        for file in files:
            if file.startswith("temp_"):
                file_path = os.path.join(root, file)
                try:
                    file_mod_time = os.path.getmtime(file_path)
                    if time.time() - file_mod_time > temp_threshold:
                        # Get channel/post/file_id from path
                        rel_path = os.path.relpath(os.path.dirname(file_path), cache_dir)
                        channel, post_id = rel_path.split(os.sep)
                        file_unique_id = file[5:]  # Remove 'temp_' prefix
                        
                        # Remove file
                        os.remove(file_path)
                        files_removed += 1
                        logger.info(f"Removed temporary file: {file_path}")
                        
                        # Also remove the temporary file entry from the in-memory list
                        updated_media_files = [
                            f for f in updated_media_files 
                            if not (f.get('channel') == channel and 
                                f.get('post_id') == int(post_id) and 
                                f.get('file_unique_id') == file_unique_id)
                        ]
                        
                except Exception as e:
                    logger.error(f"Failed to remove temporary file {file_path}: {str(e)}")

    return updated_media_files, files_removed


async def download_new_files(media_files: list, cache_dir: str) -> None:
    """
    Queue files that are not in cache yet for background download
    """
    if not media_files:
        logger.info("No media files found for download")
        return

    files_queued = 0
    for file_data in media_files:
        try:
            channel = file_data.get('channel')
            post_id = file_data.get('post_id')
            file_unique_id = file_data.get('file_unique_id')
            
            if not all([channel, post_id, file_unique_id]):
                logger.error(f"Invalid file data: {file_data}")
                continue
                
            channel_dir = os.path.join(cache_dir, str(channel))
            post_dir = os.path.join(channel_dir, str(post_id))
            os.makedirs(post_dir, exist_ok=True)
            
            # Skip if temp file exists (large videos)
            temp_path = os.path.join(post_dir, f"temp_{file_unique_id}")
            if os.path.exists(temp_path):
                continue
                
            cache_path = os.path.join(post_dir, file_unique_id)
            if not os.path.exists(cache_path):
                try:
                    await download_queue.put((channel, post_id, file_unique_id))
                    files_queued += 1
                    logger.debug(f"Queued for background download: {channel}/{post_id}/{file_unique_id}")
                except asyncio.QueueFull:
                    logger.warning(f"Download queue is full, skipping {channel}/{post_id}/{file_unique_id}")
                    break
        
        except Exception as e:
            logger.error(f"Failed to queue download for {channel}/{post_id}/{file_unique_id}: {str(e)}")
            continue

    if files_queued > 0:
        logger.info(f"Queued {files_queued} files for background download")


async def background_download_worker():
    """Worker that processes downloads from queue"""
    while True:
        try:
            channel, post_id, file_unique_id = await download_queue.get()
            logger.info(f"Background download: {channel}/{post_id}/{file_unique_id}")
            
            try:
                async with BACKGROUND_DOWNLOAD_SEMAPHORE:  # limit concurrent background downloads
                    await download_media_file(channel, post_id, file_unique_id)
                await asyncio.sleep(2)
            except Exception as e:
                logger.error(f"Background download error for {channel}/{post_id}/{file_unique_id}: {e}")
                
        except Exception as e:
            logger.error(f"Background download worker error: {e}")
        finally:
            download_queue.task_done()

async def cache_media_files() -> None:
    """Background task for cache management: removes old files and downloads new ones"""
    delay = 60
    while True:
        try:
            # Load all media file ID records from SQLite
            media_files = await asyncio.to_thread(get_all_media_file_ids_sync, DB_PATH)

            cache_dir = os.path.abspath("./data/cache")
            updated_media_files, files_removed = await asyncio.to_thread(remove_old_cached_files_sync, media_files, cache_dir)

            if files_removed > 0:
                try:
                    # Determine which entries were removed and delete them from SQLite
                    updated_set = {
                        (r['channel'], r['post_id'], r['file_unique_id'])
                        for r in updated_media_files
                    }
                    removed_entries = [
                        (r['channel'], r['post_id'], r['file_unique_id'])
                        for r in media_files
                        if (r['channel'], r['post_id'], r['file_unique_id']) not in updated_set
                    ]
                    if removed_entries:
                        await asyncio.to_thread(remove_media_file_ids_sync, DB_PATH, removed_entries)
                    logger.info(f"Removed {files_removed} old files from cache")
                except Exception as e:
                    logger.error(f"Failed to remove old entries from SQLite: {str(e)}")

            await download_new_files(updated_media_files, cache_dir)
            await asyncio.sleep(delay)  # Check every delay seconds

        except Exception as e:
            logger.error(f"Cache media files error: {str(e)}")
            await asyncio.sleep(delay)


def calculate_cache_stats() -> dict[str, Any]:
    """
    Calculate cache statistics including file count, total size in MB, and time difference in days.
    Returns a dictionary with keys: 'cache_files_count', 'cache_total_size_mb', 'cache_time_diff_days', 'channels'.
    """
    base_cache_dir = os.path.abspath("./data/cache")
    cache_files_count = 0
    cache_total_size_bytes = 0
    channels_stats = {}
    
    if os.path.isdir(base_cache_dir):
        # Recursively walk through all subdirectories
        for root, _, files in os.walk(base_cache_dir):
            for current_file in files:
                file_path = os.path.join(root, current_file)
                file_size = os.path.getsize(file_path)
                cache_files_count += 1
                cache_total_size_bytes += file_size
                
                # Calculate per-channel statistics
                rel_path = os.path.relpath(root, base_cache_dir)
                channel = rel_path.split(os.sep, maxsplit=1)[0]  # First directory is channel
                
                if channel not in channels_stats:
                    channels_stats[channel] = {
                        'files_count': 0,
                        'size_mb': 0.0
                    }
                
                channels_stats[channel]['files_count'] += 1
                channels_stats[channel]['size_mb'] = round(
                    channels_stats[channel]['size_mb'] + (file_size / (1024 * 1024)), 2
                )
                    
        cache_total_size_mb = round(cache_total_size_bytes / (1024 * 1024), 2)  # rounded size in MB
    else:
        cache_files_count = 0
        cache_total_size_mb = 0

    cache_times = []
    try:
        media_files = get_all_media_file_ids_sync(DB_PATH)
        for entry in media_files:
            if "added" in entry:
                cache_times.append(entry["added"])
    except Exception as e:
        logger.error(f"Error reading media file IDs from SQLite: {str(e)}")
    if cache_times:
        cache_time_diff_seconds = max(cache_times) - min(cache_times)
        cache_time_diff_days = round(cache_time_diff_seconds / 86400, 2)  # rounded to two decimals
    else:
        cache_time_diff_days = 0

    return {
        "cache_files_count": cache_files_count,
        "cache_total_size_mb": cache_total_size_mb,
        "cache_time_diff_days": cache_time_diff_days,
        "channels": channels_stats
    }


def is_local_request(request: Request) -> bool:
    local_hosts = ["127.0.0.1", "localhost"]
    if request and request.client and request.client.host:
        if request.client.host in local_hosts:
            return True
    return False

@app.get("/html/{channel}/{post_id}", response_class=HTMLResponse)
@app.get("/post/html/{channel}/{post_id}", response_class=HTMLResponse)
@app.get("/html/{channel}/{post_id}/{token}", response_class=HTMLResponse)  
@app.get("/post/html/{channel}/{post_id}/{token}", response_class=HTMLResponse)
async def get_post_html(channel: str, post_id: int, request: Request, token: str | None = None, debug: bool = False) -> HTMLResponse:
    if Config["token"] and not is_local_request(request):
        if token != Config["token"]:
            logger.error(f"Invalid token for HTML post: {token}, expected: {Config['token']}")
            raise HTTPException(status_code=403, detail="Invalid token")
        else:
            logger.info(f"Valid token for HTML post: {token}")
    elif Config["token"] and is_local_request(request):
        logger.info(f"Local request, skipping token check for HTML post.")
        
    try:
        parser = PostParser(client.client)
        html_content = await parser.get_post(channel, post_id, 'html', debug)
        if not html_content:
            raise HTTPException(status_code=404, detail="Post not found")
        return HTMLResponse(content=html_content)
    except Exception as e:
        error_message = f"Failed to get HTML post for channel {channel}, post_id {post_id}: {str(e)}"
        logger.error(error_message)
        raise HTTPException(status_code=500, detail=error_message) from e


@app.get("/json/{channel}/{post_id}")
@app.get("/post/json/{channel}/{post_id}")
@app.get("/json/{channel}/{post_id}/{token}")
@app.get("/post/json/{channel}/{post_id}/{token}")
async def get_post(channel: str, post_id: int, request: Request, token: str | None = None, debug: bool = False) -> Response:
    if Config["token"] and not is_local_request(request):
        if token != Config["token"]:
            logger.error(f"Invalid token for JSON post: {token}, expected: {Config['token']}")
            raise HTTPException(status_code=403, detail="Invalid token")
        else:
            logger.info(f"Valid token for JSON post: {token}")
    elif Config["token"] and is_local_request(request):
        logger.info(f"Local request, skipping token check for JSON post.")
            
    try:
        parser = PostParser(client.client)
        json_content = await parser.get_post(channel, post_id, 'json', debug)
        if not json_content:
            raise HTTPException(status_code=404, detail="Post not found")
        return Response(content=json_content, media_type="application/json")
    except Exception as e:
        error_message = f"Failed to get JSON post for channel {channel}, post_id {post_id}: {str(e)}"
        logger.error(error_message)
        raise HTTPException(status_code=500, detail=error_message) from e


@app.get("/raw_json/{channel}/{post_id}")
@app.get("/raw_json/{channel}/{post_id}/{token}")
async def get_raw_post_json(channel: str, post_id: int, request: Request, token: str | None = None) -> Response:
    if Config["token"] and not is_local_request(request):
        if token != Config["token"]:
            logger.error(f"Invalid token for raw JSON post: {token}, expected: {Config['token']}")
            raise HTTPException(status_code=403, detail="Invalid token")
        else:
            logger.info(f"Valid token for raw JSON post: {token}")
    elif Config["token"] and is_local_request(request):
        logger.info(f"Local request, skipping token check for raw JSON post.")
            
    try:
        # Convert numeric channel ID to int if needed
        channel_id: Union[str, int] = channel
        if isinstance(channel, str) and channel.startswith('-100'):
            channel_id = int(channel)
            
        message = await client.client.get_messages(channel_id, post_id)
        if not message:
            raise HTTPException(status_code=404, detail="Post not found")
            
        # Return message as plain text using Pyrogram's built-in string representation
        return Response(content=str(message), media_type="text/plain")
    except Exception as e:
        error_message = f"Failed to get raw JSON post for channel {channel}, post_id {post_id}: {str(e)}"
        logger.error(error_message)
        raise HTTPException(status_code=500, detail=error_message) from e


@app.get("/health")
@app.get("/health/{token}")
async def health_check(request: Request, token: str | None = None) -> Response:
    if Config["token"] and not is_local_request(request):
        if token != Config["token"]:
            logger.error(f"Invalid token for health check: {token}, expected: {Config['token']}")
            raise HTTPException(status_code=403, detail="Invalid token")
        else:
            logger.info(f"Valid token for health check: {token}")
    elif Config["token"] and is_local_request(request):
        logger.info(f"Local request, skipping token check for health check.")
            
    try:
        me = await client.client.get_me()
        
        # Offload heavy filesystem scanning to threadpool
        cache_stats = await asyncio.to_thread(calculate_cache_stats)
        
        config_info = {}    
        for config_key, config_value in Config.items():
            if any(sensitive in config_key.lower() for sensitive in ['token', 'tg_api_id', 'tg_api_hash']):
                config_info[config_key] = mask_sensitive_value(str(config_value)) if config_value else None
            else:
                config_info[config_key] = config_value

        data = {
            "status": "ok",
            "tg_connected": client.client.is_connected,
            "tg_name": me.username,
            "tg_id": me.id,
            "tg_phone": me.phone_number,
            "tg_first_name": me.first_name,
            "tg_last_name": me.last_name,
            "config": config_info,
            **cache_stats
        }
        return Response(content=json.dumps(data), media_type="application/json")    
    except Exception as e:
        error_message = f"Failed to get health check: {str(e)}"
        logger.error(error_message)
        raise HTTPException(status_code=500, detail=error_message) from e

@app.get("/media/{channel}/{post_id}/{file_unique_id}/{digest}", response_model=None)
@app.get("/media/{channel}/{post_id}/{file_unique_id}", response_model=None)
async def get_media(channel: str, post_id: int, file_unique_id: str, request: Request, digest: str | None = None) -> StreamingResponse|Response:
    try:
        url = f"{channel}/{post_id}/{file_unique_id}"
        if not verify_media_digest(url, digest):
            expected_digest = generate_media_digest(url)
            logger.error(f"Invalid digest for media {url}: {digest}, expected: {expected_digest}")
            raise HTTPException(status_code=403, detail="Invalid URL signature")
        #else:
        #    logger.info(f"Valid digest for media {url}: {digest}")   
            
        # Convert numeric channel ID to int if needed
        channel_id: Union[str, int] = channel
        if isinstance(channel, str) and channel.startswith('-100'):
            channel_id = int(channel)
            
        try: # Wrap the download and prepare call
            import time as _time
            _sem_wait_start = _time.monotonic()
            async with HTTP_DOWNLOAD_SEMAPHORE:  # limit concurrent live HTTP downloads
                _sem_wait = _time.monotonic() - _sem_wait_start
                if _sem_wait > 0.5:
                    logger.warning(f"diag_semaphore_wait: {channel}/{post_id}/{file_unique_id} waited {_sem_wait:.3f}s for HTTP_DOWNLOAD_SEMAPHORE")
                _dl_start = _time.monotonic()
                file_path, delete_after = await download_media_file(channel_id, post_id, file_unique_id)
                _dl_elapsed = _time.monotonic() - _dl_start
                logger.info(f"diag_download_timing: {channel}/{post_id}/{file_unique_id} download_media_file took {_dl_elapsed:.3f}s (semaphore_wait={_sem_wait:.3f}s)")
            if not file_path:
                raise HTTPException(status_code=404, detail="File not found")
            if file_path:
                return await prepare_file_response(file_path, request=request, delete_after=delete_after,
                                                   media_key=(str(channel), post_id, file_unique_id))
        except ZeroSizeFileError as e: # Catch zero-size file errors
            logger.warning(f"zero_size_file_encountered: {str(e)}. Instructing client to retry.")
            return Response(
                status_code=503, # Service Unavailable
                content="File processing resulted in zero size, please try again in 10 seconds.",
                headers={"Retry-After": "10"}
            )
            
    except HTTPException:
        raise
    except errors.RPCError as e:
        logger.error(f"Media request RPC error for {channel}/{post_id}/{file_unique_id}: {type(e).__name__} - {str(e)}")
        raise HTTPException(status_code=404, detail="File not found in Telegram") from e
    except ZeroSizeFileError as e: # Catch zero-size file errors FROM DOWNLOAD
        error_message = f"Failed to obtain valid media file for {channel}/{post_id}/{file_unique_id} after download attempt: {str(e)}"
        logger.error(error_message)
        raise HTTPException(status_code=500, detail=error_message) from e
    except Exception as e:
        error_message = f"Failed to get media for {channel}/{post_id}/{file_unique_id}: {str(e)}"
        logger.error(error_message)
        raise HTTPException(status_code=500, detail=error_message) from e

    # If the code reaches here, something is wrong with the logic above.
    # This part should theoretically be unreachable.
    logger.error(f"get_media reached unexpected state for {channel}/{post_id}/{file_unique_id}")
    raise HTTPException(status_code=500, detail="Internal server error: Unexpected state reached in media handling.")

@app.get("/rss/{channel}", response_class=Response)
@app.get("/rss/{channel}/{token}", response_class=Response)
async def get_rss_feed(channel: str, 
                        request: Request,
                        token: str | None = None, 
                        limit: int = 50, 
                        output_type: str = 'rss', 
                        exclude_flags: str | None = None,
                        exclude_text: str | None = None,
                        merge_seconds: int = 5,
                        ) -> Response:
    if Config["token"] and not is_local_request(request):
        if token != Config["token"]:
            logger.error(f"invalid_token_error: token {token}, expected {Config['token']}")
            raise HTTPException(status_code=403, detail="Invalid token")
        else:
            logger.info(f"Valid token for RSS endpoint: {token}")
    elif Config["token"] and is_local_request(request):
        logger.info(f"Local request, skipping token check for RSS endpoint.")
        
    try:
        start_time = time.time()

        if output_type == 'rss':
            rss_content = await generate_channel_rss(channel,
                                                    client=client.client, 
                                                    limit=limit, 
                                                    exclude_flags=exclude_flags,
                                                    exclude_text=exclude_text,
                                                    merge_seconds=merge_seconds)
            elapsed_time = time.time() - start_time
            logger.info(f"rss_generation_timing: channel {channel}, generated in {elapsed_time:.3f} seconds")
            return Response(content=rss_content, media_type="application/xml")
        elif output_type == 'html':
            rss_content = await generate_channel_html(channel,
                                                    client=client.client, 
                                                    limit=limit, 
                                                    exclude_flags=exclude_flags,
                                                    exclude_text=exclude_text,
                                                    merge_seconds=merge_seconds)
            elapsed_time = time.time() - start_time
            logger.info(f"html_generation_timing: channel {channel}, generated in {elapsed_time:.3f} seconds")
            return Response(content=rss_content, media_type="text/html")
        else:
            raise HTTPException(status_code=400, detail=f"invalid_output_type: {output_type}")
    except ValueError as e:
        error_message = f"invalid_parameters_error: {str(e)}"
        logger.error(error_message)
        raise HTTPException(status_code=400, detail=error_message) from e
    except errors.FloodWait as e:
        wait_time = e.value
        random_additional_wait = random.uniform(0, wait_time * 2.5)
        total_wait_time = wait_time + random_additional_wait
        if total_wait_time > 190: total_wait_time = 190

        logger.warning(f"flood_wait_error: channel {channel}, retry after {total_wait_time:.1f} seconds (base: {wait_time}s, random: {random_additional_wait:.1f}s)")
        return Response(
            status_code=429,
            headers={"Retry-After": str(int(total_wait_time))},
            content="Too many requests, please try again later"
        )
    except Exception as e:
        error_message = f"rss_generation_error: channel {channel}, error {str(e)}"
        logger.error(error_message)
        raise HTTPException(status_code=500, detail=error_message) from e

@app.get("/flags", response_model=List[str])
@app.get("/flags/{token}", response_model=List[str])
async def get_available_flags(request: Request, token: str | None = None) -> Response:
    """Returns a list of all possible flags that can be assigned to posts."""
    if Config["token"] and not is_local_request(request):
        if token != Config["token"]:
            logger.error(f"Invalid token for flags endpoint: {token}, expected: {Config['token']}")
            raise HTTPException(status_code=403, detail="Invalid token")
        else:
            logger.info(f"Valid token for flags endpoint: {token}")
    elif Config["token"] and is_local_request(request):
        logger.info(f"Local request, skipping token check for flags endpoint.")

    try:
        flags = PostParser.get_all_possible_flags()
        return Response(content=json.dumps(flags), media_type="application/json")
    except Exception as e:
        error_message = f"Failed to get flags list: {str(e)}"
        logger.error(error_message)
        raise HTTPException(status_code=500, detail=error_message) from e