"""yt-dlp extractor for berriz.in (Amazon IVS-backed K-pop fan community lives)."""

import itertools
import json
import re

from yt_dlp.extractor.common import InfoExtractor
from yt_dlp.utils import (
    ExtractorError,
    int_or_none,
    str_or_none,
    traverse_obj,
    unified_timestamp,
)

# Shared with the gallery-dl module, which handles the same URLs: /en, /ko,
# /zh-Hans, ... are optional, as are the www. and link.berriz.in/web/main
# prefixes. Safe to inline into an (?x) pattern - no literal whitespace, and
# a '#' inside a character class is not a verbose-mode comment.
_BASE = (r'https?://(?:www\.|link\.)?berriz\.in/'
         r'(?:[a-z]{2}(?:-[A-Za-z]{2,4})?/)?(?:web/main/)?')
_UUID = r'[0-9a-fA-F-]{36}'

# berriz's i18n layer intermittently ignores languageCode and answers with the
# literal string 'Unknown <locale>' - 'Unknown ja', 'Unknown zh-Hans' - in place
# of a community name. It is a server-side miss, not a signal about the request:
# the same URL fetched twice a second apart returns 'ONEUS' and then the
# placeholder. Retrying would usually clear it but costs a round trip on a field
# nothing depends on, so treat the sentinel as absent and let callers fall back
# to the community slug, which the URL always carries.
_PLACEHOLDER_NAME = re.compile(r'(?i)\Aunknown\s+[a-z]{2}(?:-[A-Za-z]{2,4})?\Z')


def _community_name(name):
    """Return a usable community name, or None for the i18n placeholder."""
    name = str_or_none(name)
    return None if name and _PLACEHOLDER_NAME.match(name.strip()) else name


# berriz's liveStatus vocabulary. END is the window between a broadcast stopping
# and its replay finishing transcoding, and nothing is downloadable during it:
# the live playback_info answers 0000 with a null playbackUrl, the replay
# endpoint 500s, and live.replay carries no duration. The status becomes REPLAY
# once the recording is published. yt-dlp's post_live means exactly this - was
# live, VOD not available yet - so END maps onto it rather than onto was_live.
_LIVE_STATUS = {
    'ON_AIR': 'is_live',
    'END': 'post_live',
    'REPLAY': 'was_live',
}


def _live_status(raw):
    """Map berriz's liveStatus onto yt-dlp's, or None when it is unknown."""
    if not raw:
        return None
    return _LIVE_STATUS.get(raw, 'was_live')


