#!/usr/bin/env python3
"""Download a DRM-protected video from Apple TV+ (Widevine via HLS)."""

import argparse
import base64
import curses
import getpass
import hashlib
import json
import os
import pickle
import random
import re
import secrets
import string
import subprocess
import sys
import time
import unicodedata
from pathlib import Path
from urllib.parse import urlparse, unquote

import requests
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes as crypto_hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

HERE = Path(__file__).parent
TOOLS = HERE / 'tools'
N_M3U8 = TOOLS / 'N_m3u8DL-RE'
MP4DEC = TOOLS / 'mp4decrypt'
SAVE_DIR = HERE / 'downloads'
TMP_DIR = '/tmp/n_m3u8dl_tmp'
COOKIES_FILE = HERE / '.apple_cookies'   # pickled cookie jar
SESSION_FILE = HERE / '.apple_session'   # JSON with token + metadata

UA = ('Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) '
      'AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36')

AUTH_HOST     = 'https://idmsa.apple.com'
OLYMPUS       = 'https://appstoreconnect.apple.com/olympus/v1'
LICENSE_URL   = 'https://play.itunes.apple.com/WebObjects/MZPlay.woa/wa/fpsRequest'
TV_API        = 'https://tv.apple.com/api/uts'
TV_HOME       = 'https://tv.apple.com'
MUSIC_MKT_KEY = '92f19b477c5c9be6ab17f3ec2b1b2b7db4d00a9a8c973e3d6c90dac08b91de71'
BUY_APPLE     = 'https://buy.itunes.apple.com'

# Populated by _get_tv_page_config() from the Apple TV page; used by _api_params()
_API_BASE: dict = {}

STOREFRONT = {
    'pl': '143478', 'us': '143441', 'gb': '143444',
    'de': '143443', 'fr': '143442', 'ca': '143455',
    'au': '143460', 'jp': '143462', 'it': '143450',
    'es': '143454', 'nl': '143452', 'se': '143456',
    'no': '143457', 'dk': '143458', 'fi': '143447',
    'br': '143503', 'mx': '143468', 'kr': '143466',
}

WIDEVINE_KEYFORMAT = 'urn:uuid:edef8ba9-79d6-4ace-a3c8-27dcd51d21ed'

LANG_NAMES = {
    'ar': 'Arabic', 'bg': 'Bulgarian', 'cs': 'Czech', 'da': 'Danish',
    'de': 'German', 'el': 'Greek', 'en': 'English', 'en-US': 'English (US)',
    'en-GB': 'English (UK)', 'es': 'Spanish', 'et': 'Estonian', 'fi': 'Finnish',
    'fr': 'French', 'hr': 'Croatian', 'hu': 'Hungarian', 'it': 'Italian',
    'ja': 'Japanese', 'ko': 'Korean', 'lt': 'Lithuanian', 'lv': 'Latvian',
    'nl': 'Dutch', 'nb': 'Norwegian', 'pl': 'Polish', 'pt': 'Portuguese',
    'pt-BR': 'Portuguese (BR)', 'ro': 'Romanian', 'ru': 'Russian',
    'sk': 'Slovak', 'sl': 'Slovenian', 'sr': 'Serbian', 'sv': 'Swedish',
    'tr': 'Turkish', 'uk': 'Ukrainian', 'zh': 'Chinese',
    'zh-CN': 'Chinese (Simp.)', 'zh-TW': 'Chinese (Trad.)', 'und': 'Undetermined',
}


def lang_name(code):
    return LANG_NAMES.get(code, code)


def _b64decode(s: str) -> bytes:
    """Decode base64 that may be standard or URL-safe, with or without padding."""
    s = s.strip().replace('-', '+').replace('_', '/')
    rem = len(s) % 4
    if rem:
        s += '=' * (4 - rem)
    return base64.b64decode(s)


# ── Widevine device ──────────────────────────────────────────────────────────

def _find_wvd():
    from pywidevine.device import Device
    for p in sorted((HERE / 'L3').glob('*.wvd')):
        try:
            d = Device.load(str(p))
            if d.system_id and d.system_id != 0:
                return p
        except Exception:
            pass
    return next((HERE / 'L3').glob('*.wvd'), None)

WVD = _find_wvd()


# ── SRP-6a (Apple's SIRP variant) ───────────────────────────────────────────

class AppleSRP:
    """SRP-6a with Apple's PBKDF2-based password hash and SHA-256 group.

    Apple SIRP uses the 2048-bit Oakley prime (same as Tom Wu SRP / srp package NG_2048).
    N is 2048 bits = 256 bytes. Verified by decoding Apple's SRP init 'b' response.
    """

    # 2048-bit SRP Oakley prime (from srp package NG_2048 / Tom Wu's SRP 2048-bit group)
    # 2048 actual bits, 512 hex chars, 256 bytes
    N = int(
        'AC6BDB41324A9A9BF166DE5E1389582FAF72B6651987EE07'
        'FC3192943DB56050A37329CBB4A099ED8193E0757767A13D'
        'D52312AB4B03310DCD7F48A9DA04FD50E8083969EDB767B0'
        'CF6095179A163AB3661A05FBD5FAAAE82918A9962F0B93B8'
        '55F97993EC975EEAA80D740ADBF4FF747359D041D5C33EA7'
        '1D281E446B14773BCA97B43A23FB801676BD207A436C6481'
        'F1D2B9078717461A5B9D32E688F87748544523B524B0D57D'
        '5EA77A2775D2ECFA032CFBDBF52FB3786160279004E57AE6'
        'AF874E7303CE53299CCC041C7BC308D82A5698F3A8D0C382'
        '71AE35F8E9DBFBB694B5C803D89F7AE435DE236D525F5475'
        '9B65E372FCD68EF20FA7111F9E4AFF73',
        16,
    )
    G = 2
    N_LEN = (N.bit_length() + 7) // 8  # 256 bytes

    def __init__(self, account: str, password: str):
        self.account = account
        self.password = password
        self._a = int.from_bytes(secrets.token_bytes(32), 'big')
        self.A = pow(self.G, self._a, self.N)

    def A_bytes(self) -> bytes:
        return self.A.to_bytes(self.N_LEN, 'big')

    def A_b64(self) -> str:
        return base64.b64encode(self.A_bytes()).decode()

    def _x(self, salt: bytes, iterations: int, protocol: str) -> int:
        """Apple SIRP x derivation (reverse-engineered from webSRPClientWorker.js).

        Step 1 — PBKDF2:
          s2k:    pbkdf_key = SHA256(password)           (32 raw bytes)
          s2k_fo: pbkdf_key = SHA256(password).hex()     (64 ASCII bytes, hex repr)
          pbkdf = PBKDF2-HMAC-SHA256(key=pbkdf_key, salt=salt, iter=iterations, dkLen=32)

        Step 2 — E() function:
          x = SHA256(salt || SHA256(b":" || pbkdf))
        """
        pw = self.password
        if protocol == 's2k_fo':
            pw = unicodedata.normalize('NFKD', pw).casefold()

        pw_hash = hashlib.sha256(pw.encode('utf-8')).digest()

        if protocol == 's2k_fo':
            # importKey gets the hex-string bytes of SHA256(password)
            pbkdf_key = pw_hash.hex().encode('ascii')  # 64 bytes
        else:
            # importKey gets the raw SHA256(password) bytes
            pbkdf_key = pw_hash  # 32 bytes

        kdf = PBKDF2HMAC(
            algorithm=crypto_hashes.SHA256(),
            length=32, salt=salt, iterations=iterations,
            backend=default_backend(),
        )
        pbkdf_bytes = kdf.derive(pbkdf_key)

        # E() function: SHA256(salt || SHA256(b":" || pbkdf))
        intermediate = hashlib.sha256(b':' + pbkdf_bytes).digest()
        x_bytes = hashlib.sha256(salt + intermediate).digest()
        return int.from_bytes(x_bytes, 'big')

    def compute_m1(self, B_b64: str, salt_b64: str, iterations: int, protocol: str):
        """Return (M1_b64, M2_b64).

        M1 = H(H(N) XOR H(g) | H(account) | salt | A | B | K)
        M2 = H(A | M1 | K)   — client's expected server proof, required by Apple
        """
        B_bytes    = _b64decode(B_b64)
        salt_bytes = _b64decode(salt_b64)

        # Pad B to N_LEN bytes in case leading zeros were dropped in transmission
        if len(B_bytes) < self.N_LEN:
            B_bytes = B_bytes.rjust(self.N_LEN, b'\x00')
        B = int.from_bytes(B_bytes, 'big')

        N_b = self.N.to_bytes(self.N_LEN, 'big')
        G_b = self.G.to_bytes(self.N_LEN, 'big')
        A_b = self.A_bytes()

        k = int.from_bytes(hashlib.sha256(N_b + G_b).digest(), 'big')
        x = self._x(salt_bytes, iterations, protocol)
        u = int.from_bytes(hashlib.sha256(A_b + B_bytes).digest(), 'big')

        gx   = pow(self.G, x, self.N)
        kgx  = (k * gx) % self.N
        base = (B - kgx) % self.N
        exp  = (self._a + u * x) % (self.N - 1)
        S    = pow(base, exp, self.N)

        K = hashlib.sha256(S.to_bytes(self.N_LEN, 'big')).digest()

        h_N = hashlib.sha256(N_b).digest()
        h_G = hashlib.sha256(G_b).digest()
        xNG = bytes(a ^ b for a, b in zip(h_N, h_G))
        h_I = hashlib.sha256(self.account.lower().encode('utf-8')).digest()

        M1 = hashlib.sha256(xNG + h_I + salt_bytes + A_b + B_bytes + K).digest()
        M2 = hashlib.sha256(A_b + M1 + K).digest()
        return base64.b64encode(M1).decode(), base64.b64encode(M2).decode()


