unshackle-services/PBS/__init__.py

from __future__ import annotations
import json
import re
import uuid
from collections.abc import Generator
from http.cookiejar import CookieJar
from typing import Any, Optional, Union
import click
from langcodes import Language
from unshackle.core.credential import Credential
from unshackle.core.manifests import HLS
from unshackle.core.search_result import SearchResult
from unshackle.core.service import Service
from unshackle.core.titles import Episode, Movie, Movies, Series, Title_T, Titles_T
from unshackle.core.tracks import Chapters, Tracks


class PBS(Service):
    """
    Service code for PBS (https://www.pbs.org)

    www.nostalgic.cc
    Authorization: Cookies
    Geofence: US
    """

    TITLE_RE = r"^(?:https?://(?:www\.)?pbs\.org/(?P<type>video|show)/)?(?P<id>[a-zA-Z0-9-]+)"
    GEOFENCE = ("US",)

    @staticmethod
    @click.command(name="PBS", short_help="https://www.pbs.org", help=__doc__)
    @click.argument("title", type=str)
    @click.option(
        "-a", "--all", "all_",
        is_flag=True,
        default=False,
        help="Parse input as slug to download all content.",
    )
    @click.pass_context
    def cli(ctx, **kwargs):
        return PBS(ctx, **kwargs)

    def __init__(self, ctx, title: str, all_: bool = False):
        super().__init__(ctx)
        self.title = title
        self.is_show = all_

        m = re.match(self.TITLE_RE, title)
        if not m:
            raise ValueError(f"Could not parse PBS URL or ID: {title!r}.")

        url_type = m.group("type")
        self.title_id = m.group("id")

        if url_type == "show":
            self.is_show = True

        self.uid: str = str(uuid.uuid4())
        self.passport: str = "no"
        self.callsign: Optional[str] = None
        self.station_id: Optional[str] = None

    def authenticate(
        self,
        cookies: Optional[CookieJar] = None,
        credential: Optional[Credential] = None,
    ) -> None:
        super().authenticate(cookies, credential)
        if cookies:
            uid = next((c.value for c in cookies if c.name == "pbs_uid"), None)
            if uid:
                self.uid = uid
                self.passport = "yes"
            self.callsign = next((c.value for c in cookies if c.name == "pbsol.station"), None)
            self.station_id = next((c.value for c in cookies if c.name == "pbsol.station_id"), None)

        self.session.headers.update({"User-Agent": self.config["client"]["web"]["user_agent"]})

    def search(self) -> Generator[SearchResult, None, None]:
        return
        yield

    def get_titles(self) -> Titles_T:
        if self.is_show:
            return self._get_show_titles(self.title_id)
        return self._get_video_title(self.title_id)

    def get_tracks(self, title: Title_T) -> Tracks:
        video_bridge = title.data.get("video_bridge") or self._fetch_video_bridge(title.id)

        if video_bridge.get("availability") != "available":
            raise ValueError(
                f"Video '{title.id}' is unavailable. (Status: {video_bridge.get('availability')!r}). "
                "Passport content requires cookies from a logged-in PBS account."
            )

        encodings = video_bridge.get("encodings", [])
        if not encodings:
            raise ValueError(f"No streams found for '{title.id}'.")

        m3u8_url = self._resolve_encoding(encodings[0])
        self.log.debug(f"HLS master: {m3u8_url}")
        tracks = HLS.from_url(url=m3u8_url, session=self.session).to_tracks(language=title.language)
        kept_subs = []
        for sub in tracks.subtitles:
            if self._subtitle_has_cues(sub):
                kept_subs.append(sub)
            else:
                self.log.warning(
                    f" - Dropping empty subtitle track ({sub.language}{' SDH' if sub.sdh else ''})."
                )
        tracks.subtitles = kept_subs

        return tracks

    def _subtitle_has_cues(self, sub: Any) -> bool:
        from urllib.parse import urljoin

        url = sub.url[0] if isinstance(sub.url, list) and sub.url else sub.url
        if not isinstance(url, str):
            return True
        try:
            resp = self.session.get(url, timeout=15)
            if not resp.ok:
                return True
            text = resp.text
        except Exception:
            return True

        if "-->" in text:
            return True
        if "#EXTM3U" not in text:
            return False

        checked = 0
        for line in text.splitlines():
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            try:
                seg = self.session.get(urljoin(url, line), timeout=15)
            except Exception:
                return True
            if seg.ok and "-->" in seg.text:
                return True
            checked += 1
            if checked >= 40:
                break
        return False

    def get_chapters(self, title: Title_T) -> Chapters:
        return Chapters()

    def get_widevine_service_certificate(self, **_: Any) -> None:
        return None

    def get_widevine_license(self, *, challenge: bytes, title: Title_T, track: Any) -> None:
        return None

    def _get_video_title(self, video_slug: str) -> Titles_T:
        video_bridge = self._fetch_video_bridge(video_slug)
        show_slug = video_bridge.get("program", {}).get("slug")
        show_title = video_bridge.get("program", {}).get("title", "Unknown Show")

        episode_data = self._find_episode_in_show(show_slug, video_slug) if show_slug else None

        if episode_data:
            parent = episode_data.get("parent", {})
            season_data = parent.get("season", {})
            show_display = season_data.get("show", {}).get("title", show_title)
            season_num = int(season_data.get("ordinal") or 0)
            ep_num = int(parent.get("ordinal") or 0)
            year = self._parse_year(episode_data.get("premiere_date"))

            return Series([
                Episode(
                    id_=video_slug,
                    service=self.__class__,
                    title=show_display,
                    season=season_num,
                    number=ep_num,
                    name=episode_data.get("title", video_slug),
                    description=episode_data.get("description_short"),
                    year=year,
                    language=Language.get("en"),
                    data={"video_bridge": video_bridge, "episode": episode_data},
                )
            ])

        self.log.warning("Unable to find episode metadata.")
        fallback_name = video_bridge.get("_page_title") or video_slug.replace("-", " ").title()
        return Series([
            Episode(
                id_=video_slug,
                service=self.__class__,
                title=show_title,
                season=0,
                number=0,
                name=fallback_name,
                language=Language.get("en"),
                data={"video_bridge": video_bridge},
            )
        ])

    def _get_show_titles(self, show_slug: str) -> Titles_T:
        self.log.info(f"Fetching season data for {show_slug}.")
        seasons = self._fetch_show_seasons(show_slug)
        specials = self._fetch_show_specials(show_slug)
        show_display = show_slug.replace("-", " ").title()
        episodes = []

        if seasons:
            self.log.info(f"Found {len(seasons)} season(s)")
            for season_cid, season_ordinal in seasons:
                self.log.info(f"Season {season_ordinal}")
                for ep in self._fetch_season_episodes(show_slug, season_cid):
                    parent = ep.get("parent", {})
                    season_data = parent.get("season", {})
                    show_display = season_data.get("show", {}).get("title") or show_display
                    season_num = int(season_data.get("ordinal") or season_ordinal)
                    ep_num = int(parent.get("ordinal") or 0)
                    year = self._parse_year(ep.get("premiere_date"))

                    episodes.append(Episode(
                        id_=ep["slug"],
                        service=self.__class__,
                        title=show_display,
                        season=season_num,
                        number=ep_num,
                        name=ep.get("title", ep["slug"]),
                        description=ep.get("description_short"),
                        year=year,
                        language=Language.get("en"),
                        data={"episode": ep},
                    ))

        if specials:
            self.log.info(f"Found {len(specials)} special(s)")
            specials.sort(key=lambda x: x.get("premiere_date") or "")
            for i, sp in enumerate(specials, start=1):
                self.log.info(f"Special {i}")
                parent = sp.get("parent", {})
                show_display = parent.get("show", {}).get("title") or show_display
                year = self._parse_year(sp.get("premiere_date"))

                episodes.append(Episode(
                    id_=sp["slug"],
                    service=self.__class__,
                    title=show_display,
                    season=0,
                    number=i,
                    name=sp.get("title", sp["slug"]),
                    description=sp.get("description_short"),
                    year=year,
                    language=Language.get("en"),
                    data={"episode": sp},
                ))

        if not episodes:
            raise ValueError(f"No episodes or specials found for {show_slug!r}.")

        return Series(episodes)

    def _fetch_video_bridge(self, video_slug: str) -> dict:
        params: dict[str, str] = {
            "uid": self.uid,
            "userPassportStatus": self.passport,
            "autoplay": "true",
            "unsafeDisableUpsellHref": "true",
        }
        if self.callsign:
            params["callsign"] = self.callsign
        if self.station_id:
            params["station_id"] = self.station_id

        r = self.session.get(
            self.config["endpoints"]["portalplayer"] + video_slug + "/",
            params=params,
        )
        r.raise_for_status()

        idx = r.text.find("window.videoBridge = ")
        if idx == -1:
            raise ValueError(
                f"videoBridge not found in portalplayer response for {video_slug!r}.\n"
                "Double check your input URL and make sure it is correct."
            )

        json_start = r.text.index("{", idx)
        video_bridge, _ = json.JSONDecoder().raw_decode(r.text, json_start)
        title_m = re.search(r"<title>\s*Video:\s*(.*?)\s*\|\s*Watch", r.text, re.DOTALL)
        if title_m:
            video_bridge["_page_title"] = " ".join(title_m.group(1).split())

        return video_bridge

    def _resolve_encoding(self, encoding_url: str) -> str:
        r = self.session.get(encoding_url, params={"format": "jsonp", "callback": "__jp0"})
        r.raise_for_status()

        m = re.search(r"__jp0\((.+)\)\s*$", r.text.strip())
        if not m:
            raise ValueError(f"Unable to parse URS JSONP response: {r.text[:200]!r}")

        data = json.loads(m.group(1))
        url = data.get("url")
        if not url:
            raise ValueError(f"URS redirect returned no URL: {data}")
        return url

    def _fetch_show_seasons(self, show_slug: str) -> list[tuple[str, int]]:
        r = self.session.get(f"https://www.pbs.org/show/{show_slug}/", timeout=15)
        if not r.ok:
            self.log.warning(f"Page returned {r.status_code} for '{show_slug}'.")
            return []
        seasons = self._parse_seasons_from_html(r.text)
        if not seasons:
            self.log.warning(f"Unable to find season data for '{show_slug}': {r.status_code}.")
        return seasons

    def _parse_seasons_from_html(self, html: str) -> list[tuple[str, int]]:
        UUID_PAT = r'[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}'
        seasons: list[tuple[str, int]] = []
        seen: set[str] = set()

        for url_m in re.finditer(
            r'content\.services\.pbs\.org/v3/pbsorg/screens/shows/[^/]+/seasons/(' + UUID_PAT + r')/',
            html,
        ):
            cid = url_m.group(1)
            if cid in seen:
                continue
            seen.add(cid)

            window = html[max(0, url_m.start() - 800):url_m.start()]
            all_ords = re.findall(r'ordinal[\\\"]*\s*:\s*(\d+)', window)
            ordinal = int(all_ords[-1]) if all_ords else (len(seasons) + 1)
            seasons.append((cid, ordinal))

        return sorted(seasons, key=lambda x: x[1], reverse=True)

    def _fetch_show_specials(self, show_slug: str) -> list[dict]:
        url = self.config["endpoints"]["show_specials"].format(show_slug=show_slug)
        r = self.session.get(url)
        if not r.ok:
            self.log.warning(f"Unable to find specials for '{show_slug}': {r.status_code}.")
            return []
        return [
            sp for sp in r.json()
            if sp.get("slug") != sp.get("parent", {}).get("slug")
        ]

    def _fetch_season_episodes(self, show_slug: str, season_cid: str) -> list[dict]:
        url = self.config["endpoints"]["show_episodes"].format(
            show_slug=show_slug,
            season_cid=season_cid,
        )
        r = self.session.get(url)
        if not r.ok:
            self.log.warning(f"Unable to find episodes for season {season_cid}: {r.status_code}.")
            return []
        return r.json()

    def _find_episode_in_show(self, show_slug: str, video_slug: str) -> Optional[dict]:
        r = self.session.get(f"https://www.pbs.org/video/{video_slug}/", timeout=15)
        seasons = self._parse_seasons_from_html(r.text) if r.ok else []
        if not seasons:
            seasons = self._fetch_show_seasons(show_slug)
        for season_cid, _ in seasons:
            for ep in self._fetch_season_episodes(show_slug, season_cid):
                if ep.get("slug") == video_slug:
                    return ep
        for sp in self._fetch_show_specials(show_slug):
            if sp.get("slug") == video_slug:
                return sp
        return None

    @staticmethod
    def _parse_year(date_str: Optional[str]) -> Optional[int]:
        if not date_str:
            return None
        try:
            return int(date_str[:4])
        except (ValueError, TypeError):
            return None