#!/usr/bin/python3
import sys
import os
import o11
import base64
import json
import datetime
import pytz
import urllib.parse
import jwt
from pywidevine.cdm import Cdm
from pywidevine.device import Device
from pywidevine.pssh import PSSH
from bs4 import BeautifulSoup

# ─── hardcoded config ────────────────────────────────────────────────────────
WVD_PATH   = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'WVD.wvd')
USER_AGENT = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/139.0.0.0 Safari/537.36'
DEVICE_INFO = 'device=web;screen=browser;os=windows;browser=chrome;browserVersion=139.0.0.0;osVersion=10;appVersion=release-R45.0.2;playerVersion=8.212.0'
F1TV_BASE  = 'https://f1tv.formula1.com'
# ─────────────────────────────────────────────────────────────────────────────

user      = o11.parse_params(sys.argv, 'user')
password  = o11.parse_params(sys.argv, 'password')

id        = o11.parse_params(sys.argv, 'id')
action    = o11.parse_params(sys.argv, 'action')

bind      = o11.parse_params(sys.argv, 'bind')
proxy     = o11.parse_params(sys.argv, 'proxy')
doh       = o11.parse_params(sys.argv, 'doh')
worker    = o11.parse_params(sys.argv, 'worker')

cdm_param = o11.parse_params(sys.argv, 'cdm')
drm       = o11.parse_params(sys.argv, 'drm')
kid       = o11.parse_params(sys.argv, 'kid')
pssh      = o11.parse_params(sys.argv, 'pssh')
challenge = o11.parse_params(sys.argv, 'challenge')

heartbeaturl    = o11.parse_params(sys.argv, 'heartbeaturl')
heartbeatparams = o11.parse_params(sys.argv, 'heartbeatparams')

o11Session = o11.session(bind=bind, proxy=proxy, worker=worker)
req = o11Session.get_session()
if doh != '':
    o11.dns(doh)

if challenge == 'cert':
    challenge = 'CAQ='

authFile = '/f1tv_' + user + '.tokens'


# ─── auth helpers ─────────────────────────────────────────────────────────────

def _load_auth():
    try:
        return json.load(open(os.path.abspath(os.path.dirname(__file__)) + authFile))
    except Exception:
        return None


def _save_auth(data):
    json.dump(data, open(os.path.abspath(os.path.dirname(__file__)) + authFile, 'w'))


def _f1_headers(token=None, ascendon=None):
    h = {
        'user-agent':      USER_AGENT,
        'x-f1-device-info': DEVICE_INFO,
    }
    if token:
        h['entitlementtoken'] = token
    if ascendon:
        h['ascendontoken'] = ascendon
    return h


def _extract_ascendon(raw):
    """
    The ascendon token can arrive in two forms:
    1. A raw JWT string
    2. A URL-encoded JSON blob (from the login-session cookie) containing
       data.subscriptionToken
    Returns the bare JWT string in either case.
    """
    # try raw JWT first
    try:
        jwt.decode(raw, options={'verify_signature': False})
        return raw
    except Exception:
        pass
    # try URL-encoded JSON
    try:
        decoded = urllib.parse.unquote(raw)
        data    = json.loads(decoded)
        return data['data']['subscriptionToken']
    except Exception:
        pass
    return raw


# ─── API calls ────────────────────────────────────────────────────────────────

def _get_entitlement(ascendon_token):
    r = req.get(
        F1TV_BASE + '/2.0/R/ENG/WEB_DASH/ALL/USER/ENTITLEMENT',
        headers=_f1_headers(ascendon=ascendon_token),
    )
    token = r.json().get('resultObj', {}).get('entitlementToken', '')
    if not token:
        print('Failed to get entitlement token', file=sys.stderr)
        print(r.text, file=sys.stderr)
    return token


def _get_events():
    r = req.get(
        F1TV_BASE + '/2.0/R/ENG/WEB_DASH/ALL/PAGE/395/REG/14',
        headers=_f1_headers(),
    )
    items = []
    for c in r.json().get('resultObj', {}).get('containers', []):
        ri = c.get('retrieveItems', {}).get('resultObj', {})
        for c2 in ri.get('containers', []):
            for e in c2.get('events', []):
                meta = e.get('metadata', {})
                if meta.get('contentSubtype') == 'LIVE':
                    items.append(meta)
    # deduplicate by contentId
    return list({item['contentId']: item for item in items}.values())


def _get_single(token, content_id):
    r = req.get(
        F1TV_BASE + '/2.0/R/ENG/WEB_HLS/ALL/CONTENT/PLAY',
        params={'contentId': content_id},
        headers=_f1_headers(token=token),
    )
    obj = r.json().get('resultObj', {})
    url       = obj.get('url', '')
    lic_url   = obj.get('laURL', '')
    lic_token = obj.get('entitlementToken', '')
    if not url:
        print('get_single failed', file=sys.stderr)
        print(r.text, file=sys.stderr)
    return url, lic_url, lic_token


def _get_pssh_from_mpd(mpd_url):
    r = req.get(mpd_url, headers={'User-Agent': USER_AGENT})
    try:
        for cp in BeautifulSoup(r.content, features='xml').find_all('ContentProtection'):
            if cp.get('schemeIdUri') == 'urn:uuid:edef8ba9-79d6-4ace-a3c8-27dcd51d21ed':
                tag = cp.find('cenc:pssh')
                if tag:
                    return tag.text
    except Exception as e:
        print('PSSH extraction failed:', e, file=sys.stderr)
    return None


def _unpack_id(raw):
    """Strip leading 'id=' if the panel passes it through."""
    return raw[3:] if raw.startswith('id=') else raw


# ─── login ────────────────────────────────────────────────────────────────────

