#!/usr/bin/env python3
"""Download a DRM-protected video from CDA.pl (Widevine DASH + DRMtoday)."""

import argparse
import json
import os
import re
import subprocess
import sys
import xml.etree.ElementTree as ET
from pathlib import Path

import requests
from bs4 import BeautifulSoup

HERE = Path(__file__).parent
CREDS_FILE = HERE / '.cda_creds'
TOOLS = HERE / 'tools'
N_M3U8 = TOOLS / 'N_m3u8DL-RE'
MP4DEC = TOOLS / 'mp4decrypt'
SAVE_DIR = HERE / 'downloads'
TMP_DIR = '/tmp/n_m3u8dl_tmp'

UA = ('Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 '
      '(KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36')
HEADERS = {
    'User-Agent': UA,
    'Content-Type': 'application/json',
    'X-Requested-With': 'XMLHttpRequest',
}

WIDEVINE_SYSTEM_ID = 'edef8ba9-79d6-4ace-a3c8-27dcd51d21ed'


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()


def login(session, username, password):
    r = session.post(
        'https://www.cda.pl/login',
        data={'username': username, 'password': password},
        headers={'User-Agent': UA},
        allow_redirects=True,
        timeout=15,
    )
    r.raise_for_status()
    if 'Zły login lub hasło!' in r.text:
        raise RuntimeError('Wrong CDA username or password.')
    if 'Zaznacz pole' in r.text:
        raise RuntimeError('Login blocked by CAPTCHA — try again later.')


def get_video_id(url):
    m = re.search(r'cda\.pl/video/([0-9a-z]+)', url)
    if not m:
        raise ValueError(f'Cannot extract video ID from URL: {url}')
    return m.group(1)


def get_stream_info(session, url, video_id):
    r = session.get(url, headers=HEADERS, timeout=15)
    r.raise_for_status()
    soup = BeautifulSoup(r.text, 'html.parser')

    title_tag = soup.find('h1')
    title = title_tag.get_text(strip=True) if title_tag else video_id

    div = soup.find('div', id=f'mediaplayer{video_id}')
    if not div:
        raise RuntimeError(
            f'mediaplayer div not found for {video_id}. '
            'Make sure you are logged in with a valid account.'
        )

    vi = json.loads(div['player_data'])['video']
    mpd_url     = vi.get('manifest')
    license_url = vi.get('manifest_drm_proxy')
    drm_header  = vi.get('manifest_drm_header')  # x-dt-custom-data for DRMtoday

    if not mpd_url:
        raise RuntimeError('No DASH manifest in player_data — video unavailable.')
    if not license_url:
        raise RuntimeError('No Widevine license URL in player_data.')

    return title, mpd_url, license_url, drm_header


def get_pssh_from_mpd(mpd_url, session):
    r = session.get(mpd_url, headers={'User-Agent': UA}, timeout=15)
    r.raise_for_status()

    root = ET.fromstring(r.text)
    # The Widevine UUID is on the ContentProtection schemeIdUri, not on the pssh element.
    # Walk ContentProtection elements, find the one with the Widevine UUID, grab its pssh child.
    for cp in root.iter('{urn:mpeg:dash:schema:mpd:2011}ContentProtection'):
        scheme = cp.get('schemeIdUri', '').lower()
        if 'edef8ba9' in scheme:
            pssh_el = cp.find('{urn:mpeg:cenc:2013}pssh')
            if pssh_el is not None and pssh_el.text:
                return pssh_el.text.strip()

    raise RuntimeError('No Widevine PSSH found in MPD.')


def fetch_content_key(pssh_b64, license_url, drm_header):
    from pywidevine.cdm import Cdm
    from pywidevine.device import Device
    from pywidevine.pssh import PSSH
    import uuid as _uuid

    device = Device.load(str(WVD))
    cdm = Cdm.from_device(device)
    session_id = cdm.open()
    challenge = cdm.get_license_challenge(session_id, PSSH(pssh_b64))

    req_headers = {
        'User-Agent': UA,
        'Content-Type': 'application/octet-stream',
    }
    if drm_header:
        req_headers['x-dt-custom-data'] = drm_header

    r = requests.post(license_url, data=challenge, headers=req_headers, timeout=15)
    if not r.ok:
        raise RuntimeError(f'License server {r.status_code}: {r.text[:300]}')

    cdm.parse_license(session_id, r.content)

    keys = {}
    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()
            keys[kid] = key.key.hex()
    cdm.close(session_id)
    return keys


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


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

    cookie_str = '; '.join(f'{k}={v}' for k, v in session.cookies.items())

    cmd = [
        str(N_M3U8),
        mpd_url,
        '--save-dir', str(SAVE_DIR),
        '--save-name', out_name,
        '--tmp-dir', TMP_DIR,
        '--decryption-binary-path', str(MP4DEC),
        '--mux-after-done', 'format=mkv',
        '--header', f'User-Agent: {UA}',
        '--header', f'Cookie: {cookie_str}',
    ] + key_args

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


def main():
    parser = argparse.ArgumentParser(
        description='Download a DRM-protected video from CDA.pl',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python3 cda_download.py -u mylogin -p mypass https://www.cda.pl/video/56706683e/vfilm

Credentials via env vars (avoids password in shell history):
  export CDA_USER=mylogin
  read -rs CDA_PASS && export CDA_PASS
  python3 cda_download.py https://www.cda.pl/video/56706683e/vfilm
""",
    )
    parser.add_argument('url', help='https://www.cda.pl/video/... URL')
    parser.add_argument('-u', '--username', default=os.environ.get('CDA_USER'),
                        help='CDA login email (or set CDA_USER env var)')
    parser.add_argument('-p', '--password', default=os.environ.get('CDA_PASS'),
                        help='CDA password (or set CDA_PASS env var)')
    parser.add_argument('-o', '--output',
                        help='Output filename without extension. Default: sanitized video title.')
    args = parser.parse_args()

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

    # Load saved credentials if not provided via flags or env vars
    if not args.username or not args.password:
        if CREDS_FILE.exists():
            saved = json.loads(CREDS_FILE.read_text())
            args.username = args.username or saved.get('username', '')
            args.password = args.password or saved.get('password', '')

    import getpass
    prompted = False
    if not args.username:
        args.username = input('CDA username (email): ').strip()
        prompted = True
    if not args.password:
        args.password = getpass.getpass('CDA password: ')
        prompted = True

    session = requests.Session()

    print(f'Logging in as {args.username}...')
    login(session, args.username, args.password)
    print('Logged in.')

    if prompted:
        CREDS_FILE.write_text(json.dumps({'username': args.username, 'password': args.password}))
        CREDS_FILE.chmod(0o600)
        print(f'Credentials saved to {CREDS_FILE}')

    video_id = get_video_id(args.url)
    print(f'Video ID: {video_id}')

    print('Fetching player data...')
    title, mpd_url, license_url, drm_header = get_stream_info(session, args.url, video_id)
    print(f'Title:       {title}')
    print(f'MPD:         {mpd_url}')
    print(f'License:     {license_url}')

    print('Extracting PSSH from MPD...')
    pssh_b64 = get_pssh_from_mpd(mpd_url, session)
    print(f'PSSH:        {pssh_b64[:48]}...')

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

    SAVE_DIR.mkdir(exist_ok=True)
    out_name = args.output if args.output else safe_filename(title)

    download_and_decrypt(mpd_url, keys, out_name, session)
    print(f'Done: {SAVE_DIR / out_name}.mkv')


if __name__ == '__main__':
    main()