import argparse
import requests
from bs4 import BeautifulSoup
import json
import logging
import coloredlogs
from cdm.wks import WvDecrypt, device_android_generic, PsshExtractor, KeyExtractor

# def search(data, target_slug):
#     if isinstance(data, list):
#         for item in data:
#             result = search(item, target_slug)
#             if result:
#                 return result
#     elif isinstance(data, dict):
#         for key, value in data.items():
#             if key == "slug" and value == target_slug:
#                 return data.get("productId")
#             else:
#                 result = search(value, target_slug)
#                 if result:
#                     return result
#     return None

parser = argparse.ArgumentParser(description='PYWKS-NPO')

parser.add_argument('-url', dest='url', required=True, help='NPO Video URL')
parser.add_argument("-logger", default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], help="Logger level")
args = parser.parse_args()
LOG_FORMAT = "{asctime} [{levelname[0]}] {name} : {message}"
LOG_DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
LOG_STYLE = "{"
coloredlogs.install(level=args.logger, fmt=LOG_FORMAT, datefmt=LOG_DATE_FORMAT, style=LOG_STYLE)
logger = logging.getLogger("NPO")

response_target_id = requests.get(args.url)
content = response_target_id.content
try:
    target_slug = args.url.split("/")[7]
except IndexError:
    logger.error("Invalid URL format. Example: https://npo.nl/start/serie/wie-is-de-mol/seizoen-24/wie-is-de-mol_56/afspelen")
    exit()
soup = BeautifulSoup(content, 'html.parser')
script_tag = soup.find('script', {'id': '__NEXT_DATA__'})
script_content = script_tag.contents[0] if script_tag else logger.error("Script tag not found.")
data_dict = json.loads(script_content)
target_product_id = search(data_dict, target_slug)
if not target_product_id:
    logger.error("Failed to retrieve target product ID.")
    exit()
response_csrf = requests.get('https://npo.nl/start/api/auth/session')
cookies = {'__Host-next-auth.csrf-token': response_csrf.cookies.get_dict()["__Host-next-auth.csrf-token"],'__Secure-next-auth.callback-url': 'https://npo.nl'}
json_product_id = {'productId': target_product_id}
response_token = requests.post('https://npo.nl/start/api/domain/player-token', cookies=cookies, json=json_product_id)
headers = {'authorization': response_token.json()["token"]}
json_auth = {'profileName': 'dash', 'drmType': 'widevine', 'referrerUrl': args.url}
response = requests.post('https://prod.npoplayer.nl/stream-link', headers=headers, json=json_auth)
stream_data = response.json().get('stream', {})
if not stream_data.get('streamURL'):
    logger.error("Failed to retrieve MPD URL. Invalid or expired authentication token.")
    exit()
mpd_url = stream_data.get('streamURL')
logger.info(f"MPD URL: {mpd_url}")
license_url = "https://npo-drm-gateway.samgcloud.nepworldwide.nl/authentication"
response = requests.get(mpd_url, headers=headers)
pssh_extractor = PsshExtractor(response.text)
pssh_value = pssh_extractor.extract_pssh()
logger.info(f"PSSH: {pssh_value}")
headers_license = {'x-custom-data': stream_data.get('drmToken'),'origin': 'https://start-player.npo.nl','referer': 'https://start-player.npo.nl/'}
cert_b64 = None
key_extractor = KeyExtractor(pssh_value, cert_b64, license_url, headers_license)
keys = key_extractor.get_keys()
wvdecrypt = WvDecrypt(init_data_b64=pssh_value, cert_data_b64=cert_b64, device=device_android_generic)
raw_challenge = wvdecrypt.get_challenge()
for key in keys:
    if isinstance(key, list) and key:
        for key_str in key:
            logger.info(f"\u251C KEY: {key_str}")