# ── Hashcash (Apple proof-of-work) ───────────────────────────────────────────

def _hashcash(bits: int, challenge: str) -> str:
    """Compute SHA-1 hashcash stamp satisfying 'bits' leading zero bits."""
    bits = int(bits)
    n_zero = bits // 8
    n_extra = bits % 8
    chars = string.ascii_letters + string.digits
    for _ in range(5_000_000):
        nonce = ''.join(random.choices(chars, k=16))
        stamp = f'1:{bits}:{challenge}:{nonce}'
        h = hashlib.sha1(stamp.encode()).digest()
        ok = all(b == 0 for b in h[:n_zero])
        if ok and n_extra and (h[n_zero] >> (8 - n_extra)):
            ok = False
        if ok:
            return stamp
    raise RuntimeError('Hashcash PoW failed — this should not happen')


# ── Apple ID authentication ──────────────────────────────────────────────────

def _get_service_key(session: requests.Session) -> str:
    """Fetch Apple auth widget key for Apple TV+ login."""
    # 1. Try appleid.apple.com — the canonical Apple ID sign-in page embeds it
    try:
        r = session.get(
            'https://appleid.apple.com/',
            headers={'User-Agent': UA},
            timeout=10,
        )
        if r.ok:
            m = re.search(r'"widgetKey"\s*:\s*"([^"]+)"', r.text)
            if m:
                return m.group(1)
            m2 = re.search(r"widgetKey\s*=\s*['\"]([^'\"]+)['\"]", r.text)
            if m2:
                return m2.group(1)
    except Exception:
        pass
    # 2. Try idmsa sign-in page response header
    try:
        r = session.get(
            f'{AUTH_HOST}/appleauth/auth/signin',
            params={'widgetKey': ''},
            headers={'User-Agent': UA},
            timeout=10,
        )
        key = r.headers.get('X-Apple-Widget-Key', '')
        if key:
            return key
    except Exception:
        pass
    # 3. Try App Store Connect olympus endpoint (may not be the right key for TV)
    try:
        r = session.get(
            f'{OLYMPUS}/app/config?hostname=itunesconnect.apple.com',
            timeout=10,
        )
        if r.ok:
            key = r.json().get('authServiceKey', '')
            if key:
                return key
    except Exception:
        pass
    raise RuntimeError(
        'Cannot determine Apple auth service key. '
        'Check your internet connection or update the hardcoded key.'
    )


def _auth_headers(service_key: str, session_id: str = '', scnt: str = '') -> dict:
    h = {
        'Content-Type': 'application/json',
        'Accept': 'application/json',
        'X-Apple-Widget-Key': service_key,
        'X-Requested-With': 'XMLHttpRequest',
        'User-Agent': UA,
        'Origin': 'https://appleid.apple.com',
        'Referer': 'https://appleid.apple.com/',
    }
    if session_id:
        h['X-Apple-ID-Session-Id'] = session_id
    if scnt:
        h['scnt'] = scnt
    return h


def _get_hashcash_params(session: requests.Session, service_key: str):
    """GET /appleauth/auth/signin to obtain PoW parameters."""
    r = session.get(
        f'{AUTH_HOST}/appleauth/auth/signin',
        params={'widgetKey': service_key},
        headers=_auth_headers(service_key),
        timeout=10,
    )
    bits      = r.headers.get('X-Apple-HC-Bits', '0')
    challenge = r.headers.get('X-Apple-HC-Challenge', '')
    return bits, challenge


def apple_login(session: requests.Session, username: str, password: str) -> dict:
    """Full Apple ID SIRP login with 2FA support.

    Returns a dict with 'media_user_token', 'myacinfo', 'itua'.
    Raises RuntimeError on failure.
    """
    print('Fetching Apple auth service key...')
    service_key = _get_service_key(session)
    print(f'  Service key: {service_key[:20]}...')

    srp = AppleSRP(username, password)

    # ── Step 1: SRP init ────────────────────────────────────────────────────
    print('Starting SRP handshake...')
    r = session.post(
        f'{AUTH_HOST}/appleauth/auth/signin/init',
        json={
            'a': srp.A_b64(),
            'accountName': username,
            'protocols': ['s2k', 's2k_fo'],
        },
        headers=_auth_headers(service_key),
        timeout=15,
    )
    if not r.ok:
        raise RuntimeError(f'SRP init failed (HTTP {r.status_code}): {r.text[:200]}')
    init_data = r.json()
    print(f'  Init keys:   {list(init_data.keys())}')

    # ── Step 2: Compute M1 + M2 ────────────────────────────────────────────
    print('Computing SRP response...')
    b_raw    = _b64decode(init_data['b'])
    salt_raw = _b64decode(init_data['salt'])
    protocol = init_data.get('protocol', 's2k')
    iters    = init_data['iteration']
    print(f'  Server B:    {len(b_raw)} bytes ({len(b_raw)*8}-bit group)')
    print(f'  Salt:        {len(salt_raw)} bytes = {salt_raw.hex()}')
    print(f'  Iterations:  {iters}')
    print(f'  Protocol:    {protocol}')
    print(f'  B (b64 raw): {init_data["b"][:32]}...')
    print(f'  c:           {init_data.get("c", "")[:32]}...')
    m1_b64, m2_b64 = srp.compute_m1(
        B_b64      = init_data['b'],
        salt_b64   = init_data['salt'],
        iterations = iters,
        protocol   = protocol,
    )
    c = init_data['c']
    print(f'  M1: {m1_b64}')
    print(f'  M2: {m2_b64[:16]}...')

    # ── Step 3: Hashcash PoW ────────────────────────────────────────────────
    hc_bits, hc_challenge = _get_hashcash_params(session, service_key)
    hc = ''
    if hc_bits and hc_challenge and int(hc_bits) > 0:
        print(f'Computing hashcash ({hc_bits} bits)...')
        hc = _hashcash(int(hc_bits), hc_challenge)

    # ── Step 4: SRP complete ─────────────────────────────────────────────────
    print('Submitting credentials...')
    headers = _auth_headers(service_key)
    if hc:
        headers['X-Apple-HC'] = hc
    r = session.post(
        f'{AUTH_HOST}/appleauth/auth/signin/complete',
        params={'isRememberMeEnabled': 'false'},
        json={
            'accountName': username,
            'c': c,
            'm1': m1_b64,
            'm2': m2_b64,
            'rememberMe': False,
        },
        headers=headers,
        timeout=20,
    )

    session_id = r.headers.get('X-Apple-ID-Session-Id', '')
    scnt       = r.headers.get('scnt', '')

    if r.status_code == 200:
        # Login succeeded without 2FA
        print('Login succeeded (no 2FA required).')
        return _extract_tokens(session, service_key, session_id, scnt)

    if r.status_code == 409:
        # 2FA required
        return _handle_2fa(session, service_key, session_id, scnt)

    body = r.text[:500]
    print(f'  Response headers: {dict(r.headers)}')
    raise RuntimeError(f'Login failed (HTTP {r.status_code}): {body}')