class _BerrizAuthMixin:
    """Shared endpoints, headers and access-token handling.

    berriz answers an expired or invalid bz_a with a plain-text HTTP 401 on
    every endpoint, including the ones that work fine anonymously, so the
    failure has to be caught at the HTTP layer - there is no JSON error code
    to match on. Access tokens last about an hour, which makes this the
    normal state of an exported cookies file rather than an edge case.
    """

    _API = 'https://svc-api.berriz.in'
    _ACCOUNT_API = 'https://account.berriz.in'
    # Public SPA client id, lifted from the site's JS bundle.
    _CLIENT_ID = 'e8faf56c-575a-42d2-933d-7b2e279ad827'
    _HEADERS = {
        'Accept': 'application/json',
        'Origin': 'https://berriz.in',
        'Referer': 'https://berriz.in/',
    }

    def _logged_in(self):
        cookies = self._get_cookies('https://berriz.in/')
        return 'bz_a' in cookies or 'bz_r' in cookies

    # Only failures are remembered: a dead bz_r will not come back to life, but
    # a successful refresh may need repeating on a playlist walk that outlives
    # the one-hour access token.
    _refresh_failed = False

    def _refresh_token(self, item_id):
        """Trade the bz_r refresh cookie for a fresh bz_a. Returns True on success."""
        if self._refresh_failed:
            return False
        if 'bz_r' not in self._get_cookies('https://berriz.in/'):
            self._refresh_failed = True
            return False
        try:
            response = self._download_json(
                f'{self._ACCOUNT_API}/auth/v1/token:refresh', item_id,
                note='Refreshing access token', errnote='Token refresh failed',
                data=json.dumps({'clientId': self._CLIENT_ID}).encode(),
                headers={**self._HEADERS, 'Content-Type': 'application/json'},
                fatal=False)
        except ExtractorError:
            self._refresh_failed = True
            return False
        # yt-dlp's cookiejar absorbs the Set-Cookie response itself, so a 0000
        # here means the jar already holds the new bz_a. Only bz_a is rotated;
        # bz_r comes back unchanged, so refreshing never costs the session.
        if traverse_obj(response, 'code') == '0000':
            return True
        self._refresh_failed = True
        return False

    def _download_api(self, path, item_id, note, query=None, fatal=True,
                      expected_status=None):
        """Fetch and parse an API response, refreshing an expired bz_a once.

        The refresh is scoped to this one request rather than the extractor as
        a whole, so a token that expires midway through a long playlist run is
        still recovered.
        """
        allowed = expected_status or ()
        if isinstance(allowed, int):
            allowed = (allowed,)

        for attempt in (1, 2):
            page, urlh = self._download_webpage_handle(
                f'{self._API}{path}', item_id,
                note=note if attempt == 1 else f'{note} (after token refresh)',
                headers=self._HEADERS,
                query={'languageCode': 'en', **(query or {})},
                fatal=False, expected_status=(401, *allowed)) or (None, None)

            if urlh is not None and urlh.status == 401:
                if attempt == 1 and self._refresh_token(item_id):
                    continue
                self.raise_login_required(
                    'berriz rejected the access token, which expires after about '
                    'an hour. Re-export your cookies, or use '
                    '--cookies-from-browser BROWSER to keep them fresh',
                    method='cookies')
            break

        if page is None:
            if fatal:
                raise ExtractorError(
                    'Unable to download berriz API response', video_id=item_id)
            return None
        return self._parse_json(page, item_id, fatal=fatal)

    def _call_api(self, path, item_id, note, query=None, fatal=True):
        """Return the `data` payload. berriz reports its own errors as HTTP 200 + code."""
        response = self._download_api(
            path, item_id, note, query=query, fatal=fatal) or {}
        code = response.get('code')
        if code != '0000':
            if fatal:
                raise ExtractorError(
                    response.get('message') or f'berriz returned {code}',
                    expected=True, video_id=item_id)
            return None
        return response.get('data')

    def _community_name_for(self, slug):
        """Look up a community's display name by slug.

        Unlike the per-media endpoint, this one honours languageCode reliably,
        so it is the repair path when the media payload hands back a
        placeholder. Only called when that happens, so the common case still
        costs a single request.
        """
        return _community_name(traverse_obj(self._call_api(
            f'/service/v1/community/id/{slug}', slug,
            'Resolving community name', fatal=False), 'communityName'))


