from __future__ import annotations
import json
import re
from collections.abc import Generator
from http.cookiejar import CookieJar
from typing import Optional, Union
import click
from langcodes import Language
from unshackle.core.constants import AnyTrack
from unshackle.core.credential import Credential
from unshackle.core.manifests import DASH, 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 Chapter, Tracks


class PBSK(Service):
    """
    Service code for the PBS Kids streaming service (https://pbskids.org)
    
    Author: n0stal6ic
    Authorization: None
    Security: FHD@L3
    """

    TITLE_RE = r"^(?:https?://(?:www\.)?pbskids\.org)?(?P<path>/videos/watch/[^?#\s]+)"
    GEOFENCE = ("US",)
    _HLS_PROFILES = ("hls-16x9-1080p", "hls-16x9-720p")
    _DASH_DRM_PROFILES = ("dash-16x9-1080p-drm", "dash-16x9-720p-drm")

    @staticmethod
    @click.command(name="PBSK", short_help="https://pbskids.org")
    @click.argument("title", type=str)
    @click.pass_context
    def cli(ctx, **kwargs):
        return PBSK(ctx, **kwargs)

    def __init__(self, ctx, title: str):
        super().__init__(ctx)
        self.title = title
        self._license_url: Optional[str] = None

    def authenticate(self, cookies: Optional[CookieJar] = None, credential: Optional[Credential] = None) -> None:
        super().authenticate(cookies, credential)

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

    def get_titles(self) -> Titles_T:
        match = re.match(self.TITLE_RE, self.title.strip())
        if not match:
            raise ValueError(f"Unrecognized URL: {self.title!r}")

        path = match.group("path").rstrip("/")
        page_url = f"https://pbskids.org{path}"

        self.log.info("Detecting buildId")
        resp = self.session.get(page_url)
        resp.raise_for_status()

        nd_match = re.search(
            r'<script[^>]+id="__NEXT_DATA__"[^>]*>(.*?)</script>',
            resp.text,
            re.DOTALL,
        )
        if not nd_match:
            raise RuntimeError("__NEXT_DATA__ not found.")

        next_data = json.loads(nd_match.group(1))
        build_id = next_data.get("buildId")
        if not build_id:
            raise RuntimeError("buildId not found.")
        self.log.debug(f"BuildID: {build_id}")

        api_url = f"https://pbskids.org/_next/data/{build_id}/en-US{path}.json"
        self.log.info(f"Fetching JSON from: {api_url}")
        resp = self.session.get(api_url)
        resp.raise_for_status()
        data = resp.json()

        video_data = data["pageProps"]["videoData"]
        asset = video_data["mediaManagerAsset"]

        video_id = str(video_data["id"])
        video_type = video_data.get("videoType", "short")
        title_text = video_data.get("title") or asset.get("title", "Unknown")
        description = asset.get("description_short") or asset.get("description_long", "")
        season_number = asset.get("season_number")
        episode_number = asset.get("episode_number")

        show_title: Optional[str] = None
        for prop in video_data.get("properties", []):
            show_title = prop.get("title")
            if show_title:
                break

        title_data = {
            "asset": asset,
            "drm_enabled": bool(asset.get("drm_enabled", False)),
        }

        if (
            video_type == "fullEpisode"
            and show_title
            and season_number is not None
            and episode_number is not None
        ):
            return Series(
                [
                    Episode(
                        id_=video_id,
                        service=self.__class__,
                        title=show_title,
                        season=season_number,
                        number=episode_number,
                        name=title_text,
                        description=description,
                        year=None,
                        language=Language.get("en"),
                        data=title_data,
                    )
                ]
            )
        else:
            return Movies(
                [
                    Movie(
                        id_=video_id,
                        service=self.__class__,
                        name=title_text,
                        description=description,
                        year=None,
                        language=Language.get("en"),
                        data=title_data,
                    )
                ]
            )

    def get_tracks(self, title: Title_T) -> Tracks:
        asset = title.data["asset"]
        drm_enabled = title.data["drm_enabled"]
        videos = asset.get("videos", [])

        if drm_enabled:
            return self._get_drm_tracks(title, videos)
        else:
            return self._get_clear_tracks(title, videos)

    def _pick_stream(self, videos: list[dict], preferred_profiles: tuple[str, ...]) -> Optional[dict]:
        seen_urls: set[str] = set()
        for profile in preferred_profiles:
            for video in videos:
                if video.get("profile") == profile:
                    url = video.get("url", "")
                    if url and url not in seen_urls:
                        seen_urls.add(url)
                        return video
        return None

    def _resolve_redirect(self, url: str) -> str:
        resp = self.session.head(url, allow_redirects=True)
        return str(resp.url)

    def _get_clear_tracks(self, title: Title_T, videos: list[dict]) -> Tracks:
        video = self._pick_stream(videos, self._HLS_PROFILES)
        if not video:
            raise RuntimeError("No HLS stream found.")

        stream_url = self._resolve_redirect(video["url"])
        self.log.debug(f"HLS stream URL: {stream_url}")

        return HLS.from_url(url=stream_url, session=self.session).to_tracks(language=title.language)

    def _get_drm_tracks(self, title: Title_T, videos: list[dict]) -> Tracks:
        video = self._pick_stream(videos, self._DASH_DRM_PROFILES)
        if not video:
            raise RuntimeError("No Widevine DASH stream found.")

        self._license_url = video.get("widevine_license")
        if not self._license_url:
            raise RuntimeError("No Widevine license URL detected from DASH stream.")

        stream_url = self._resolve_redirect(video["url"])
        self.log.debug(f"DASH stream URL: {stream_url}")
        self.log.debug(f"Widevine license URL: {self._license_url}")

        return DASH.from_url(url=stream_url, session=self.session).to_tracks(language=title.language)

    def get_chapters(self, title: Title_T) -> list[Chapter]:
        return []

    def get_widevine_service_certificate(self, **_: any) -> Optional[str]:
        if self.config:
            return self.config.get("certificate")
        return None

    def get_widevine_license(
        self, *, challenge: bytes, title: Title_T, track: AnyTrack
    ) -> Optional[Union[bytes, str]]:
        if not self._license_url:
            raise RuntimeError("Widevine license URL is not set.")

        resp = self.session.post(
            url=self._license_url,
            data=challenge,
            headers={"Content-Type": "application/octet-stream"},
        )
        resp.raise_for_status()
        return resp.content