def _handle_2fa(session: requests.Session, service_key: str,
                session_id: str, scnt: str) -> dict:
    """Handle two-factor authentication (trusted device or SMS)."""
    headers = _auth_headers(service_key, session_id, scnt)

    # Get available 2FA methods
    r = session.get(
        f'{AUTH_HOST}/appleauth/auth',
        headers=headers,
        timeout=15,
    )
    auth_data = r.json() if r.ok else {}
    phone_numbers = auth_data.get('trustedPhoneNumbers', [])

    print('\n2FA is required.')

    # Prefer trusted-device code if available
    security_devices = auth_data.get('trustedDevices', []) or auth_data.get('securityCode', {})
    if security_devices or not phone_numbers:
        code = input('Enter 6-digit verification code from your trusted device: ').strip()
        r = session.post(
            f'{AUTH_HOST}/appleauth/auth/verify/trusteddevice/securitycode',
            json={'securityCode': {'code': code}},
            headers=headers,
            timeout=15,
        )
    else:
        print('Available phone numbers for SMS verification:')
        for i, ph in enumerate(phone_numbers):
            num = ph.get('numberWithDialCode', ph.get('number', f'#{i+1}'))
            print(f'  [{i+1}] {num}')
        choice = input('Select number (1): ').strip() or '1'
        ph_idx = int(choice) - 1
        ph_id = phone_numbers[ph_idx]['id']

        # Trigger SMS
        session.put(
            f'{AUTH_HOST}/appleauth/auth/verify/phone',
            json={'phoneNumber': {'id': ph_id}, 'mode': 'sms'},
            headers=headers,
            timeout=15,
        )
        code = input('Enter SMS verification code: ').strip()
        r = session.post(
            f'{AUTH_HOST}/appleauth/auth/verify/phone/securitycode',
            json={
                'phoneNumber': {'id': ph_id},
                'securityCode': {'code': code},
                'mode': 'sms',
            },
            headers=headers,
            timeout=15,
        )

    if r.status_code not in (200, 204):
        raise RuntimeError(f'2FA verification failed (HTTP {r.status_code}): {r.text[:200]}')

    print('2FA verified.')

    # Trust this session so we don't need 2FA again for a while
    try:
        session.post(
            f'{AUTH_HOST}/appleauth/auth/2sv/trust',
            headers=headers,
            timeout=10,
        )
    except Exception:
        pass

    return _extract_tokens(session, service_key, session_id, scnt)


def _extract_tokens(session: requests.Session, service_key: str,
                    session_id: str, scnt: str) -> dict:
    """Get the olympus session and extract media-user-token."""
    try:
        r = session.get(f'{OLYMPUS}/session', timeout=15)
    except Exception:
        pass

    mut = session.cookies.get('media-user-token', '')
    myacinfo = session.cookies.get('myacinfo', '')
    itua = session.cookies.get('itua', '')

    if not mut:
        # media-user-token may be in response headers on some flows
        for cookie in session.cookies:
            if 'media' in cookie.name.lower() and 'token' in cookie.name.lower():
                mut = cookie.value
                break

    return {
        'media_user_token': mut,
        'myacinfo': myacinfo,
        'itua': itua,
    }


# ── MusicKit / media-user-token ──────────────────────────────────────────────

def get_media_user_token(session: requests.Session, dev_token: str,
                          username: str, password: str) -> str:
    """Obtain media-user-token via MusicKit OAuth + SIRP.

    Flow:
      authenticateMusicKitRequest → IDMSA OAuth page (scnt, oAuthToken, liteSessionId)
      → SIRP with MUSIC_MKT_KEY → POST /auth → getMusicUserToken

    Uses a separate session so main session cookies are not touched.
    Raises RuntimeError if the flow cannot complete (e.g. 2FA needed).
    """
    # 1 — try simple direct-exchange endpoints first (work when myacinfo is already
    #     valid for the buy.itunes.apple.com domain)
    for path in ('/commerce/account/webPlayerToken',
                 '/commerce/account/getMusicUserToken'):
        try:
            r = session.get(
                f'{BUY_APPLE}{path}',
                params={'jwtToken': dev_token},
                headers={'Accept': 'application/json'},
                timeout=10,
            )
            if r.ok:
                data = r.json()
                for key in ('musicUserToken', 'music-user-token', 'userToken', 'token'):
                    if data.get(key):
                        return data[key]
        except Exception:
            pass

    # 2 — full OAuth + SIRP flow (requires credentials)
    if not username or not password:
        raise RuntimeError(
            'Credentials required for automated media-user-token flow.\n'
            'Options:\n'
            '  a) Re-run with -u <email> -p <password> so the flow can complete.\n'
            '  b) Provide the token directly: --media-token <token> or APPLE_MUT env var.\n'
            '     To get it: open tv.apple.com in a browser, then\n'
            '     DevTools → Application → Cookies → .apple.com → media-user-token'
        )

    ms = requests.Session()
    ms.headers['User-Agent'] = UA
    for c in session.cookies:
        ms.cookies.set(c.name, c.value, domain=c.domain, path=c.path)

    # 2a — kick off MusicKit OAuth
    r1 = ms.post(
        f'{BUY_APPLE}/commerce/account/authenticateMusicKitRequest',
        data={'jwtToken': dev_token, 'isWebPlayer': 'true', 'LogoURL': ''},
        headers={'Content-Type': 'application/x-www-form-urlencoded',
                 'Accept': 'text/html,*/*'},
        allow_redirects=False, timeout=10,
    )
    if r1.status_code not in (301, 302, 303, 307, 308):
        raise RuntimeError(
            f'authenticateMusicKitRequest returned HTTP {r1.status_code}'
        )
    idmsa_url = r1.headers.get('Location', '')

    # 2b — fetch IDMSA OAuth page
    r2 = ms.get(idmsa_url, timeout=10, headers={'Accept': 'text/html'})
    html = r2.text

    page_scnt  = (re.search(r'name="scnt"\s+value="([^"]+)"',       html) or ['', ''])[1] or \
                 (re.search(r'name="scnt"\s+value=\'([^\']+)\'',     html) or ['', ''])[1]
    oauth_tok  = (re.search(r'name="oAuthToken"\s+value="([^"]+)"',  html) or ['', ''])[1]
    if not page_scnt:
        m = re.search(r'"scnt"\s*:\s*"([^"]+)"', html)
        page_scnt = m.group(1) if m else ''

    lite_id = validate_host = None
    am = re.search(r'"appReturnUrl"\s*:\s*"([^"]+)"', html)
    if am:
        app_ret = unquote(am.group(1))
        lm = re.search(r'liteSessionId=([^&"]+)', app_ret)
        if lm:
            lite_id = lm.group(1)
        parsed = urlparse(app_ret)
        validate_host = f'{parsed.scheme}://{parsed.netloc}'

    if not all([page_scnt, oauth_tok, lite_id, validate_host]):
        raise RuntimeError(
            f'Missing OAuth params from IDMSA page '
            f'(scnt={bool(page_scnt)}, oAuthToken={bool(oauth_tok)}, '
            f'liteSessionId={lite_id}, host={validate_host})'
        )

    # 2c — SIRP with Music widget key (fresh session, no pre-existing scnt)
    srp = AppleSRP(username, password)

    init_r = ms.post(
        f'{AUTH_HOST}/appleauth/auth/signin/init',
        json={'a': srp.A_b64(), 'accountName': username,
              'protocols': ['s2k', 's2k_fo']},
        headers=_auth_headers(MUSIC_MKT_KEY),
        timeout=15,
    )
    if not init_r.ok:
        raise RuntimeError(
            f'Music SIRP init failed ({init_r.status_code}): {init_r.text[:200]}'
        )
    init_data   = init_r.json()
    sirp_scnt   = init_r.headers.get('scnt', page_scnt)
    sirp_sid    = init_r.headers.get('X-Apple-ID-Session-Id', '')

    m1, m2 = srp.compute_m1(
        B_b64=init_data['b'], salt_b64=init_data['salt'],
        iterations=init_data['iteration'],
        protocol=init_data.get('protocol', 's2k'),
    )

    hc_bits, hc_chal = _get_hashcash_params(ms, MUSIC_MKT_KEY)
    hc = _hashcash(int(hc_bits), hc_chal) if hc_bits and hc_chal and int(hc_bits) > 0 else ''

    comp_headers = _auth_headers(MUSIC_MKT_KEY, sirp_sid, sirp_scnt)
    if hc:
        comp_headers['X-Apple-HC'] = hc

    comp_r = ms.post(
        f'{AUTH_HOST}/appleauth/auth/signin/complete',
        params={'isRememberMeEnabled': 'false'},
        json={'accountName': username, 'c': init_data['c'],
              'm1': m1, 'm2': m2, 'rememberMe': False},
        headers=comp_headers, timeout=20,
    )

    if comp_r.status_code == 409:
        # Apple requires 2FA for the Music widget key session — handle it inline
        music_sid  = comp_r.headers.get('X-Apple-ID-Session-Id', sirp_sid)
        music_scnt = comp_r.headers.get('scnt', sirp_scnt)
        music_hdrs = _auth_headers(MUSIC_MKT_KEY, music_sid, music_scnt)

        auth_r    = ms.get(f'{AUTH_HOST}/appleauth/auth', headers=music_hdrs, timeout=15)
        auth_data = auth_r.json() if auth_r.ok else {}
        phones    = auth_data.get('trustedPhoneNumbers', [])
        sec_devs  = auth_data.get('trustedDevices', []) or auth_data.get('securityCode', {})

        print('\n2FA required for Music authorisation.')
        if sec_devs or not phones:
            code = input('Enter 6-digit verification code from your trusted device: ').strip()
            v_r = ms.post(
                f'{AUTH_HOST}/appleauth/auth/verify/trusteddevice/securitycode',
                json={'securityCode': {'code': code}},
                headers=music_hdrs, timeout=15,
            )
        else:
            print('Available phone numbers for SMS verification:')
            for i, ph in enumerate(phones):
                print(f'  [{i+1}] {ph.get("numberWithDialCode", ph.get("number", f"#{i+1}"))}')
            choice = input('Select number (1): ').strip() or '1'
            ph_id = phones[int(choice) - 1]['id']
            ms.put(f'{AUTH_HOST}/appleauth/auth/verify/phone',
                   json={'phoneNumber': {'id': ph_id}, 'mode': 'sms'},
                   headers=music_hdrs, timeout=15)
            code = input('Enter SMS verification code: ').strip()
            v_r = ms.post(
                f'{AUTH_HOST}/appleauth/auth/verify/phone/securitycode',
                json={'phoneNumber': {'id': ph_id},
                      'securityCode': {'code': code}, 'mode': 'sms'},
                headers=music_hdrs, timeout=15,
            )

        if v_r.status_code not in (200, 204):
            raise RuntimeError(
                f'Music 2FA verification failed (HTTP {v_r.status_code}): {v_r.text[:200]}\n'
                'Provide the token via --media-token or APPLE_MUT env var.\n'
                'To get it: open tv.apple.com in your browser (logged in),\n'
                'open DevTools → Application → Cookies → .apple.com → media-user-token'
            )
        print('Music 2FA verified.')
        try:
            ms.post(f'{AUTH_HOST}/appleauth/auth/2sv/trust', headers=music_hdrs, timeout=10)
        except Exception:
            pass

    elif not comp_r.ok:
        raise RuntimeError(
            f'Music SIRP complete failed ({comp_r.status_code}): {comp_r.text[:200]}'
        )

    # 2d — submit IDMSA OAuth form with the original OAuth scnt (not the SIRP auth scnt)
    ms.post(
        'https://idmsa.apple.com/auth',
        data={'rememberMe': 'false', 'oAuthToken': oauth_tok, 'scnt': page_scnt},
        headers={
            'Content-Type': 'application/x-www-form-urlencoded',
            'Accept': 'text/html,*/*',
            'Referer': idmsa_url,
            'Origin': 'https://idmsa.apple.com',
            'X-Apple-Widget-Key': MUSIC_MKT_KEY,
        },
        allow_redirects=True, timeout=15,
    )

    # 2e — getMusicUserToken
    mut_r = ms.get(
        f'{validate_host}/commerce/account/getMusicUserToken',
        params={'liteSessionId': lite_id, 'jwtToken': dev_token},
        headers={'Accept': 'application/json'},
        timeout=10,
    )
    if mut_r.ok:
        data = mut_r.json()
        for key in ('musicUserToken', 'music-user-token', 'userToken'):
            if data.get(key):
                return data[key]

    raise RuntimeError(
        f'getMusicUserToken failed ({mut_r.status_code}): {mut_r.text[:300]}\n'
        'Provide token via --media-token or APPLE_MUT env var.\n'
        'To get it: open tv.apple.com in your browser (logged in),\n'
        'open DevTools → Application → Cookies → .apple.com → media-user-token'
    )


