services/range_response.py

import os
import re
import hashlib
from starlette.responses import Response
from starlette.types import Receive, Scope, Send


class RangeFileResponse(Response):
    def __init__(
        self,
        path: str,
        media_type: str = "video/mp4",
        content_disposition_type: str = "inline",
        filename: str | None = None,
    ):
        self.path = path
        self.file_size = os.path.getsize(path)
        self.media_type = media_type
        self.content_disposition_type = content_disposition_type
        self.filename = filename or os.path.basename(path)
        stat = os.stat(path)
        self.etag = f'"{hashlib.md5(f"{stat.st_mtime}-{stat.st_size}".encode()).hexdigest()}"'
        self.last_modified = stat.st_mtime
        super().__init__()

    async def __call__(self, scope: Scope, receive: Receive, send: Send):
        if scope.get("method") == "HEAD":
            headers = {
                "content-type": self.media_type,
                "content-disposition": f'{self.content_disposition_type}; filename="{self.filename}"',
                "accept-ranges": "bytes",
                "cache-control": "private, max-age=3600",
                "etag": self.etag,
                "content-length": str(self.file_size),
            }
            await send({
                "type": "http.response.start",
                "status": 200,
                "headers": [(k.encode(), v.encode()) for k, v in headers.items()],
            })
            await send({"type": "http.response.body", "body": b"", "more_body": False})
            return

        request_headers = dict(
            (k.decode().lower(), v.decode()) for k, v in scope.get("headers", [])
        )
        range_header = request_headers.get("range", "")

        headers = {
            "content-type": self.media_type,
            "content-disposition": f'{self.content_disposition_type}; filename="{self.filename}"',
            "accept-ranges": "bytes",
            "cache-control": "private, max-age=3600",
            "etag": self.etag,
        }

        if range_header:
            range_match = re.match(r"bytes=(\d*)-(\d*)", range_header)
            if range_match:
                start_str, end_str = range_match.groups()

                if not start_str and not end_str:
                    headers["content-range"] = f"bytes */{self.file_size}"
                    headers["content-length"] = "0"
                    await send({
                        "type": "http.response.start",
                        "status": 416,
                        "headers": [
                            (k.encode(), v.encode()) for k, v in headers.items()
                        ],
                    })
                    await send({"type": "http.response.body", "body": b"", "more_body": False})
                    return

                if start_str and end_str:
                    start = int(start_str)
                    end = int(end_str)
                elif start_str:
                    start = int(start_str)
                    end = self.file_size - 1
                else:
                    suffix_len = int(end_str)
                    start = max(0, self.file_size - suffix_len)
                    end = self.file_size - 1

                if start >= self.file_size:
                    headers["content-range"] = f"bytes */{self.file_size}"
                    headers["content-length"] = "0"
                    await send({
                        "type": "http.response.start",
                        "status": 416,
                        "headers": [
                            (k.encode(), v.encode()) for k, v in headers.items()
                        ],
                    })
                    await send({"type": "http.response.body", "body": b"", "more_body": False})
                    return

                end = min(end, self.file_size - 1)
                content_length = end - start + 1

                headers["content-range"] = f"bytes {start}-{end}/{self.file_size}"
                headers["content-length"] = str(content_length)

                await send({
                    "type": "http.response.start",
                    "status": 206,
                    "headers": [
                        (k.encode(), v.encode()) for k, v in headers.items()
                    ],
                })

                chunk_size = 1024 * 1024  # 1MB chunks for better streaming performance
                with open(self.path, "rb") as f:
                    f.seek(start)
                    remaining = content_length
                    while remaining > 0:
                        read_size = min(chunk_size, remaining)
                        data = f.read(read_size)
                        if not data:
                            break
                        remaining -= len(data)
                        await send({
                            "type": "http.response.body",
                            "body": data,
                            "more_body": remaining > 0,
                        })
                return

        headers["content-length"] = str(self.file_size)

        await send({
            "type": "http.response.start",
            "status": 200,
            "headers": [
                (k.encode(), v.encode()) for k, v in headers.items()
            ],
        })

        chunk_size = 1024 * 1024  # 1MB chunks for better streaming performance
        with open(self.path, "rb") as f:
            while True:
                data = f.read(chunk_size)
                if not data:
                    break
                more = bool(f.tell() < self.file_size)
                await send({
                    "type": "http.response.body",
                    "body": data,
                    "more_body": more,
                })