live/backend/services/stream_service.py

from sqlalchemy import select, update, delete
from sqlalchemy.ext.asyncio import AsyncSession
from models.stream import Stream, StreamStatus
from services.media_service import media_service
from services.load_balancer import lb
from services.multicast_detector import multicast_detector
from services.transcoder import transcode_mgr
from config import settings
import json
import re
from datetime import datetime


# MediaMTX rejects any HLS path that is not alphanumeric ("invalid path name: can
# contain only alphanumeric characters"), and the label is also used as a directory
# name, so every rendition needs a path-safe slug of its own.
def hls_path_slug(label: str) -> str:
    """Filesystem- and MediaMTX-safe rendition path for a display label."""
    slug = re.sub(r"[^A-Za-z0-9]", "", label or "")
    return slug or "rendition"


def with_hls_path(rendition: dict) -> dict:
    """Attach the `path` an admin-configured rendition is stored and served under."""
    return {**rendition, "path": hls_path_slug(rendition.get("label", ""))}


class StreamService:
    def __init__(self, db: AsyncSession):
        self.db = db

    @staticmethod
    def effective_renditions(stream: Stream) -> list:
        if stream.renditions_json:
            try:
                renditions = json.loads(stream.renditions_json)
                if isinstance(renditions, list) and renditions:
                    return renditions
            except (TypeError, ValueError, json.JSONDecodeError):
                pass
        return settings.transcode_profiles.get(stream.transcode_profile, settings.qualities)

    @staticmethod
    def source_label() -> str:
        """Display name of the stream-copy rendition, or '' when it is disabled.

        A blank label would leak into the HLS path and produce a URL like
        /hls/<key>//index.m3u8, so it degrades to "off" instead.
        """
        return settings.source_rendition_label.strip() if settings.source_rendition_enabled else ""

    @classmethod
    def source_slug(cls) -> str:
        """Path segment for that rendition, safe for MediaMTX and the filesystem."""
        if not cls.source_label():
            return ""
        return hls_path_slug(settings.source_rendition_slug)

    def watch_renditions(self, stream: Stream) -> list:
        """Renditions a viewer may pick: the stream-copy rendition first, then the configured ones.

        Each entry carries both the label shown to viewers and the `path` used in
        HLS URLs — MediaMTX only accepts alphanumeric HLS paths, so a Chinese label
        such as "原画" cannot double as a directory name.

        The source entry carries no size: the input resolution is only known to
        ffmpeg, so the player measures it from the media it ends up decoding.
        """
        renditions = [with_hls_path(r) for r in self.effective_renditions(stream)]
        label, slug = self.source_label(), self.source_slug()
        if not label or not slug or not renditions or renditions[0]["path"] == slug:
            return renditions
        return [{"label": label, "path": slug, "width": None, "height": None, "copy": True}] + renditions

    def transcode_renditions(self, stream: Stream) -> list:
        """Renditions to hand to the transcoder, including the stream-copy source rendition.

        Carries the same `path` values as watch_renditions so the files ffmpeg
        writes are the files the player asks for.

        The copy entry is the top rendition with its label replaced, so it carries
        every key the transcoder reads. ffmpeg ignores those numbers for a stream
        copy, but the supervisor in main.py compares TranscodeManager's recorded
        signature against this list every tick — an entry missing a field would
        both crash the start and make the signature drift.
        """
        renditions = [with_hls_path(r) for r in self.effective_renditions(stream)]
        label, slug = self.source_label(), self.source_slug()
        if not label or not slug or not renditions or renditions[0]["path"] == slug:
            return renditions
        return [{**renditions[0], "label": label, "path": slug, "copy": True}] + renditions

    async def create_stream(self, title: str, description: str | None = None,
                            cover_url: str | None = None,
                            scheduled_at: str | None = None,
                            replay_file: str | None = None,
                            replay_url: str | None = None,
                            chapters: list[dict] | None = None,
                            transcode_profile: str = "standard",
                            renditions_json: str | None = None,
                            homepage_visible: bool = True,
                            remote_control_enabled: bool = False) -> Stream:
        import uuid
        scheduled_datetime = None
        if scheduled_at:
            try:
                scheduled_datetime = datetime.fromisoformat(scheduled_at.replace("Z", "+00:00"))
            except ValueError as exc:
                raise ValueError("预定开播时间格式无效") from exc
        stream = Stream(
            title=title,
            description=description,
            cover_url=cover_url,
            replay_file=replay_file,
            replay_url=replay_url,
            chapters_json=json.dumps(chapters or [], ensure_ascii=False),
            transcode_profile=transcode_profile,
            renditions_json=renditions_json,
            homepage_visible=homepage_visible,
            remote_control_enabled=remote_control_enabled,
            stream_key=uuid.uuid4().hex,
            status=StreamStatus.pending,
            scheduled_at=scheduled_datetime,
        )
        self.db.add(stream)
        await self.db.commit()
        await self.db.refresh(stream)
        return stream

    async def get_all_streams(self):
        result = await self.db.execute(
            select(Stream).order_by(Stream.created_at.desc())
        )
        return result.scalars().all()

    async def get_public_streams(self):
        result = await self.db.execute(
            select(Stream)
            .where(Stream.homepage_visible.is_(True))
            .order_by(Stream.created_at.desc())
        )
        return result.scalars().all()

    async def get_stream(self, stream_id: int):
        result = await self.db.execute(
            select(Stream).where(Stream.id == stream_id)
        )
        return result.scalar_one_or_none()

    async def get_stream_by_key(self, stream_key: str):
        result = await self.db.execute(
            select(Stream).where(Stream.stream_key == stream_key)
        )
        return result.scalar_one_or_none()

    async def update_stream(self, stream_id: int, **kwargs):
        stmt = update(Stream).where(Stream.id == stream_id).values(**kwargs)
        await self.db.execute(stmt)
        await self.db.commit()

    async def delete_stream(self, stream_id: int):
        stream = await self.get_stream(stream_id)
        if stream:
            transcode_mgr.stop_transcode(stream.stream_key)
            await self.db.delete(stream)
            await self.db.commit()

    async def set_stream_live(self, stream_key: str):
        stream = await self.get_stream_by_key(stream_key)
        if stream:
            from datetime import datetime
            await self.update_stream(
                stream.id,
                status=StreamStatus.live,
                started_at=datetime.utcnow(),
            )
            rtmp_url = media_service.get_rtmp_url(stream_key)
            transcode_mgr.start_transcode(stream_key, rtmp_url, self.transcode_renditions(stream))

    async def set_stream_ended(self, stream_key: str):
        stream = await self.get_stream_by_key(stream_key)
        if stream:
            from datetime import datetime
            await self.update_stream(
                stream.id,
                status=StreamStatus.ended,
                ended_at=datetime.utcnow(),
            )
            transcode_mgr.stop_transcode(stream_key)

    async def get_watch_urls(self, stream_key: str, client_ip: str | None = None):
        stream = await self.get_stream_by_key(stream_key)
        if not stream:
            return {"lines": [], "qualities": [], "multicast_available": False, "selected_node": None}
        qualities = self.watch_renditions(stream)
        best_node = lb.select_best_node(client_ip)
        if not best_node:
            best_node = lb.get_enabled_nodes()[0] if lb.get_enabled_nodes() else None

        lines = []
        if best_node:
            base_hls = f"/hls/{stream_key}"
            line_info = {
                "node_id": best_node["id"],
                "node_name": best_node["name"],
                "weight": best_node["weight"],
                "type": "hls",
                "url": f"{base_hls}/index.m3u8",
                "qualities": [],
            }
            for quality in qualities:
                # The URL segment is the rendition's path slug, not its label:
                # labels are display text and may be Chinese.
                line_info["qualities"].append({
                    "label": quality["label"],
                    "width": quality["width"],
                    "height": quality["height"],
                    "url": f"{base_hls}/{quality['path']}/index.m3u8",
                })
            lines.append(line_info)

        return {
            "lines": lines,
            "qualities": [{"label": q["label"], "width": q["width"], "height": q["height"]} for q in qualities],
            "multicast_available": multicast_detector.is_available(),
            "selected_node": {
                "id": best_node["id"],
                "name": best_node["name"],
            } if best_node else None,
        }