# ── Session persistence ──────────────────────────────────────────────────────

def load_session(session: requests.Session) -> dict:
    """Restore cookies and session metadata from disk. Returns metadata dict."""
    meta = {}
    if COOKIES_FILE.exists():
        try:
            with open(COOKIES_FILE, 'rb') as f:
                jar = pickle.load(f)
            session.cookies.update(jar)
        except Exception as e:
            print(f'[warn] Could not load cookies: {e}')
    if SESSION_FILE.exists():
        try:
            meta = json.loads(SESSION_FILE.read_text())
        except Exception:
            pass
    return meta


def save_session(session: requests.Session, meta: dict):
    """Persist cookies and session metadata to disk."""
    with open(COOKIES_FILE, 'wb') as f:
        pickle.dump(session.cookies, f)
    COOKIES_FILE.chmod(0o600)
    SESSION_FILE.write_text(json.dumps(meta, indent=2))
    SESSION_FILE.chmod(0o600)


def is_session_valid(session: requests.Session, meta: dict) -> bool:
    """Check that the myacinfo auth cookie is present (set after Apple ID login)."""
    return bool(session.cookies.get('myacinfo', ''))


# ── Developer token + utsk (public API key extracted from page HTML) ──────────

def _get_tv_page_config(session: requests.Session, country: str = 'us') -> tuple:
    """Fetch tv.apple.com and extract (dev_token, utsk). Also updates _API_BASE.

    Returns (dev_token: str, utsk: str).
    utsk is the user session key required as URL param in content API calls.
    """
    global _API_BASE
    r = session.get(
        f'{TV_HOME}/{country}',
        headers={'User-Agent': UA, 'Accept': 'text/html'},
        timeout=15,
    )
    if not r.ok:
        raise RuntimeError(f'Apple TV page returned HTTP {r.status_code}')

    dev_token = ''
    utsk = ''

    m = re.search(
        r'<script[^>]+id=["\']serialized-server-data["\'][^>]*>(.*?)</script>',
        r.text, re.DOTALL,
    )
    if m:
        try:
            data = json.loads(m.group(1))
            # Structure: {"data": [{"data": {"configureParams": ..., "configuration": ...}}]}
            entry = data['data'][0]['data']
            dev_token = entry['configureParams']['developerToken']
            cfg = entry['configuration']
            utsk = cfg.get('utskProps', {}).get('utsk', '')
            # Use the live requiredParamsMap.Default as base for all API calls
            default_params = (cfg.get('applicationProps', {})
                                 .get('requiredParamsMap', {})
                                 .get('Default', {}))
            if default_params:
                # Strip utsk — we inject it per-country in _api_params
                _API_BASE = {k: v for k, v in default_params.items()
                             if k not in ('sf', 'locale', 'utsk')}
                _API_BASE['utsk'] = utsk
        except (KeyError, IndexError, json.JSONDecodeError):
            pass

    if not dev_token:
        # Fallback: regex anywhere in the page
        m2 = re.search(r'"developerToken"\s*:\s*"([^"]{40,})"', r.text)
        if not m2:
            raise RuntimeError(
                'Cannot extract developerToken from Apple TV page. '
                'The page structure may have changed.'
            )
        dev_token = m2.group(1)

    return dev_token, utsk


# ── Apple TV content API ──────────────────────────────────────────────────────

LOCALE_MAP = {
    'us': 'en-US', 'gb': 'en-GB', 'ca': 'en-CA', 'au': 'en-AU',
    'pl': 'pl-PL', 'de': 'de-DE', 'fr': 'fr-FR', 'it': 'it-IT',
    'es': 'es-ES', 'nl': 'nl-NL', 'se': 'sv-SE', 'no': 'nb-NO',
    'dk': 'da-DK', 'fi': 'fi-FI', 'br': 'pt-BR', 'mx': 'es-MX',
    'jp': 'ja-JP', 'kr': 'ko-KR',
}


def _sf(country: str) -> str:
    return STOREFRONT.get(country.lower(), '143441')


def _locale(country: str) -> str:
    return LOCALE_MAP.get(country.lower(), 'en-US')


def _api_params(country: str, **extra) -> dict:
    base = dict(_API_BASE) if _API_BASE else {
        'utscf': 'OjAAAAAAAAA~',
        'caller': 'web',
        'v': '94',
        'pfm': 'web',
    }
    # Always override country-specific fields
    base['sf'] = _sf(country)
    base['locale'] = _locale(country)
    base.update(extra)
    return base


def _parse_url(url: str):
    """Return (country, kind, content_id) from an Apple TV URL."""
    m = re.search(
        r'tv\.apple\.com/([a-z]{2})/(\w+)/[^/]+/(umc\.cmc\.[a-z0-9]+)',
        url,
    )
    if not m:
        raise ValueError(f'Unrecognised Apple TV URL: {url}')
    country, kind, cid = m.group(1), m.group(2), m.group(3)
    if kind in ('episode', 'season'):
        kind = 'show'
    return country, kind, cid


def _check_api(r: requests.Response):
    if not r.ok:
        raise RuntimeError(
            f'API {r.status_code} for {r.url.split("?")[0]}: {r.text[:300]}'
        )


def get_show_metadata(session: requests.Session, cid: str, country: str) -> dict:
    r = session.get(
        f'{TV_API}/v3/shows/{cid}',
        params=_api_params(country),
        timeout=15,
    )
    if r.status_code == 404:
        # cid might be an episode ID; try to look up its parent show
        ep_r = session.get(f'{TV_API}/v3/episodes/{cid}',
                           params=_api_params(country), timeout=15)
        if ep_r.ok:
            show_id = ep_r.json().get('data', {}).get('content', {}).get('showId', '')
            if show_id:
                r = session.get(f'{TV_API}/v3/shows/{show_id}',
                                params=_api_params(country), timeout=15)
    _check_api(r)
    return r.json().get('data', {})


def get_movie_metadata(session: requests.Session, cid: str, country: str) -> dict:
    r = session.get(
        f'{TV_API}/v3/movies/{cid}',
        params=_api_params(country),
        timeout=15,
    )
    _check_api(r)
    return r.json().get('data', {})