class BerrizIE(_BerrizAuthMixin, InfoExtractor):
    IE_NAME = 'berriz'
    IE_DESC = 'berriz.in live streams, replays and VODs'
    _VALID_URL = rf'''(?x)
        {_BASE}
        (?P<community>[^/?#]+)/
        (?P<kind>live|media|vod|archive)/
        (?:(?P<replay>replay)/)?
        (?P<id>{_UUID})'''

    _TESTS = [{
        'url': 'https://berriz.in/en/kiiikiii/live/019fef22-1b10-2950-8dc0-6e9b20226ce4/',
        'info_dict': {
            'id': '019fef22-1b10-2950-8dc0-6e9b20226ce4',
            'ext': 'mp4',
            'title': str,
            'live_status': 'is_live',
            'channel': 'KiiiKiii',
        },
        'skip': 'Requires a logged-in berriz account; live streams are ephemeral',
    }, {
        'url': 'https://link.berriz.in/en/web/main/kiiikiii/live/019fef22-1b10-2950-8dc0-6e9b20226ce4',
        'only_matching': True,
    }]

    def _playback_paths(self, media_id, kind, is_replay):
        live = f'/service/v1/medias/live/{media_id}/playback_info'
        replay = f'/service/v1/medias/live/replay/{media_id}/playback_info'
        vod = f'/service/v1/medias/{media_id}/playback_info'

        if is_replay:
            return [replay, vod, live]
        if kind == 'live':
            return [live, replay, vod]
        return [vod, replay, live]

    @staticmethod
    def _playback_url(data):
        """Lives carry a tokenized IVS URL; replays/VODs a plain CDN HLS manifest."""
        return traverse_obj(data, ('replay', 'hls', 'playbackUrl')) \
            or traverse_obj(data, ('hls', 'playbackUrl')) \
            or traverse_obj(data, 'playbackUrl')

    def _extract_playback_data(self, media_id, kind, is_replay, live_status=None):
        # Paths are ordered most-likely-first, so the earliest error is the one
        # that describes the URL the user actually asked for. Keeping the last
        # one instead reports whatever the least relevant endpoint said - an
        # on-air live ends on the VOD path's 'This is a deleted media'.
        first_error = None

        for path in self._playback_paths(media_id, kind, is_replay):
            for attempt in (1, 2):
                response = self._download_api(
                    path, media_id, 'Downloading playback info', fatal=False,
                    expected_status=(401, 403, 404))

                if not response:
                    break

                code = response.get('code')
                if code == '0000':
                    data = response.get('data') or {}
                    if self._playback_url(data):
                        return data

                # A 0000 that got this far means the payload simply had no URL
                # in it - the shape berriz uses for a stream that is not
                # currently playable. Its message is the literal 'SUCCESS',
                # which is worse than useless as a failure reason.
                message = None if code == '0000' else (response.get('message') or code)
                if message and first_error is None:
                    first_error = message

                # This is the HTTP 200 + error code form of a stale bz_a, which
                # _download_api cannot see. Refresh once, then replay the path;
                # _refresh_token remembers a dead bz_r, so this cannot spin.
                if code == 'FS_ER4020' or 'log in' in (message or '').lower():
                    if attempt == 1 and self._refresh_token(media_id):
                        continue
                    self.raise_login_required(
                        'berriz rejected the session. Log in at berriz.in and pass '
                        '--cookies-from-browser BROWSER', method='cookies')
                break

        # Every path having failed on a just-ended broadcast is the expected
        # outcome, not a breakage, so say so instead of surfacing the raw
        # endpoint errors ('An error occurred during the service').
        if live_status == 'post_live':
            raise ExtractorError(
                'This broadcast has ended and its replay is not published yet. '
                'Try this URL again later', expected=True, video_id=media_id)

        raise ExtractorError(
            f'Could not get a playback URL: {first_error or "unknown error"}',
            expected=bool(first_error), video_id=media_id)

    def _real_extract(self, url):
        media_id, slug, kind, replay = self._match_valid_url(url).group(
            'id', 'community', 'kind', 'replay')
        is_replay = bool(replay)

        # Metadata is public; fetch it first so failures still report what the video is.
        # The artist list sits beside `media` rather than inside it, so keep the
        # whole `data` payload instead of narrowing straight to the media object.
        data = traverse_obj(self._download_api(
            f'/service/v1/medias/live/{media_id}', media_id,
            'Downloading media metadata', fatal=False,
            expected_status=(401, 404)), 'data') or {}
        if not data.get('media'):
            data = traverse_obj(self._download_api(
                f'/service/v1/medias/{media_id}', media_id,
                'Downloading media metadata', fatal=False,
                expected_status=(401, 404)), 'data') or {}
        meta = data.get('media') or {}

        live_status_raw = traverse_obj(meta, ('live', 'liveStatus'))
        live_status = _live_status(live_status_raw)
        if live_status_raw:
            is_live = live_status_raw == 'ON_AIR'
        else:
            # Both metadata calls are non-fatal, so meta can be empty. Falling
            # through with is_live=False would hand a live stream to the HLS
            # parser as a finished VOD; the URL kind is the best hint left.
            is_live = kind == 'live' and not is_replay

        if not self._logged_in():
            self.raise_login_required(
                'berriz requires an account to fetch stream URLs. Use '
                '--cookies-from-browser BROWSER after logging in at berriz.in',
                method='cookies')

        playback = self._extract_playback_data(media_id, kind, is_replay, live_status)
        replay_info = playback.get('replay') or {}
        if replay_info.get('isDrm') or replay_info.get('drmInfo'):
            self.report_drm(media_id)

        playback_url = self._playback_url(playback)

        # Lives resolve to a pre-signed IVS playlist, replays to a plain CDN manifest;
        # neither needs cookies past this point.
        # Note: with live=True yt-dlp deliberately numbers formats instead of naming
        # them (bandwidth drifts mid-broadcast), so select by resolution on lives.
        formats, subtitles = self._extract_m3u8_formats_and_subtitles(
            playback_url, media_id, 'mp4', m3u8_id='hls', live=is_live)

        # liveReport carries a cumulative total, not a concurrent-viewer gauge,
        # so it only ever populates view_count - and only when berriz itself
        # shows it. With exposeTotalViewCount false the site hides the counter,
        # and an on-air live reports a flat 0 that would otherwise be published
        # as a real tally.
        view_count = None
        if traverse_obj(meta, ('live', 'exposeTotalViewCount')):
            view_count = traverse_obj(
                meta, ('live', 'liveReport', 'totalViewCount'), expected_type=int_or_none)

        channel = (_community_name(meta.get('communityName'))
                   or self._community_name_for(slug) or slug)

        return {
            'id': media_id,
            'title': meta.get('title') or media_id,
            'formats': formats,
            'subtitles': subtitles,
            'is_live': is_live,
            'live_status': live_status,
            'duration': int_or_none(
                replay_info.get('duration')
                or traverse_obj(meta, ('live', 'replay', 'duration'))),
            'thumbnail': meta.get('thumbnailUrl'),
            'timestamp': unified_timestamp(meta.get('publishedAt')),
            'channel': channel,
            'channel_id': str_or_none(meta.get('communityId')),
            'uploader': channel,
            # The slug is the one community identifier that cannot come back
            # localised, mistranslated or missing.
            'uploader_id': slug,
            'view_count': view_count,
            'artists': traverse_obj(data, ('communityArtists', ..., 'name')) or None,
        }


