#!/usr/bin/env python3
"""
TVP VOD batch downloader — Widevine L3 key extraction + DASH download.

Usage:
    python3 tvp_download.py <product_id> [<product_id> ...]
    python3 tvp_download.py --series <series_page_url>

Requirements:
    pip install pywidevine requests
    Tools: N_m3u8DL-RE, mp4decrypt (in TOOLS_DIR)
"""

import sys, os, subprocess, json, re, requests
from pathlib import Path

HERE = Path(__file__).parent
TOOLS = HERE / 'tools'
N_M3U8 = TOOLS / 'N_m3u8DL-RE'
MP4DEC = TOOLS / 'mp4decrypt'
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()

SAVE_DIR = HERE / 'downloads'
TMP_DIR  = '/tmp/n_m3u8dl_tmp'

API = 'https://vod.tvp.pl/api'
HEADERS = {
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 Chrome/120 Safari/537.36',
    'Referer': 'https://vod.tvp.pl/',
    'Accept': 'application/json',
}

# ──────────────────────────────────────────────────────────────────────────────

def get_playlist(product_id):
    """Returns dict with sources + drm info."""
    for vtype in ('MOVIE', 'LIVE'):
        r = requests.get(
            f'{API}/products/{product_id}/videos/playlist',
            params={'platform': 'BROWSER', 'videoType': vtype},
            headers=HEADERS, timeout=15)
        if r.ok:
            return r.json()
    raise RuntimeError(f"Playlist fetch failed for {product_id}: {r.status_code} {r.text[:120]}")


def get_product_info(product_id):
    r = requests.get(f'{API}/products/{product_id}', params={'platform': 'BROWSER'},
                     headers=HEADERS, timeout=15)
    if r.ok:
        return r.json()
    return {}


def fetch_content_key(pssh_b64, license_url):
    """Use pywidevine to fetch the CONTENT key from the license server."""
    from pywidevine.cdm import Cdm
    from pywidevine.device import Device
    from pywidevine.pssh import PSSH

    device = Device.load(str(WVD))
    cdm    = Cdm.from_device(device)
    session_id = cdm.open()

    challenge = cdm.get_license_challenge(session_id, PSSH(pssh_b64))

    r = requests.post(license_url,
                      data=challenge,
                      headers={'User-Agent': HEADERS['User-Agent'],
                               'Content-Type': 'application/octet-stream'},
                      timeout=15)
    if not r.ok:
        raise RuntimeError(f"License server {r.status_code}: {r.text[:200]}")

    cdm.parse_license(session_id, r.content)

    import uuid as _uuid
    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 extract_pssh_from_mpd(mpd_url):
    """Fetch MPD and extract Widevine PSSH base64."""
    r = requests.get(mpd_url, timeout=15)
    r.raise_for_status()
    # <cenc:pssh>...</cenc:pssh>
    m = re.search(r'<cenc:pssh>([^<]+)</cenc:pssh>', r.text)
    if m:
        return m.group(1).strip()
    # try default_KID only (build minimal PSSH)
    m2 = re.search(r'default_KID="([^"]+)"', r.text)
    if m2:
        kid_hex = m2.group(1).replace('-', '')
        import base64, struct
        # build minimal Widevine PSSH
        kid_bytes = bytes.fromhex(kid_hex)
        # proto: field 2 (key_id) = bytes
        proto = b'\x12\x10' + kid_bytes
        box_data = (
            b'\x00\x00\x00\x01'   # version=1
            + b'\x00\x00\x00\x00' # flags
            + bytes.fromhex('edef8ba979d64acea3c827dcd51d21ed')  # Widevine system_id
            + b'\x00\x00\x00\x00' * 4  # KID count placeholder
            + struct.pack('>I', len(proto)) + proto
        )
        return base64.b64encode(box_data).decode()
    raise RuntimeError("No PSSH found in MPD")


def make_save_name(info):
    """Build sXXeXX_title filename from product info."""
    # Season + episode from webUrl (contains e.g. ,S01E13,)
    web_url = info.get('webUrl', '')
    m = re.search(r'S(\d+)E(\d+)', web_url, re.IGNORECASE)
    season  = int(m.group(1)) if m else 1
    episode = info.get('number') or (int(m.group(2)) if m else 0)

    # Title: prefer slug (already ASCII-transliterated), strip the "odc-N--" prefix
    slug = info.get('slug', '')
    clean = re.sub(r'^odc-\d+-+', '', slug).strip('-')
    if not clean:
        # fallback: strip "odc. N – " from title, keep rest
        clean = re.sub(r'^odc\.\s*\d+\s*[–\-]\s*', '', info.get('title', str(info.get('id')))).strip()
        clean = re.sub(r'[^\w ]', '_', clean)
    clean = clean.replace('-', '_').strip('_')

    return f"s{season:02d}e{episode:02d}_{clean}"