def get_episodes(session: requests.Session, show_id: str, country: str,
                 season_ids: list | None = None) -> list:
    """Fetch all available episodes for a show across all seasons.

    season_ids: list of season ID strings from the show metadata 'seasons' dict.
    If None or a single season, fetches without seasonId filter.
    """
    def _fetch_season(sid: str | None) -> list:
        eps = []
        skip = 0
        while True:
            extra = {'count': '20', 'skip': str(skip)}
            if sid:
                extra['seasonId'] = sid
            params = _api_params(country, **extra)
            r = session.get(
                f'{TV_API}/v3/shows/{show_id}/episodes',
                params=params,
                timeout=20,
            )
            _check_api(r)
            data = r.json().get('data', {})
            items = data.get('episodes', [])
            if not items:
                break
            eps.extend(items)
            if len(items) < 20:
                break
            skip += 20
        return eps

    if not season_ids or len(season_ids) <= 1:
        return _fetch_season(season_ids[0] if season_ids else None)

    all_episodes = []
    for sid in season_ids:
        all_episodes.extend(_fetch_season(sid))
    return all_episodes


def get_playback_info(session: requests.Session, content_id: str, country: str) -> dict:
    """Get subscriber playback data (requires media-user-token).

    v3/episodes and v3/movies are the current endpoints; v2/view/product/*/personalized
    returns 404 as of 2026.
    """
    endpoints = [
        f'{TV_API}/v3/episodes/{content_id}',
        f'{TV_API}/v3/movies/{content_id}',
    ]
    for url in endpoints:
        r = session.get(url, params=_api_params(country), timeout=15)
        if r.status_code in (401, 403):
            raise RuntimeError(
                f'Playback info denied (HTTP {r.status_code}). '
                'media-user-token may be missing or expired — '
                'provide via --media-token or APPLE_MUT env var.'
            )
        if r.ok:
            data = r.json().get('data', {})
            if _has_playback_data(data):
                return data
        elif r.status_code != 404:
            _check_api(r)
    raise RuntimeError(
        f'No playback endpoint returned usable data for {content_id}. '
        'Content may not be available in your region or the ID format changed.'
    )


def _has_playback_data(data: dict) -> bool:
    """Return True if the API response contains a playables section with entitlement info."""
    if not data:
        return False
    pl = data.get('playables', {})
    if isinstance(pl, dict) and pl:
        first = next(iter(pl.values()), {})
        return 'isEntitledToPlay' in first or 'hlsUrl' in first
    return False


def _extract_hls_url(playback: dict) -> str:
    """Dig through the response to find the HLS URL.

    v3/episodes|movies: 'playables' is a dict keyed by playable_id;
    each value has an 'assets' sub-dict containing 'hlsUrl' when entitled.
    """
    pl = playback.get('playables', {})
    if isinstance(pl, dict):
        for pid, pval in pl.items():
            if not isinstance(pval, dict):
                continue
            # Primary location: assets.hlsUrl (v3 entitled response)
            url = pval.get('assets', {}).get('hlsUrl', '')
            if url and url.startswith('http'):
                return url
            # Fallback: hlsUrl directly on playable
            url = pval.get('hlsUrl', '')
            if url and url.startswith('http'):
                return url
            # Give a clear entitlement error if we know the reason
            if not pval.get('isEntitledToPlay', True):
                raise RuntimeError(
                    'Not entitled to play this content (isEntitledToPlay=false). '
                    'Ensure media-user-token is set via --media-token or APPLE_MUT env var '
                    'and that you have an active Apple TV+ subscription.'
                )

    raise RuntimeError(
        f'No hlsUrl found in playback response. '
        f'Keys present: {list(playback.keys())}\n'
        f'Raw snippet: {json.dumps(playback)[:500]}'
    )


def _extract_play_params(playback: dict) -> dict:
    """Extract adamId, svcId, and licenseUrl needed for license requests.

    v3 responses: 'playables[pid].assets' contains assetAdamId,
    fpsKeyServerUrl, and fpsKeyServerQueryParameters.svcId.
    """
    pl = playback.get('playables', {})
    if isinstance(pl, dict):
        for pid, pval in pl.items():
            if not isinstance(pval, dict):
                continue
            assets = pval.get('assets', {})
            fps_q  = assets.get('fpsKeyServerQueryParameters', {})
            adam   = (assets.get('assetAdamId', '')
                      or fps_q.get('adamId', '')
                      or pval.get('externalId', ''))
            svc    = (fps_q.get('svcId', '')
                      or pval.get('channelId', 'tvs.sbd.4000'))
            lic    = (assets.get('fpsKeyServerUrl', '')
                      or LICENSE_URL)
            if adam:
                return {
                    'adamId': adam,
                    'svcId': svc,
                    'licenseUrl': lic,
                    'isExternal': fps_q.get('isExternal', True),
                }
    # Fallback: v2 personalized shape
    try:
        items = (playback.get('content', {})
                         .get('playables', {})
                         .get('items', [{}]))
        pp = items[0].get('playParams', {})
        return {
            'adamId': pp.get('id', pp.get('catalogId', '')),
            'svcId':  pp.get('channelId', 'tvs.sbd.4000'),
            'licenseUrl': LICENSE_URL,
            'isExternal': True,
        }
    except Exception:
        return {'adamId': '', 'svcId': 'tvs.sbd.4000', 'licenseUrl': LICENSE_URL,
                'isExternal': True}


# ── HLS / DRM ────────────────────────────────────────────────────────────────

def fetch_m3u8(url: str, session: requests.Session) -> str:
    r = session.get(url, headers={'User-Agent': UA}, timeout=15)
    _check_api(r)
    return r.text


def _parse_hls_attrs(attr_str: str) -> dict:
    """Parse HLS attribute list: KEY=VALUE,KEY="VALUE",..."""
    result = {}
    for m in re.finditer(r'([\w-]+)=("(?:[^"\\]|\\.)*"|[^,]+)', attr_str):
        result[m.group(1)] = m.group(2).strip('"')
    return result


def _pssh_from_m3u8_text(text: str) -> str:
    """Scan M3U8 text for a Widevine EXT-X-KEY/EXT-X-SESSION-KEY; returns full data: URI.

    Apple TV+ uses: data:text/plain;keyInfo=<token>;base64,<pssh>
    The full URI (including keyInfo) must be forwarded to the license server.
    """
    for line in text.splitlines():
        if not (line.startswith('#EXT-X-KEY') or line.startswith('#EXT-X-SESSION-KEY')):
            continue
        if WIDEVINE_KEYFORMAT not in line:
            continue
        attrs = _parse_hls_attrs(line.split(':', 1)[1])
        uri = attrs.get('URI', '')
        if 'base64,' in uri:
            return uri  # return full URI, keyInfo included
    return ''


def get_pssh_from_m3u8(m3u8_text: str, session: requests.Session | None = None,
                        base_url: str = '') -> list[str]:
    """Extract Widevine PSSH data: URIs from HLS manifests.

    Returns a list of unique PSSH URIs (full data: URIs with keyInfo preserved).
    Apple TV+ uses per-codec-group KIDs: dvh1 (Dolby Vision), hvc1 (HEVC), and
    avc1 (H.264) streams each carry a different KID.  We sample one variant per
    codec family so that all KIDs are covered and every selected stream can be
    decrypted.

    base_url: URL the m3u8_text was fetched from, needed to resolve relative
              variant URLs.
    """
    # Check master first (uncommon but possible)
    master_uri = _pssh_from_m3u8_text(m3u8_text)
    if master_uri:
        return [master_uri]

    if session is None:
        raise RuntimeError('No Widevine PSSH found in HLS master manifest.')

    base_dir = base_url.split('?')[0].rsplit('/', 1)[0] + '/' if base_url else ''
    lines = m3u8_text.splitlines()

    # Collect the first variant URL for each top-level codec family.
    # Apple uses the same KID within a family (all dvh1 bitrates share one key,
    # all hvc1 share another, etc.), so one sample per family is enough.
    codec_seen: set[str] = set()
    variant_urls: list[str] = []
    i = 0
    while i < len(lines):
        line = lines[i]
        if line.startswith('#EXT-X-STREAM-INF:'):
            attrs = _parse_hls_attrs(line[len('#EXT-X-STREAM-INF:'):])
            codecs = attrs.get('CODECS', '')
            # Codec family = first 3 chars of the first codec (dvh, hvc, avc, …)
            family = codecs.split('.')[0][:3].lower() if codecs else 'unk'
            if family not in codec_seen:
                codec_seen.add(family)
                for j in range(i + 1, len(lines)):
                    candidate = lines[j].strip()
                    if candidate and not candidate.startswith('#'):
                        var_url = (candidate if candidate.startswith('http')
                                   else base_dir + candidate)
                        variant_urls.append(var_url)
                        break
        i += 1

    # Fetch sampled variants, collect unique PSSH URIs (deduplicated by b64 payload)
    seen_b64: set[str] = set()
    pssh_uris: list[str] = []
    for var_url in variant_urls:
        try:
            var_r = session.get(var_url, headers={'User-Agent': UA}, timeout=15)
            if var_r.ok:
                uri = _pssh_from_m3u8_text(var_r.text)
                if uri:
                    b64 = uri.split('base64,', 1)[1] if 'base64,' in uri else uri
                    if b64 not in seen_b64:
                        seen_b64.add(b64)
                        pssh_uris.append(uri)
        except Exception:
            pass

    if not pssh_uris:
        raise RuntimeError('No Widevine PSSH found in HLS manifest or any variant playlist.')
    return pssh_uris