def login():
    """
    Accept the ascendon token (login-session cookie value from F1TV) via the
    'password' panel field. Fetches the entitlement token and saves both to
    the .tokens file.
    """
    print('logging in...', file=sys.stderr)

    raw_token = password.strip() if password else ''
    if not raw_token:
        print(
            'No ascendon token supplied.\n'
            'Set the Password field in the panel to the value of the '
            '"login-session" cookie from https://f1tv.formula1.com\n'
            '(DevTools → Application → Cookies → login-session)',
            file=sys.stderr,
        )
        return 'error'

    ascendon_token = _extract_ascendon(raw_token)
    ent_token      = _get_entitlement(ascendon_token)
    if not ent_token:
        return 'error'

    _save_auth({
        'ascendonToken':    ascendon_token,
        'entitlementToken': ent_token,
    })
    print('logged in successfully', file=sys.stderr)


# ─── action handlers ──────────────────────────────────────────────────────────

def do_action():
    if action == 'login':
        result = login()
        sys.exit(0 if result != 'error' else 1)

    auth = _load_auth()
    if not auth:
        return 'error'

    ascendon_token = auth.get('ascendonToken', '')
    # always refresh the entitlement token — it is short-lived
    token = _get_entitlement(ascendon_token)
    if not token:
        return 'error'

    if action in ('channels', 'events'):
        events = _get_events()
        if not events:
            print('No live events found', file=sys.stderr)
            return 'error'

        key     = 'Channels' if action == 'channels' else 'Events'
        output  = {key: []}
        now_ts  = int(datetime.datetime.now(pytz.UTC).timestamp())

        for c in events:
            content_id = str(c['contentId'])
            item = {
                'Name':            c.get('title', 'Unknown'),
                'Mode':            'live',
                'SessionManifest': True,
                'ManifestScript':  'id=' + content_id,
                'CdmType':         'widevine',
                'UseCdm':          True,
                'Cdm':             'id=' + content_id,
                'Video':           'best',
            }
            if action == 'events':
                item['Autostart'] = True
                item['Start']     = now_ts
                item['End']       = now_ts + 4 * 3600
            output[key].append(item)
        print(json.dumps(output, indent=2))

    elif action == 'heartbeat':
        sys.exit()

    elif action == 'manifest':
        content_id = _unpack_id(id)
        url, lic_url, lic_token = _get_single(token, content_id)
        if not url:
            return 'error'

        pssh_data = _get_pssh_from_mpd(url)

        # cache lic_url and lic_token for cdm step
        _save_auth({**auth,
                    'entitlementToken': token,
                    'manifest_cache': {
                        'lic_url':   lic_url,
                        'lic_token': lic_token,
                        'pssh':      pssh_data or '',
                    }})

        output = {
            'Cdn':         [{'Name': 'default', 'ManifestUrl': url}],
            'ManifestUrl': url,
            'Headers': {
                'Manifest': {'User-Agent': USER_AGENT},
                'Media':    {'User-Agent': USER_AGENT},
            },
            'Heartbeat': {'Url': '', 'Params': '', 'PeriodMs': 5 * 60 * 1000},
        }
        print(json.dumps(output))

    elif action == 'cdm' and cdm_param == 'internal':
        cache     = auth.get('manifest_cache', {})
        lic_url   = cache.get('lic_url', '')
        lic_token = cache.get('lic_token', token)
        if not lic_url:
            print('cdm=internal: no lic_url cached — run manifest action first',
                  file=sys.stderr)
            return 'error'

        h = {
            'entitlementtoken': lic_token,
            'origin':           F1TV_BASE,
            'user-agent':       USER_AGENT,
            'content-type':     'application/x-www-form-urlencoded',
        }
        r = req.post(lic_url, headers=h, data=base64.b64decode(challenge))
        response_b64 = base64.b64encode(r.content).decode()
        if response_b64.startswith('CA'):
            print(response_b64)
        else:
            print('License request failed', file=sys.stderr)
            print(r.text, file=sys.stderr)
            return 'error'

    elif action == 'cdm' and cdm_param == 'external':
        cache     = auth.get('manifest_cache', {})
        lic_url   = cache.get('lic_url', '')
        lic_token = cache.get('lic_token', token)
        pssh_b64  = pssh if pssh else cache.get('pssh', '')

        if not lic_url:
            print('cdm=external: no lic_url cached — run manifest action first',
                  file=sys.stderr)
            return 'error'
        if not pssh_b64:
            print('cdm=external: no PSSH available', file=sys.stderr)
            return 'error'

        try:
            wv_device = Device.load(WVD_PATH)
            wv_cdm    = Cdm.from_device(wv_device)
            sid       = wv_cdm.open()
            wv_chal   = wv_cdm.get_license_challenge(sid, PSSH(pssh_b64))
        except Exception as e:
            print('CDM init failed:', e, file=sys.stderr)
            return 'error'

        h = {
            'entitlementtoken': lic_token,
            'origin':           F1TV_BASE,
            'user-agent':       USER_AGENT,
            'content-type':     'application/x-www-form-urlencoded',
        }
        licence = req.post(lic_url, headers=h, data=wv_chal)
        try:
            wv_cdm.parse_license(sid, licence.content)
            for key in wv_cdm.get_keys(sid):
                if key.type != 'SIGNING':
                    print(f'{key.kid.hex}:{key.key.hex()}')
            wv_cdm.close(sid)
        except Exception as e:
            print('License parsing failed:', e, file=sys.stderr)
            print(licence.text, file=sys.stderr)
            return 'error'

    else:
        print('invalid action: ' + action, file=sys.stderr)


if do_action() == 'error':
    login()
    do_action()