mfinder/plugins/stream_server.py

import urllib.parse
import re
from fastapi import APIRouter, Request, HTTPException, Header
from fastapi.responses import StreamingResponse, Response
from mfinder.db.files_sql import SESSION, Files

router = APIRouter()

def get_file_by_name(file_name: str) -> Files:
    """Safely retrieves file details from the SQL database without session leaking."""
    try:
        session = SESSION()
        return session.query(Files).filter_by(file_name=file_name).first()
    except Exception as e:
        print(f"Error querying file details for download: {e}")
        return None
    finally:
        try:
            SESSION.close()
        except Exception:
            pass

@router.get("/dl/{file_name}")
async def download_file(request: Request, file_name: str, range: str = Header(None)):
    """
    Asynchronously streams a Telegram file directly to the browser chunk-by-chunk.
    Supports HTTP Range requests (206 Partial Content) for accelerated multi-thread downloads (IDM/ADM/Chrome).
    Uses a dedicated stream_client to keep main bot commands 100% fast and responsive for 50+ concurrent users.
    """
    decoded_file_name = urllib.parse.unquote(file_name)
    
    file_details = get_file_by_name(decoded_file_name)
    if not file_details:
        raise HTTPException(status_code=404, detail="File not found in database.")
    
    file_id = file_details.file_id
    file_size = int(file_details.file_size)
    mime_type = file_details.mime_type or "application/octet-stream"
    
    # Use dedicated stream_client if available to isolate file streaming from bot commands
    stream_client = getattr(request.app.state, "stream_client", None) or getattr(request.app.state, "client", None)
    if not stream_client:
        raise HTTPException(status_code=503, detail="Telegram client gateway not initialized.")
        
    start_bytes = 0
    end_bytes = file_size - 1

    if range:
        range_match = re.match(r"bytes=(\d+)-(\d+)?", range)
        if range_match:
            start_bytes = int(range_match.group(1))
            if range_match.group(2):
                end_bytes = int(range_match.group(2))

    if start_bytes >= file_size or end_bytes >= file_size or start_bytes > end_bytes:
        headers = {"Content-Range": f"bytes */{file_size}"}
        return Response(status_code=416, headers=headers)

    content_length = (end_bytes - start_bytes) + 1

    chunk_offset = start_bytes // (1024 * 1024)
    skip_bytes = start_bytes % (1024 * 1024)

    async def media_generator():
        import asyncio
        queue = asyncio.Queue(maxsize=3)
        
        async def producer():
            try:
                async for chunk in stream_client.stream_media(file_id, offset=chunk_offset):
                    await queue.put(chunk)
                await queue.put(None)
            except Exception as e:
                print(f"Error in producer: {e}")
                await queue.put(None)
                
        producer_task = asyncio.create_task(producer())
        bytes_sent = 0
        try:
            first = True
            while True:
                chunk = await queue.get()
                if chunk is None:
                    break
                    
                if first:
                    first = False
                    if skip_bytes > 0:
                        chunk = chunk[skip_bytes:]
                        
                if bytes_sent + len(chunk) > content_length:
                    chunk = chunk[:content_length - bytes_sent]
                    
                if not chunk:
                    break
                    
                yield chunk
                bytes_sent += len(chunk)
                
                if bytes_sent >= content_length:
                    break
        except Exception as e:
            print(f"Error occurred while streaming file '{decoded_file_name}': {e}")
        finally:
            producer_task.cancel()
            try:
                await producer_task
            except asyncio.CancelledError:
                pass
            except Exception as ex:
                print(f"Error clean cancelling producer: {ex}")
            
    headers = {
        "Content-Disposition": f"attachment; filename=\"{decoded_file_name}\"; filename*=UTF-8''{urllib.parse.quote(decoded_file_name)}",
        "Content-Length": str(content_length),
        "Accept-Ranges": "bytes"
    }

    if range:
        headers["Content-Range"] = f"bytes {start_bytes}-{end_bytes}/{file_size}"
        return StreamingResponse(media_generator(), status_code=206, media_type=mime_type, headers=headers)
    else:
        return StreamingResponse(media_generator(), status_code=200, media_type=mime_type, headers=headers)