def parse_m3u8_tracks(m3u8_text: str) -> tuple:
    """Parse video streams, audio tracks, and subtitle tracks from HLS master.

    Returns (video_streams, audio_tracks, sub_tracks).
    """
    video_streams = []
    audio_tracks  = []
    sub_tracks    = []

    lines = m3u8_text.splitlines()
    for i, line in enumerate(lines):
        if line.startswith('#EXT-X-STREAM-INF:'):
            attrs = _parse_hls_attrs(line[len('#EXT-X-STREAM-INF:'):])
            bw  = int(attrs.get('BANDWIDTH', 0)) // 1000
            res = attrs.get('RESOLUTION', '')
            codecs = attrs.get('CODECS', '')
            if res and 'x' in res:
                w, h = map(int, res.split('x'))
                codec_tag = codecs.split(',')[0] if codecs else ''
                video_streams.append({
                    'width': w, 'height': h, 'kbps': bw, 'codec': codec_tag
                })

        elif line.startswith('#EXT-X-MEDIA:'):
            attrs = _parse_hls_attrs(line[len('#EXT-X-MEDIA:'):])
            typ  = attrs.get('TYPE', '')
            lang = attrs.get('LANGUAGE', 'und')
            name = attrs.get('NAME', lang)
            uri  = attrs.get('URI', '')
            grp  = attrs.get('GROUP-ID', '')
            if typ == 'AUDIO' and uri:
                audio_tracks.append({
                    'lang': lang,
                    'lang_disp': lang_name(lang),
                    'name': name,
                    'group_id': grp,
                    'default': attrs.get('DEFAULT', 'NO') == 'YES',
                })
            elif typ == 'SUBTITLES':
                forced = attrs.get('FORCED', 'NO') == 'YES'
                sub_tracks.append({
                    'lang': lang,
                    'lang_disp': lang_name(lang),
                    'name': name,
                    'forced': forced,
                    'uri': uri,
                })

    # Deduplicate video by resolution, keep highest bitrate per resolution
    best: dict = {}
    for v in video_streams:
        key = (v['width'], v['height'])
        if key not in best or v['kbps'] > best[key]['kbps']:
            best[key] = v
    video_streams = sorted(best.values(), key=lambda x: (-x['width'], -x['kbps']))

    return video_streams, audio_tracks, sub_tracks


# ── Widevine license via Apple's FPS endpoint ─────────────────────────────────

def fetch_content_key(pssh_uris: list[str] | str, play_params: dict,
                      session: requests.Session) -> dict:
    """Fetch Widevine content keys via Apple's license server.

    pssh_uris: one or more full data: URIs from HLS EXT-X-KEY lines
               (data:text/plain;keyInfo=<tok>;base64,<pssh>).  Apple TV+ uses
               per-codec-group KIDs, so callers should pass one URI per codec
               family (dvh1, hvc1, avc1) to obtain all needed keys.
    """
    from pywidevine.cdm import Cdm
    from pywidevine.device import Device
    from pywidevine.pssh import PSSH
    import uuid as _uuid

    if isinstance(pssh_uris, str):
        pssh_uris = [pssh_uris]

    lic_url = play_params.get('licenseUrl', LICENSE_URL)
    lic_headers = {
        'User-Agent': UA,
        'Content-Type': 'application/json',
        'Accept': 'application/json',
    }
    mut = session.cookies.get('media-user-token', '') or session.headers.get('media-user-token', '')
    if mut:
        lic_headers['media-user-token'] = mut
    myacinfo = session.cookies.get('myacinfo', '')
    if myacinfo:
        lic_headers['Cookie'] = f'myacinfo={myacinfo}; media-user-token={mut}'

    device = Device.load(str(WVD))
    all_keys: dict = {}

    for pssh_uri in pssh_uris:
        cdm        = Cdm.from_device(device)
        session_id = cdm.open()
        try:
            pssh_b64 = pssh_uri.split('base64,', 1)[1] if 'base64,' in pssh_uri else pssh_uri
            challenge = cdm.get_license_challenge(session_id, PSSH(pssh_b64))
            challenge_b64 = base64.b64encode(challenge).decode()

            body = {
                'streaming-request': {
                    'version': 1,
                    'streaming-keys': [{
                        'challenge':    challenge_b64,
                        'key-system':   'com.widevine.alpha',
                        'uri':          pssh_uri,
                        'id':           1,
                        'lease-action': 'start',
                        'adamId':       play_params.get('adamId', ''),
                        'isExternal':   play_params.get('isExternal', True),
                        'svcId':        play_params.get('svcId', 'tvs.sbd.4000'),
                    }],
                }
            }

            r = session.post(lic_url, json=body, headers=lic_headers, timeout=20)
            if not r.ok:
                raise RuntimeError(f'License server {r.status_code}: {r.text[:300]}')

            resp = r.json()
            keys_resp = (resp.get('streaming-response', {})
                             .get('streaming-keys', [{}]))
            license_b64 = keys_resp[0].get('license', '')
            if not license_b64:
                raise RuntimeError(f'No license in response: {resp}')

            cdm.parse_license(session_id, base64.b64decode(license_b64))
            for key in cdm.get_keys(session_id):
                if str(key.type) not in ('SIGNING', 'OPERATOR_SESSION'):
                    kid = key.kid.hex if isinstance(key.kid, _uuid.UUID) else key.kid.hex()
                    all_keys[kid] = key.key.hex()
        finally:
            cdm.close(session_id)

    return all_keys


# ── Interactive pickers (reused from sky_download style) ─────────────────────

def checkbox_picker(title, items, default_all=True, single_select=False, default_idx=None):
    if not items:
        return []
    if default_idx is not None:
        selected = [i == default_idx for i in range(len(items))]
    else:
        selected = [default_all] * len(items)
    state = {'cursor': 0, 'offset': 0}

    def _run(stdscr):
        curses.curs_set(0)
        try:
            curses.start_color(); curses.use_default_colors()
            curses.init_pair(1, curses.COLOR_BLACK, curses.COLOR_CYAN)
        except Exception:
            pass
        while True:
            stdscr.erase()
            h, w = stdscr.getmaxyx()
            list_rows = max(1, h - 5)
            cursor = state['cursor']; offset = state['offset']
            if cursor < offset: offset = cursor
            elif cursor >= offset + list_rows: offset = cursor - list_rows + 1
            state['offset'] = offset
            try:
                stdscr.addstr(0, 0, title[:w-1], curses.A_BOLD)
                stdscr.addstr(1, 0, ('─' * min(len(title), w-1))[:w-1])
                hint = ('↑↓ move   SPACE select   ENTER=confirm' if single_select
                        else '↑↓ move   SPACE toggle   a=all   n=none   ENTER=confirm')
                stdscr.addstr(2, 0, hint[:w-1])
            except curses.error:
                pass
            for i in range(min(list_rows, len(items))):
                idx = i + offset
                if idx >= len(items): break
                mark = ('(◉)' if selected[idx] else '( )') if single_select else ('[✓]' if selected[idx] else '[ ]')
                line = f'  {mark}  {items[idx]}'[:w-1]
                try:
                    stdscr.addstr(i + 4, 0, line,
                                  curses.color_pair(1) if idx == cursor else 0)
                except curses.error:
                    pass
            if not single_select:
                try:
                    stdscr.addstr(h-1, 0, f'  {sum(selected)} of {len(items)} selected'[:w-1], curses.A_DIM)
                except curses.error:
                    pass
            stdscr.refresh()
            key = stdscr.getch()
            if key in (curses.KEY_UP, ord('k')): state['cursor'] = max(0, cursor-1)
            elif key in (curses.KEY_DOWN, ord('j')): state['cursor'] = min(len(items)-1, cursor+1)
            elif key == ord(' '):
                if single_select:
                    for j in range(len(selected)): selected[j] = (j == cursor)
                else:
                    selected[cursor] = not selected[cursor]
            elif not single_select and key == ord('a'):
                for j in range(len(selected)): selected[j] = True
            elif not single_select and key == ord('n'):
                for j in range(len(selected)): selected[j] = False
            elif key in (curses.KEY_ENTER, 10, 13):
                break

    try:
        curses.wrapper(_run)
    except Exception as e:
        print(f'[warning] Interactive picker failed ({e}), using default.')
        if single_select:
            return [default_idx if default_idx is not None else 0]
        return list(range(len(items)))
    return [i for i, s in enumerate(selected) if s]