class _BerrizListBaseIE(_BerrizAuthMixin, InfoExtractor):
    """Community lookup, cursor pagination and entry building for the list pages."""

    # The list endpoints accept at least this much; it keeps the request count down.
    _PAGE_SIZE = 100

    def _resolve_community(self, slug):
        community = self._call_api(
            f'/service/v1/community/id/{slug}', slug, 'Resolving community', fatal=False)
        if not traverse_obj(community, 'communityId'):
            raise ExtractorError(f'No berriz community named {slug!r}', expected=True)
        return community

    def _paginate(self, path, item_id, query=None, note='page'):
        """Walk a cursor-paginated list endpoint, yielding raw content items."""
        cursor = None
        seen_cursors = set()
        for page in itertools.count(1):
            data = self._call_api(
                path, item_id, f'Downloading {note} {page}',
                query={'pageSize': self._PAGE_SIZE, **(query or {}),
                       **({'next': cursor} if cursor is not None else {})}) or {}
            yield from data.get('contents') or []
            cursor = traverse_obj(data, ('cursor', 'next'))
            # Single-shot endpoints such as on-air return neither key at all.
            if not data.get('hasNext') or cursor is None:
                return
            # hasNext with a cursor that never advances would loop forever.
            if cursor in seen_cursors:
                self.report_warning(
                    f'Pagination stalled on a repeated cursor after {page} pages; '
                    'the list may be incomplete', video_id=item_id)
                return
            seen_cursors.add(cursor)

    def _live_entry(self, content, slug, channel=None):
        media = content.get('media') or {}
        media_id = media.get('mediaId')
        if not media_id:
            return None
        return self.url_result(
            f'https://berriz.in/en/{slug}/live/{media_id}', BerrizIE, media_id,
            media.get('title'),
            duration=traverse_obj(
                media, ('live', 'replay', 'duration'), expected_type=int_or_none),
            thumbnail=media.get('thumbnailUrl'),
            timestamp=unified_timestamp(media.get('publishedAt')),
            live_status=_live_status(traverse_obj(media, ('live', 'liveStatus'))),
            channel=channel,
            channel_id=str_or_none(media.get('communityId')),
            artists=traverse_obj(content, ('communityArtists', ..., 'name')) or None)


