import subprocess
import sys
import requests
import xml.etree.ElementTree as ET
import re
from pywidevine.cdm import Cdm
from pywidevine.device import Device
from pywidevine.pssh import PSSH


def call_tasks(subcommand, *args):
    command = [sys.executable, 'tasks.py', subcommand] + list(args)

    try:
        result = subprocess.run(command, capture_output=True, text=True, check=True)
        return result.stdout.strip()
    except subprocess.CalledProcessError as e:
        print(f"Error executing {subcommand}: {e}")
        return None


def parse_output(output):
    dash_link = None
    widevine_license_server = None
    token = None

    lines = output.splitlines()

    for line in lines:
        if line.startswith("DASH Link:"):
            dash_link = line.split("DASH Link: ", 1)[1].strip()
        elif line.startswith("Widevine License Server:"):
            widevine_license_server = line.split("Widevine License Server: ", 1)[1].strip()
        elif line.startswith("Token:"):
            token = line.split("Token: ", 1)[1].strip()

    return dash_link, widevine_license_server, token


def download_mpd(dash_link):
    try:
        response = requests.get(dash_link)
        response.raise_for_status()
        return response.text
    except requests.RequestException as e:
        print(f"Error downloading MPD file: {e}")
        return None


def extract_pssh(mpd_content):
    pssh_strings = set()
    root = ET.fromstring(mpd_content)

    # More flexible approach to find <pssh> elements regardless of namespaces
    for elem in root.iter():
        if elem.tag.endswith('pssh') and elem.text:
            pssh_strings.add(elem.text.strip())

    return pssh_strings


KEY_REGEX = re.compile(r'\b([0-9a-fA-F]{32}):([0-9a-fA-F]{32})\b')


def extract_keys_from_cdm_output(text):
    """
    Extract and return list of keys matching 32hex:32hex from CDM output text.
    """
    matches = KEY_REGEX.findall(text)
    return [f"{kid.lower()}:{key.lower()}" for kid, key in matches]


def fetch_keys_from_pssh(pssh_b64, license_url, token=None):
    try:
        pssh = PSSH(pssh_b64)

        # Load device (adjust path to your provisioned .wvd file)
        device = Device.load("cdm.wvd")  # <-- Replace with your actual path
        cdm = Cdm.from_device(device)

        session_id = cdm.open()

        # Get license challenge
        challenge = cdm.get_license_challenge(session_id, pssh)

        headers = {'x-vudrm-token': token} if token else {}

        # Send license request
        response = requests.post(license_url, data=challenge, headers=headers)
        response.raise_for_status()

        # Parse license
        cdm.parse_license(session_id, response.content)

        # Extract keys from CDM object
        keys_text = ""
        for key in cdm.get_keys(session_id):
            keys_text += f"[{key.type}] {key.kid.hex}:{key.key.hex()}\n"

        # Close session
        cdm.close(session_id)

        # Extract and print only keys matching 32hex:32hex pattern
        keys = extract_keys_from_cdm_output(keys_text)
        if keys:
            for k in keys:
                print(k)
        else:
            print("No 32hex:32hex keys found in this session.")

    except Exception as e:
        print(f"Error processing PSSH: {e}")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python wrapper.py <subcommand> [args...]")
        sys.exit(1)

    subcommand = sys.argv[1]
    args = sys.argv[2:]

    # Step 1: Run tasks.py and get output
    output = call_tasks(subcommand, *args)

    if output is not None:
        # Step 2: Parse output
        dash_link, widevine_license_server, token = parse_output(output)

        # Step 3: Download MPD file
        mpd_content = download_mpd(dash_link)

        if mpd_content:
            # Step 4: Extract PSSHs
            pssh_strings = extract_pssh(mpd_content)

            
            for pssh in pssh_strings:
                print("")

            # Step 5: Fetch keys using each PSSH
            print("")
            for pssh in sorted(pssh_strings):
                print()
                fetch_keys_from_pssh(pssh, widevine_license_server, token)