def select_video_interactive(video_streams):
    if not video_streams:
        return ['--select-video', 'best']
    items = []
    for v in video_streams:
        c = v.get('codec', '')
        if 'hvc' in c or 'hev' in c or 'dvh' in c:
            cn = 'H.265'
        elif 'avc' in c:
            cn = 'H.264'
        else:
            cn = c.split('.')[0] if c else '?'
        items.append(f'{v["width"]}×{v["height"]:<9} {v["kbps"]:>5} Kbps  {cn}')
    print()
    sel = checkbox_picker('Select video quality', items, single_select=True, default_idx=0)
    if not sel:
        return ['--select-video', 'best']
    v = video_streams[sel[0]]
    print(f'Video: {v["width"]}×{v["height"]} {v["kbps"]} Kbps')
    return ['--select-video', f'res={v["width"]}x{v["height"]}:for=best']


def select_audio_interactive(audio_tracks):
    if not audio_tracks:
        return ['--drop-audio', 'all']
    seen = set()
    langs = []
    for t in audio_tracks:
        if t['lang'] not in seen:
            seen.add(t['lang'])
            langs.append(t['lang'])
    items = [f'{l:<9} {lang_name(l)}' for l in langs]
    print()
    sel = checkbox_picker('Select audio languages', items, default_all=False)
    if not sel:
        return ['--drop-audio', 'all']
    sel_langs = [langs[i] for i in sel]
    print(f'Audio: {", ".join(sel_langs)}')
    return ['--select-audio', f'lang={"|".join(sel_langs)}:for=best']


def select_subs_interactive(sub_tracks):
    if not sub_tracks:
        return ['--drop-subtitle', 'all']
    seen = set(); langs = []
    for t in sub_tracks:
        if t['lang'] not in seen:
            seen.add(t['lang']); langs.append(t['lang'])
    items = []
    for l in langs:
        forced_only = all(t['forced'] for t in sub_tracks if t['lang'] == l)
        note = '  (forced only)' if forced_only else ''
        items.append(f'{l:<9} {lang_name(l)}{note}')
    print()
    sel = checkbox_picker('Select subtitle tracks', items, default_all=False)
    if not sel:
        return ['--drop-subtitle', 'all']
    sel_langs = [langs[i] for i in sel]
    print(f'Subtitles: {", ".join(sel_langs)}')
    return ['--select-subtitle', f'lang={"|".join(sel_langs)}:for=best']


# ── Download + decrypt ────────────────────────────────────────────────────────

def safe_filename(title: str) -> str:
    title = re.sub(r'[^\w\s-]', '', title)
    return re.sub(r'[\s-]+', '_', title).strip('_')


def download_and_decrypt(hls_url: str, keys: dict, out_name: str,
                         video_args: list, audio_args: list,
                         cookie_header: str = '', sub_args: list | None = None):
    if sub_args is None:
        sub_args = ['--drop-subtitle', 'all']

    key_args = []
    for kid, key in keys.items():
        key_args += ['--key', f'{kid}:{key}']

    extra_headers = [f'User-Agent: {UA}']
    if cookie_header:
        extra_headers.append(f'Cookie: {cookie_header}')

    cmd = [
        str(N_M3U8), hls_url,
        '--save-dir', str(SAVE_DIR),
        '--save-name', out_name,
        '--tmp-dir', TMP_DIR,
        '--decryption-binary-path', str(MP4DEC),
        '--mux-after-done', 'format=mkv',
        '--concurrent-downloads', '16',
    ]
    for h in extra_headers:
        cmd += ['--header', h]
    cmd += video_args + audio_args + sub_args + key_args

    print(f'N_m3u8DL-RE: {" ".join(video_args + audio_args + sub_args)}')
    print('Downloading and decrypting...')
    result = subprocess.run(cmd)
    if result.returncode != 0:
        raise RuntimeError(f'N_m3u8DL-RE failed (exit {result.returncode})')


# ── Episode helpers ───────────────────────────────────────────────────────────

def _parse_episode(item: dict) -> dict | None:
    """Normalise an episode item from the API into a flat dict."""
    # The API may return episodes in different shapes
    cid  = item.get('id', '')
    title = item.get('title', item.get('name', ''))
    ep_num = item.get('episodeNumber', item.get('number', 0))
    sn_num = item.get('seasonNumber', 1)

    # Some responses nest under 'attributes'
    attrs = item.get('attributes', {})
    if attrs:
        cid    = cid or attrs.get('id', '')
        title  = title or attrs.get('title', attrs.get('name', ''))
        ep_num = ep_num or attrs.get('episodeNumber', attrs.get('number', 0))
        sn_num = sn_num or attrs.get('seasonNumber', 1)

    if not cid or not title:
        return None

    try: ep_num = int(ep_num)
    except (TypeError, ValueError): ep_num = 0
    try: sn_num = int(sn_num)
    except (TypeError, ValueError): sn_num = 1

    return {
        'id': cid,
        'title': title,
        'season': sn_num,
        'episode': ep_num,
    }


def select_episode_interactive(show_title: str, episodes: list) -> list:
    items = ['[All episodes]'] + [
        f'S{e["season"]:02d}E{e["episode"]:02d}  {e["title"]}' for e in episodes
    ]
    print(f'\nShow: {show_title}  ({len(episodes)} episodes)')
    sel = checkbox_picker('Select episode to download', items,
                          single_select=True, default_idx=0)
    idx = sel[0] if sel else 0
    if idx == 0:
        print('Downloading all episodes.')
        return list(episodes)
    ep = episodes[idx - 1]
    print(f'Downloading S{ep["season"]:02d}E{ep["episode"]:02d}: {ep["title"]}')
    return [ep]


def episode_out_name(ep: dict) -> str:
    return f's{ep["season"]:02d}e{ep["episode"]:02d}_{safe_filename(ep["title"])}'


def download_episode(session, ep: dict, country: str,
                     video_args: list, audio_args: list,
                     cookie_header: str, sub_args: list | None = None,
                     keys_only: bool = False):
    print(f'\n── S{ep["season"]:02d}E{ep["episode"]:02d}: {ep["title"]} ──')
    pb = get_playback_info(session, ep['id'], country)
    hls_url = _extract_hls_url(pb)
    play_params = _extract_play_params(pb)

    m3u8_text = fetch_m3u8(hls_url, session)
    pssh_uris = get_pssh_from_m3u8(m3u8_text, session, base_url=hls_url)

    print('Fetching content key...')
    keys = fetch_content_key(pssh_uris, play_params, session)
    for kid, key in keys.items():
        print(f'  {kid}:{key}')

    if keys_only:
        return

    out_name = episode_out_name(ep)
    download_and_decrypt(hls_url, keys, out_name, video_args, audio_args, cookie_header,
                         sub_args=sub_args)
    print(f'Saved: {SAVE_DIR / out_name}.mkv')


# ── CLI ───────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description='Download a DRM-protected video from Apple TV+',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python3 apple_download.py https://tv.apple.com/pl/show/star-city/umc.cmc.2l8p785osmtmiyk64bh6tfde1
  python3 apple_download.py -a en -s pl https://tv.apple.com/pl/movie/f1-the-movie/umc.cmc.3t6dvnnr87zwd4wmvpdx5came
  python3 apple_download.py -u me@icloud.com https://tv.apple.com/...