def download(product_id, save_name=None):
    info = get_product_info(product_id)
    title = info.get('title', str(product_id))
    ptype = info.get('type', 'UNKNOWN')
    print(f"\n{'='*60}")
    print(f"Product {product_id}: {title!r} [{ptype}]")

    playlist = get_playlist(product_id)
    mpd_url  = playlist['sources']['DASH'][0]['src']
    lic_url  = playlist.get('drm', {}).get('WIDEVINE', {}).get('src')

    print(f"  MPD: {mpd_url[:80]}...")

    key_args = []
    if lic_url:
        # DRM-protected: extract PSSH from MPD, fetch Widevine content key.
        print(f"  License: {lic_url}")
        try:
            pssh = extract_pssh_from_mpd(mpd_url)
            # re-fetch fresh playlist (MPD URL may expire quickly)
            playlist = get_playlist(product_id)
            mpd_url  = playlist['sources']['DASH'][0]['src']
        except Exception as e:
            raise RuntimeError(f"PSSH extraction failed: {e}")

        print(f"  PSSH: {pssh[:40]}...")
        keys = fetch_content_key(pssh, lic_url)
        if not keys:
            raise RuntimeError("No CONTENT keys returned by license server")

        for kid, key in keys.items():
            print(f"  Key: {kid}:{key}")
            key_args += ['--key', f'{kid}:{key}']
    else:
        # No DRM in playlist — clear (unencrypted) stream, download directly.
        print("  DRM: none (clear stream) — skipping key extraction")

    if save_name is None:
        save_name = make_save_name(info)

    SAVE_DIR.mkdir(parents=True, exist_ok=True)

    cmd = [
        str(N_M3U8), mpd_url,
        *key_args,
        '--save-dir', str(SAVE_DIR),
        '--save-name', save_name,
        '--select-video', 'best',
        '--select-audio', 'best',
        '--tmp-dir', TMP_DIR,
        '--decryption-binary-path', str(MP4DEC),
        '--mux-after-done', 'format=mp4:bin=ffmpeg',
    ]

    print(f"\n  Running: N_m3u8DL-RE {save_name}")
    result = subprocess.run(cmd)
    if result.returncode != 0:
        print(f"  [!] Download failed (exit {result.returncode})")
        return False

    out_file = SAVE_DIR / f"{save_name}.mp4"
    if out_file.exists():
        size_mb = out_file.stat().st_size // 1024 // 1024
        print(f"  [+] Saved: {out_file.name} ({size_mb} MB)")
        return True
    else:
        print(f"  [!] Output file not found: {out_file}")
        return False


def get_series_episodes(series_id):
    """
    Get all episode product IDs for a series.
    series_id: the SERIAL product id (e.g. 2726890 for "Wojna zastępcza")
    Returns list of (episode_number, product_id, title) sorted by episode number.
    """
    VOD_BASE = f'{API}/products/vods'
    # Get seasons
    r = requests.get(f'{VOD_BASE}/serials/{series_id}/seasons',
                     params={'platform': 'BROWSER'}, headers=HEADERS, timeout=15)
    if not r.ok:
        raise RuntimeError(f"Seasons fetch failed: {r.status_code} {r.text[:100]}")
    seasons = r.json()
    if not seasons:
        raise RuntimeError(f"No seasons found for series {series_id}")

    episodes = []
    for season in seasons:
        season_id = season['id']
        r2 = requests.get(f'{VOD_BASE}/serials/{series_id}/seasons/{season_id}/episodes',
                          params={'platform': 'BROWSER'}, headers=HEADERS, timeout=15)
        if r2.ok:
            for ep in r2.json():
                episodes.append((ep.get('number', 0), ep['id'], ep.get('title', '')))

    episodes.sort(key=lambda x: x[0])
    return episodes


# ──────────────────────────────────────────────────────────────────────────────

if __name__ == '__main__':
    if not WVD:
        print(f"[!] No .wvd file found in {HERE / 'L3'}!")
        sys.exit(1)

    if not N_M3U8.exists() or not MP4DEC.exists():
        print(f"[!] Tools missing: {N_M3U8} / {MP4DEC}")
        sys.exit(1)

    args = sys.argv[1:]
    if not args:
        print("Usage:")
        print("  tvp_download.py <product_id> [<product_id> ...]   # download episodes by ID")
        print("  tvp_download.py --series <series_id>               # download all episodes of a series")
        sys.exit(1)

    product_ids = []
    if args[0] == '--series':
        series_id = int(args[1])
        print(f"Fetching episode list for series {series_id}...")
        episodes = get_series_episodes(series_id)
        print(f"Found {len(episodes)} episodes:")
        for num, pid, title in episodes:
            print(f"  E{num:02d} | id={pid} | {title}")
        product_ids = [pid for _, pid, _ in episodes]
    else:
        product_ids = [int(a) for a in args if a.isdigit()]

    ok = failed = 0
    for pid in product_ids:
        try:
            if download(pid):
                ok += 1
            else:
                failed += 1
        except Exception as e:
            print(f"  [!] Error for {pid}: {e}")
            import traceback; traceback.print_exc()
            failed += 1

    print(f"\nDone: {ok} succeeded, {failed} failed")