from devine.core.manifests import dash, hls
from devine.core.tracks import Tracks
import requests
from . import Decrypt

# Define a function to get PSSH from MPD
def get_pssh(mpd_or_m3u8_url):
    try:
        tracks = Tracks()
        client = requests.Session()
        License_cURL = Decrypt.load_license_curl()
        if License_cURL.headers:
            client.headers.update(License_cURL.headers)
        if License_cURL.cookies:
            client.cookies.update(License_cURL.cookies)
        mpd = dash.DASH.from_url(url=mpd_or_m3u8_url, session=client).to_tracks(language="en")
        tracks.add(mpd)
        PSSH = []
        for track in tracks:
            representation = track.data['dash']['representation']
            adaptation_set = track.data['dash']['adaptation_set']
            track_drm = dash.DASH.get_drm(
                representation.findall("ContentProtection") +
                adaptation_set.findall("ContentProtection")
            )
            for drm in track_drm:
                PSSH.append(drm.pssh)
        pssh_list = []
        for pssh in PSSH:
            if str(pssh) not in pssh_list:
                pssh_list.append(str(pssh))
        return pssh_list
    except:
        try:
            tracks = Tracks()
            client = requests.Session()
            License_cURL = Decrypt.load_license_curl()
            client.headers.update(License_cURL.headers)
            client.cookies.update(License_cURL.cookies)
            hls = hls.from_url(url=mpd_or_m3u8_url, session=client).to_tracks(language="en")
            tracks.add(hls)
            PSSH = []
            for track in tracks:
                representation = track.data['dash']['representation']
                adaptation_set = track.data['dash']['adaptation_set']
                track_drm = dash.DASH.get_drm(
                    representation.findall("ContentProtection") +
                    adaptation_set.findall("ContentProtection")
                )
                for drm in track_drm:
                    PSSH.append(drm.pssh)
            pssh_list = []
            for pssh in PSSH:
                if str(pssh) not in pssh_list:
                    pssh_list.append(str(pssh))
            return pssh_list
        except:
            return