Session is cached in .apple_cookies + .apple_session after first login.
Delete those files to force re-login.
""",
    )
    parser.add_argument('url', help='https://tv.apple.com/... URL')
    parser.add_argument('-u', '--username', default=os.environ.get('APPLE_USER'))
    parser.add_argument('-p', '--password', default=os.environ.get('APPLE_PASS'))
    parser.add_argument('-t', '--media-token', default=os.environ.get('APPLE_MUT'),
                        metavar='TOKEN',
                        help='media-user-token (from browser DevTools → Cookies → .apple.com)')
    parser.add_argument('-o', '--output', help='Output filename without extension')
    parser.add_argument('-v', '--video', default=None, metavar='RES',
                        help='Video resolution e.g. 1920x1080, "best", "worst"')
    parser.add_argument('-a', '--audio', default=None, metavar='LANGS',
                        help='Audio: comma-separated lang codes, "all", or "none"')
    parser.add_argument('-s', '--subs', default=None, metavar='LANGS',
                        help='Subtitles: comma-separated lang codes, "all", or "none"')
    parser.add_argument('--debug', action='store_true',
                        help='Print raw API responses for debugging')
    parser.add_argument('--keys-only', action='store_true',
                        help='Fetch and print Widevine keys then exit without downloading')
    args = parser.parse_args()

    if not WVD:
        sys.exit('ERROR: No valid .wvd file found in L3/.')

    # Parse URL
    try:
        country, kind, cid = _parse_url(args.url)
    except ValueError as e:
        sys.exit(f'ERROR: {e}')
    print(f'Country: {country.upper()}  Type: {kind}  ID: {cid}')

    session = requests.Session()
    session.headers.update({'User-Agent': UA})

    # ── Session restore ───────────────────────────────────────────────────────
    meta = load_session(session)

    if is_session_valid(session, meta):
        print('Using cached session.')
    else:
        # Need to login
        if not args.username:
            args.username = input('Apple ID (email): ').strip()
        if not args.password:
            args.password = getpass.getpass('Apple ID password: ')

        print(f'Logging in as {args.username}...')
        try:
            token_data = apple_login(session, args.username, args.password)
        except RuntimeError as e:
            sys.exit(f'ERROR: Login failed: {e}')

        meta.update(token_data)
        meta['username'] = args.username
        save_session(session, meta)
        print('Session saved.')

    # ── media-user-token — set before fetching TV page so utsk is authenticated ─
    mut = (args.media_token
           or meta.get('media_user_token', '')
           or session.cookies.get('media-user-token', ''))
    if not mut:
        # We need dev_token to drive the OAuth flow; fetch the TV page first (anonymous).
        print('Fetching Apple TV page tokens (anonymous)...')
        try:
            dev_token, _ = _get_tv_page_config(session, country)
            session.headers['Authorization'] = f'Bearer {dev_token}'
        except RuntimeError as e:
            sys.exit(f'ERROR: {e}')

        print('Obtaining media-user-token...')
        try:
            mut = get_media_user_token(
                session, dev_token,
                username=args.username or meta.get('username', ''),
                password=args.password or '',
            )
            print('media-user-token obtained.')
        except RuntimeError as e:
            print(f'[warn] Could not obtain media-user-token: {e}')

    if mut:
        session.cookies.set('media-user-token', mut, domain='.apple.com')
        session.headers['media-user-token'] = mut
        meta['media_user_token'] = mut
        save_session(session, meta)

    # Fetch TV page with media-user-token set so SSR returns authenticated utsk
    print('Fetching Apple TV page tokens...')
    try:
        dev_token, utsk = _get_tv_page_config(session, country)
        session.headers['Authorization'] = f'Bearer {dev_token}'
        print(f'Developer token: {dev_token[:40]}...')
        if utsk:
            anon = '::::::' in utsk
            print(f'utsk: {utsk[:32]}...  ({"anonymous" if anon else "authenticated"})')
        else:
            print('[warn] utsk not found in TV page — API calls may fail')
    except RuntimeError as e:
        sys.exit(f'ERROR: {e}')

    use_interactive = (args.video is None and args.audio is None
                       and args.subs is None and sys.stdin.isatty())

    # Cookie header string for N_m3u8DL-RE
    cookie_header = '; '.join(f'{c.name}={c.value}' for c in session.cookies)

    SAVE_DIR.mkdir(exist_ok=True)

    if kind == 'show':
        # ── Series flow ───────────────────────────────────────────────────────
        print(f'Fetching show metadata for {cid}...')
        show_data = get_show_metadata(session, cid, country)
        if args.debug:
            print('SHOW DATA:', json.dumps(show_data, indent=2)[:2000])

        content = show_data.get('content', {})
        show_title = content.get('title') or show_data.get('title', cid)
        # Update cid in case the URL pointed to an episode that was resolved to its show
        cid = content.get('id', cid)
        season_ids = list(show_data.get('seasons', {}).keys())
        print(f'Show: {show_title}  ({len(season_ids)} season(s))')

        print('Fetching episode list...')
        ep_raw = get_episodes(session, cid, country, season_ids)
        if args.debug:
            print('EPISODES RAW:', json.dumps(ep_raw[:2], indent=2))

        episodes = [e for e in (_parse_episode(r) for r in ep_raw) if e]
        if not episodes:
            # Sometimes episodes are nested differently
            for item in ep_raw:
                for sub in item.get('episodes', item.get('items', [])):
                    ep = _parse_episode(sub)
                    if ep:
                        episodes.append(ep)

        episodes.sort(key=lambda e: (e['season'], e['episode']))

        if not episodes:
            sys.exit(f'ERROR: No episodes found. Use --debug to inspect the API response.')

        if sys.stdin.isatty():
            selected_eps = select_episode_interactive(show_title, episodes)
        else:
            selected_eps = episodes

        # Fetch first episode's manifest to offer track selection
        first = selected_eps[0]
        print(f'\nFetching track list from S{first["season"]:02d}E{first["episode"]:02d}...')
        pb0 = get_playback_info(session, first['id'], country)
        hls0 = _extract_hls_url(pb0)
        m3u8_0 = fetch_m3u8(hls0, session)
        video_streams, audio_tracks, sub_tracks = parse_m3u8_tracks(m3u8_0)

        if use_interactive:
            video_args = select_video_interactive(video_streams)
            audio_args = select_audio_interactive(audio_tracks)
            sub_args   = select_subs_interactive(sub_tracks)
        else:
            v = args.video
            if v and v not in ('best', 'worst'):
                video_args = ['--select-video', f'res={v}:for=best']
            else:
                video_args = ['--select-video', v or 'best']
            a = args.audio or 'all'
            if a == 'all':
                audio_args = ['--select-audio', 'all']
            elif a == 'none':
                audio_args = ['--drop-audio', 'all']
            else:
                langs = '|'.join(x.strip() for x in a.split(','))
                audio_args = ['--select-audio', f'lang={langs}:for=best']
            s = args.subs
            if not s:
                sub_args = ['--drop-subtitle', 'all']
            elif s == 'all':
                sub_args = ['--select-subtitle', 'all']
            else:
                sub_langs = '|'.join(x.strip() for x in s.split(','))
                sub_args = ['--select-subtitle', f'lang={sub_langs}:for=best']

        print(f'\nDownloading {len(selected_eps)} episode(s)...')
        for ep in selected_eps:
            try:
                download_episode(session, ep, country, video_args, audio_args, cookie_header,
                                 sub_args=sub_args, keys_only=args.keys_only)
            except Exception as e:
                print(f'  ERROR on {ep["title"]}: {e}')
        print(f'\nAll done. Files in {SAVE_DIR}')

    else:
        # ── Movie flow ────────────────────────────────────────────────────────
        print(f'Fetching movie metadata for {cid}...')
        movie_data = get_movie_metadata(session, cid, country)
        if args.debug:
            print('MOVIE DATA:', json.dumps(movie_data, indent=2)[:2000])

        movie_title = (movie_data.get('content', {}).get('title')
                       or movie_data.get('title', cid))
        print(f'Movie: {movie_title}')

        print('Getting playback info...')
        pb = get_playback_info(session, cid, country)
        if args.debug:
            print('PLAYBACK:', json.dumps(pb, indent=2)[:2000])

        hls_url = _extract_hls_url(pb)
        play_params = _extract_play_params(pb)
        print(f'HLS: {hls_url[:80]}...')
        print(f'adamId={play_params["adamId"]}  svcId={play_params["svcId"]}')

        print('Fetching manifest...')
        m3u8_text = fetch_m3u8(hls_url, session)

        print('Extracting Widevine PSSH...')
        pssh_uris = get_pssh_from_m3u8(m3u8_text, session, base_url=hls_url)
        print(f'PSSH count: {len(pssh_uris)} codec group(s)')

        video_streams, audio_tracks, sub_tracks = parse_m3u8_tracks(m3u8_text)

        if use_interactive:
            video_args = select_video_interactive(video_streams)
            audio_args = select_audio_interactive(audio_tracks)
            sub_args   = select_subs_interactive(sub_tracks)
        else:
            v = args.video
            if v and v not in ('best', 'worst'):
                video_args = ['--select-video', f'res={v}:for=best']
            else:
                video_args = ['--select-video', v or 'best']
            a = args.audio or 'all'
            if a == 'all':
                audio_args = ['--select-audio', 'all']
            elif a == 'none':
                audio_args = ['--drop-audio', 'all']
            else:
                langs = '|'.join(x.strip() for x in a.split(','))
                audio_args = ['--select-audio', f'lang={langs}:for=best']
            s = args.subs
            if not s:
                sub_args = ['--drop-subtitle', 'all']
            elif s == 'all':
                sub_args = ['--select-subtitle', 'all']
            else:
                sub_langs = '|'.join(x.strip() for x in s.split(','))
                sub_args = ['--select-subtitle', f'lang={sub_langs}:for=best']

        print('Fetching content key...')
        keys = fetch_content_key(pssh_uris, play_params, session)
        if not keys:
            raise RuntimeError('License server returned no content keys.')
        for kid, key in keys.items():
            print(f'  {kid}:{key}')

        if args.keys_only:
            return

        out_name = args.output or safe_filename(movie_title)
        download_and_decrypt(hls_url, keys, out_name, video_args, audio_args, cookie_header,
                             sub_args=sub_args)
        print(f'Done: {SAVE_DIR / out_name}.mkv')


if __name__ == '__main__':
    main()