#! /usr/bin/env python
# -*- coding: utf-8 -*-
'''
Created on 2023/6/27 18:43:40

@author: LiuKuan
@copyright: Apache License, Version 2.0
'''
import laok
from Crypto.Cipher import AES
# ===============================================================================
r'''
'''
# ===============================================================================
def download_tslist(url):
    aes = None
    index_file = laok.requests.req_save_auto(url, skip_exist=True)
    url_file = laok.path_replace_basename(index_file, 'url.txt')
    laok.file_write_text(url_file, url, skip_exist=True)
    for line in laok.file_read_lines(index_file, skip_empty=True):
        if line.startswith('#EXT-X-KEY'):
            kdata = [v.strip() for v in line[11:].split(',')]
            METHOD = ''
            IV = b'0000000000000000'
            KEY = b''

            for d in kdata:
                if d.startswith('METHOD='):
                    METHOD = d[7:]

                elif d.startswith('URI='):
                    URI = laok.url_join(url, d[4:].strip('""') )
                    key_file = laok.requests.req_save_auto(URI, skip_exist=True)
                    KEY = laok.file_read_bin(key_file)
                # elif d.startswith('IV='):
                #     iv = d[3:]
                #     IV = iv[2:18].encode()

            if METHOD.startswith('NONE'):
                aes = None
            else:
                aes = AES.new(KEY, AES.MODE_CBC, IV)

        elif line.startswith('#EXT-X-ENDLIST'):
            break
        elif line[0] != '#':
            if line.startswith('https://'):
                continue

            URI = laok.url_join(url, line)
            ts_file = laok.url_make_file(URI)
            if laok.path_exist(ts_file):
                continue

            print(f'[download]: {ts_file}', end=' ')
            laok.requests.req_save(ts_file, URI, skip_exist=True)
            print('success!', end=' ' if aes else '\n')

            if aes:
                print(f'[decrypt] method={METHOD}')
                ts_data = laok.file_read_bin(ts_file)
                ts_data = aes.decrypt(ts_data)
                laok.file_write_bin(ts_file, ts_data)

    return index_file

def down_m3u8_mp4(url):
    #下载所有 ts
    index_file = download_tslist(url)
    # 合并为 mp4
    merge_tslist(index_file)

def merge_tslist(index_file):
    base_name_list = []
    for data in laok.file_read_lines(index_file, skip_empty=True, strip=True):
        if data.startswith('#'):
            continue
        base_name_list.append(laok.path_basename(data))

    ts_list = list(laok.files_cur(laok.path_parent(index_file), '.ts'))
    # print(base_name_list)

    ts_list.sort(key=lambda x : base_name_list.index(laok.path_basename(x)) )
    # print(ts_list)
    dst_file = laok.path_replace_basename(index_file, 'all.mp4')
    laok.file_list_merge(ts_list, dst_file)

def decrypt_ts(ts_file, key):
    iv = b'0000000000000000'
    aes = AES.new(key, AES.MODE_CBC, iv)
    ts_data = laok.file_read_bin(ts_file)
    # print(ts_data)
    # print('\n\n============== lskf =======\n\n')
    ts_data = aes.decrypt(ts_data)
    # print(ts_data)
    ts_file = laok.path_deco_stem(ts_file, suffix='_lk')
    laok.file_write_bin(ts_file, ts_data)


if __name__ == '__main__':
    import fire as _fire
    _fire.Fire({
        'down_m3u8_mp4': down_m3u8_mp4,
        'download_tslist': download_tslist,
        'merge_tslist': merge_tslist,
        'decrypt_ts': decrypt_ts,
    })