class BerrizCommunityIE(_BerrizListBaseIE):
    """Resolve an artist page to whatever they currently have on air."""

    IE_NAME = 'berriz:community'
    IE_DESC = 'berriz.in artist page (current live stream)'
    # When nothing follows the language prefix the regex backtracks and hands
    # the locale itself to <id>, turning /en/ into a lookup for a community
    # named "en". Reject an <id> that is a bare locale at the end of the URL.
    # A genuine two-letter slug at the site root is indistinguishable from a
    # locale and loses out; every other extractor here is unaffected, since
    # each requires a further path segment to disambiguate.
    _VALID_URL = rf'{_BASE}(?![a-z]{{2}}(?:-[A-Za-z]{{2,4}})?/?(?:[?#]|$))(?P<id>[^/?#]+)/?(?:[?#]|$)'

    _TESTS = [{
        'url': 'https://berriz.in/en/kiiikiii',
        'only_matching': True,
    }]

    def _real_extract(self, url):
        slug = self._match_id(url)
        community = self._resolve_community(slug)

        live = self._call_api(
            f'/service/v1/community/main/live/{community["communityId"]}', slug,
            'Checking for a live stream', fatal=False)
        media_id = traverse_obj(live, ('media', 'mediaId'))
        if not media_id:
            raise ExtractorError(f'{slug} is not currently live', expected=True)

        # This endpoint omits communityName from the media object entirely, so
        # the name resolved above is what keeps this from falling back to the
        # lowercase slug.
        return self.url_result(
            f'https://berriz.in/en/{slug}/live/{media_id}', BerrizIE,
            media_id, traverse_obj(live, ('media', 'title')),
            channel=(_community_name(traverse_obj(live, ('media', 'communityName')))
                     or _community_name(community.get('communityName')) or slug))


class BerrizCommunityLiveIE(_BerrizListBaseIE):
    """The community's live tab: every live and replay, newest first."""

    IE_NAME = 'berriz:community:live'
    IE_DESC = 'berriz.in artist live tab (all lives and replays)'
    _VALID_URL = rf'''(?x)
        {_BASE}
        (?P<id>[^/?#]+)/live/?(?:[?#]|$)'''

    _TESTS = [{
        'url': 'https://berriz.in/en/kiiikiii/live/',
        'info_dict': {
            'id': 'kiiikiii',
            'title': 'KiiiKiii - Live',
        },
        'playlist_mincount': 40,
    }, {
        'url': 'https://link.berriz.in/en/web/main/kiiikiii/live',
        'only_matching': True,
    }]

    def _real_extract(self, url):
        slug = self._match_id(url)
        community = self._resolve_community(slug)
        community_id = community['communityId']
        name = _community_name(community.get('communityName')) or slug

        def entries():
            seen = set()
            # Streams that are still on air are absent from the ended list, so walk both.
            for content in itertools.chain(
                self._paginate(
                    f'/service/v1/community/{community_id}/medias/live/on-air',
                    slug, note='on-air page'),
                self._paginate(
                    f'/service/v1/community/{community_id}/medias/live/end',
                    slug, note='replay page'),
            ):
                entry = self._live_entry(content, slug, name)
                if entry and entry['id'] not in seen:
                    seen.add(entry['id'])
                    yield entry

        return self.playlist_result(entries(), slug, f'{name} - Live')


class BerrizProfileLiveIE(_BerrizListBaseIE):
    """A single member's live tab, keyed by the community user id in the URL."""

    IE_NAME = 'berriz:profile:live'
    IE_DESC = 'berriz.in member profile live tab (that member\'s lives and replays)'
    _VALID_URL = rf'''(?x)
        {_BASE}
        (?P<community>[^/?#]+)/profile/
        (?P<id>\d+)/live/?(?:[?#]|$)'''

    _TESTS = [{
        'url': 'https://berriz.in/en/kiiikiii/profile/158606553891123/live/',
        'info_dict': {
            'id': 'kiiikiii-158606553891123',
            'title': 'KiiiKiii - LEESOL - Live',
        },
        'playlist_mincount': 5,
    }, {
        'url': 'https://link.berriz.in/en/web/main/kiiikiii/profile/158606553891123/live',
        'only_matching': True,
    }]

    def _real_extract(self, url):
        slug, user_id = self._match_valid_url(url).group('community', 'id')
        community = self._resolve_community(slug)
        community_id = community['communityId']
        channel = _community_name(community.get('communityName')) or slug

        # The URL carries a user id, but only the roster maps it to a display name.
        artist = traverse_obj(self._call_api(
            f'/service/v1/community/main/{community_id}', user_id,
            'Downloading artist roster', fatal=False),
            ('artists', lambda _, v: str(v['userId']) == user_id, any)) or {}
        member = artist.get('name') or user_id

        def entries():
            for content in self._paginate(
                f'/service/v1/community/archive/user/{user_id}/artist/live', user_id,
                query={'communityId': community_id},
            ):
                entry = self._live_entry(content, slug, channel)
                if entry:
                    yield entry

        return self.playlist_result(
            entries(), f'{slug}-{user_id}', f'{channel} - {member} - Live',
            thumbnail=artist.get('imageUrl'))