#!/usr/bin/env python3.6
# -*- encoding: utf-8 -*-
'''
@File    :   m3u8b.py
@Time    :   2020/04/23 12:54:28
@Author  :   Tang Jing 
@Version :   1.0.0
@Contact :   yeihizhi@163.com
@License :   (C)Copyright 2020
@Desc    :   None
'''

# here put the import lib
import re
import os
import io
import shutil
import time
# from Crypto.Cipher import AES 解密后CKPLAYER播放有问题.
from threading import Thread
from threading import Lock
from threading import Event as ThreadEvent
from urllib.parse import urlparse
from io import StringIO, BytesIO
from TDhelper.Cache.pools import pools
from TDhelper.network.http.http_helper import m_http
from TDhelper.network.http.status.M3U8_STATUS import M3U8_STATUS

# code start
cfg = {
    "cfg": {
        'timeout': 5,  # 超时(秒)
        'reconnect': 5  # 连接重试次数
    },
    "save-folder": "tmp\\ts",  # 临时文件件
    "download-poolsize": 20,  # 下载线程
    "buffer-index": 5  # 下载多少文件后生成索引文件
}


class m3u8:
    def __init__(self, m3u8Cfg: dict = None):
        self.setting = m3u8Cfg if m3u8Cfg else cfg  # 如果不传配置参数，则使用默认参数.
        self._state = True  # 运行状态
        self._http = m_http()
        # requests 连接失败重连次数
        self._reconnect = self.setting['cfg']['reconnect']
        # requests超时时间
        self._timeout = self.setting['cfg']['timeout']
        # 文件保存文件夹
        self._save_folder = self.setting['save-folder']
        self._m3u8_file_hash = None
        self._create_m3u8_index_file = False  # 是否已经建立M3U8播放索引
        self._del_state = False  # 是否正在删除过期文件,LIVE模式需要使用.
        # 密钥
        self._m3u8_key = None
        # 下载文件线程池配置
        self._lock__ = Lock()  # 线程锁
        self._download_event__ = ThreadEvent()  # 下载线程同步
        self._complete_event__ = ThreadEvent()  # 下载完成通知事件
        self._event_create_m3u8_file__ = ThreadEvent()  # 生成m3u8索引文件事件
        self._download_complate_total__ = 0  # 统计下载完成数量
        self._download_thread_active_ = 0  # 下载线程活跃数,用于判断文件是否下载完成.
        self._downlaod_buffer_length = 0  # 下载缓存长度, 用于计算快偏移量
        self._download_pool_size = self.setting['download-poolsize']
        self._download_pool = pools(self._download_pool_size)

        # 初始化下载线程池
        for i in range(0, self._download_pool_size):
            self._download_pool.push(m_http())
        self._save_folder = self.__checkfolder__(self._save_folder)

        # 下载缓存索引，下载索引到达这个值时才生成播放索引文件
        self._ts_save_folder = None
        self._m3u8_buffer_index = self.setting['buffer-index']
        self._current_m3u8_file = None
        # live:直播(需要连续多次下载M3U8文件),whole:完整播放(点播不需要多次下载M3U8文件)
        self._play_mode__ = 'whole'
        # M3U8 TAGS
        self._prefix_url = None
        self._EXT_X_TARGETDURATION__ = None
        self._EXTINF__ = []
        self._EXT_X_VERSION__ = None
        self._EXT_X_BYTERANGE__ = None
        self._EXT_X_KEY__ = None
        self._EXT_X_MAP__ = None
        self._EXT_X_PROGRAM_DATE_TIME__ = None
        self._EXT_X_DATERANGE__ = None
        self._EXT_X_MEDIA_SEQUENCE__ = None
        self._EXT_X_DISCONTINUITY_SEQUENCE__ = None
        self._EXT_X_PLAYLIST_TYPE__ = None
        self._EXT_X_MEDIA__ = None
        self._EXT_X_STREAM_INF__ = None
        self._EXT_X_I_FRAME_STREAM_INF__ = None
        self._EXT_X_SESSION_DATA__ = None
        self._EXT_X_SESSION_KEY__ = None
        self._EXT_X_START__ = None
        self._EXT_X_DISCONTINUITY__ = None
        self._EXT_X_ENDLIST__ = None
        self._EXT_X_I_FRAMES_ONLY__ = None
        self._EXT_X_INDEPENDENT_SEGMENTS__ = None

        # 自己定义M3U8 处理属性
        self._EXT_X_M3U8_FILE_LIST__ = []  # 直播类有m3u8有多个m3u8文件地址
        self._EXT_X_M3U8_PLAY_LIST__ = []  # m3u8播放列表TS文件列表
        self._EXT_X_M3U8_PLAY_LIST_INDEX__ = 0  # 下载m3u8播放列表TS文件列表索引
        self._LIVE_DOWNLOAD_FILE = []  # 直播下载文件

    def __resetVar__(self):
        '''
            重置变量
        '''
        self._ts_save_folder = None
        self._EXT_X_KEY__ = None
        self._download_event__.clear()
        self._complete_event__.clear()
        self._event_create_m3u8_file__.clear()
        self._download_complate_total__ = 0
        self._create_m3u8_index_file = False
        self._EXT_X_M3U8_FILE_LIST__ = []
        self._EXT_X_M3U8_PLAY_LIST__ = []
        self._EXT_X_M3U8_PLAY_LIST_INDEX__ = 0
        self._EXTINF__ = []
        self._download_complate_total__ = 0
        self._download_thread_active_ = 0
        self._downlaod_buffer_length = 0
        self._del_state = False
        if self._play_mode__ == "whole":
            self._LIVE_DOWNLOAD_FILE = []
            self._m3u8_file_hash = None

    def Get(self, m3u8_url, ts_save_path: str = None, download_start_offset: int = 0):
        try:
            self._EXT_X_M3U8_PLAY_LIST_INDEX__ = download_start_offset
            m_count = 0
            while True:
                if not self._state:
                    self._event_create_m3u8_file__.set()
                    self._download_event__.set()
                    self._complete_event__.set()
                    break
                self.__resetVar__()
                if ts_save_path:
                    self._ts_save_folder = ts_save_path
                else:
                    return 'ts save folder is none.', M3U8_STATUS.ERROR
                ret, status = self.__getM3U8file__(m3u8_url)
                if status == M3U8_STATUS.WAIT:
                    time.sleep(1)
                elif status == M3U8_STATUS.SUCCESS:
                    if self._play_mode__ == 'live':
                        self._ts_save_folder = 'live'
                        if m_count == 0:
                            m_count += 1
                            if os.path.exists(os.path.join(self._save_folder, self._ts_save_folder)):
                                for item in os.listdir(os.path.join(self._save_folder, self._ts_save_folder)):
                                    os.remove(os.path.join(self._save_folder, *(self._ts_save_folder,item))) #直播类，删除已下载TS文件.
                    self.__get_key__()  # 下载KEY文件.模块内有判断是否需要下载.
                    ret, status = self.__download_ts__()
                    if self._play_mode__ == 'live':
                        continue
                    else:
                        if status == M3U8_STATUS.SUCCESS:
                            return 'SUCCESS', M3U8_STATUS.SUCCESS
                        return ret, status
                elif status == M3U8_STATUS.ERROR:
                    return ret, M3U8_STATUS.ERROR
                else:
                    return 'Unknow error.', M3U8_STATUS.ERROR
        finally:
            self._event_create_m3u8_file__.set()
            self._complete_event__.set()
            self._download_event__.set()

    def getComplate(self, func):
        self._complete_event__.wait()
        func((), {})
        self._complete_event__.clear()

    def GetM3U8File(self):
        '''
            获取生成的M3U8索引文件
        '''
        self._event_create_m3u8_file__.wait()  # 阻塞等待生成播放索引.
        self._event_create_m3u8_file__.clear()  # 重置状态.
        return os.path.join(self._save_folder, *(self._ts_save_folder, 'index.m3u8',)), M3U8_STATUS.SUCCESS

    def close(self):
        self._state = False
        self._event_create_m3u8_file__.set()
        self._complete_event__.set()
        self._download_event__.set()

    def seek(self, time):
        ''''
        寻道时间定位
        '''
        total_time = 0
        index = 0
        for item in self._EXT_X_M3U8_PLAY_LIST__:
            total_time += float(item['duration'])
            if total_time >= time:
                self._EXT_X_M3U8_PLAY_LIST_INDEX__ = index
                break
            index += 1

    def __checkfolder__(self, folder):
        # 检查临时文件夹路径是否完整
        self._lock__.acquire()
        try:
            if not re.search(r':\\', folder, re.I | re.M):
                folder = os.path.join(os.getcwd(), folder)
            # 如果文件夹不存在则建立文件夹
            if not os.path.exists(folder):
                os.mkdir(folder)
        except OSError as e:
            raise e
        self._lock__.release()
        return folder

    def __getM3U8file__(self, uri: str):
        '''
            下载m3u8文件
        '''
        self._prefix_url = uri.rsplit('/', 1)[0]  # 获取URI前缀
        self._current_m3u8_file = uri  # 当前下载M3U8 URI
        response, status = self.__request__(uri)
        if status == 200:
            if response == self._m3u8_file_hash:
                return uri, M3U8_STATUS.WAIT  # m3u8文件没有变化.返回等待状态.
            self._m3u8_file_hash = response
            self.__analysis__(response)
            if self._EXT_X_STREAM_INF__:
                self._EXT_X_STREAM_INF__ = None  # 重置STREAM_INF属性
                # 下载下一层M3U8.
                return self.__getM3U8file__(self.__spell_uri__(
                    self._EXT_X_M3U8_FILE_LIST__[0]))
            else:
                if self._EXT_X_ENDLIST__:
                    self._play_mode__ = 'whole'
                    self._m3u8_buffer_index = self.setting['buffer-index']
                else:
                    self._play_mode__ = 'live'
                    self._m3u8_buffer_index = 1
                return uri, M3U8_STATUS.SUCCESS
        else:
            self._complete_event__.set()
            self._event_create_m3u8_file__.set()
            return uri, M3U8_STATUS.ERROR

    def __download_ts__(self):
        if self._EXTINF__:
            # 生成ts文件下载列表.
            m_index = 0
            for item in self._EXTINF__:
                m_file = item.split(',')[1]
                duration = item.split(',')[0]
                if not re.match(r"^(http|https)://", m_file, re.I | re.M):
                    download_uri = self.__spell_uri__(m_file)
                else:
                    download_uri = m_file
                local_file_name = m_file.rsplit('/', 1)
                if len(local_file_name) == 1:
                    local_file_name = local_file_name[0]
                else:
                    local_file_name = local_file_name[1]
                m_download_struct = {'index': m_index, 'url': download_uri, 'offset': None,
                                     'duration': duration, 'filename': local_file_name, 'state': False}
                self._EXT_X_M3U8_PLAY_LIST__.append(m_download_struct)
                m_index += 1
            # 开始下载
            for i in range(0, self._download_pool_size):
                Thread(target=self.__get_ts_thread__).start()
            self._download_event__.wait()  # 阻塞，等待下载线程完成下载
            self._download_event__.clear()
            self._complete_event__.set()
            return 'SUCCESS', M3U8_STATUS.SUCCESS
        else:
            return 'EXTINF IS NONE.', -1

    def __get_ts_thread__(self):
        '''
        TS文件下载
        '''
        self._lock__.acquire()
        if self._download_thread_active_ < self._download_pool_size:
            self._download_thread_active_ += 1
        else:
            raise Exception('download thread error: thread active total count({0}) Over length ({1}).'.format(
                self._download_thread_active_, self._download_pool_size))
        # thiread_id = self._download_thread_active_
        self._lock__.release()
        m_http_control = None
        m_http_control = self._download_pool.pop()
        m_index = 0
        if m_http_control:
            while True:
                if not self._state:
                    # 调用close方法，退出下载线程
                    self._download_pool.push(m_http_control)
                    break
                ts_struct = None  # TS文件下载信息结构体
                # 从播放列表取TS文件地址
                if self._EXT_X_M3U8_PLAY_LIST_INDEX__ < len(self._EXT_X_M3U8_PLAY_LIST__):
                    ts_struct = self._EXT_X_M3U8_PLAY_LIST__[
                        self._EXT_X_M3U8_PLAY_LIST_INDEX__]
                    m_index = self._EXT_X_M3U8_PLAY_LIST_INDEX__
                    self._lock__.acquire()
                    self._EXT_X_M3U8_PLAY_LIST_INDEX__ += 1
                    self._lock__.release()
                else:
                    # 下载完毕跳出下载线程
                    self._download_pool.push(m_http_control)
                    break
                if ts_struct:
                    if ts_struct['state']:
                        continue
                    if os.path.exists(os.path.join(self._save_folder, *(self._ts_save_folder, ts_struct['filename']))):
                        continue
                    file_bytes, status = m_http_control.download(
                        ts_struct['url'])
                    if status != 200:
                        i = 0
                        while True:
                            file_bytes, status = m_http_control.download(
                                ts_struct['url'])
                            i += 1
                            if i > self._reconnect or status == 200:  # if status==200:
                                break
                    else:
                        print('\rdownload:{1}({0}%)'.format(self._EXT_X_M3U8_PLAY_LIST_INDEX__/len(
                            self._EXT_X_M3U8_PLAY_LIST__)*100, self._current_m3u8_file), end='')
                        # 建立播放索引.
                        if ts_struct['index'] >= self._m3u8_buffer_index:
                            self.__create_m3u8_play_file__()
                        # 更改TS下载列表对应文件下载状态.
                        self._EXT_X_M3U8_PLAY_LIST__[m_index]['state'] = True
                        self._lock__.acquire()
                        ts_struct['offset'] = (
                            self._downlaod_buffer_length, len(file_bytes))
                        self._downlaod_buffer_length += len(file_bytes)
                        self._lock__.release()
                        self._EXT_X_M3U8_PLAY_LIST__[
                            ts_struct['index']] = ts_struct
                        # 用索引号替代原TS文件名称 ts_struct['url'].rsplit('/', 1)[1]
                        tmp_file_name = str(
                            ts_struct['filename']).split('?')[0]
                        # 直播类过期TS文件删除，模块内判断是否是直播.
                        self.__del_live_ts__(tmp_file_name)
                        '''if self._m3u8_key:
                        # 解密视频文件(CKPlayer播放器解密后的TS不能连续下载)不知道原因
                        cryptor = AES.new(bytes(self._m3u8_key['key'], encoding='utf-8'), AES.MODE_CBC, bytes(
                            self._m3u8_key['key'], encoding='utf-8'))
                        file_bytes = cryptor.decrypt(file_bytes)'''
                        # 生成TS文件.
                        m_saveTsPath = os.path.join(
                            self._save_folder, self._ts_save_folder)
                        self.__checkfolder__(m_saveTsPath)
                        with open(os.path.join(m_saveTsPath, tmp_file_name), mode="wb") as tmp_stream:
                            tmp_stream.write(file_bytes)
                            tmp_stream.flush()
                            tmp_stream.close()

                else:
                    continue  # 下载信息为空则进入下个下载循环.
            # 下载完成退出逻辑.
            self._lock__.acquire()
            self._download_thread_active_ -= 1
            event_state = False
            if self._download_thread_active_ == 0:
                event_state = True
            if event_state:
                self._download_event__.set()
            self._lock__.release()
        else:
            raise Exception('httpContent is none.')

    def __del_live_ts__(self, ts_file_name):
        '''
            删除直播过期文件.
        '''
        self._lock__.acquire()
        if self._del_state:
            return None
        self._del_state = True
        if self._play_mode__ == 'live':
            if len(self._LIVE_DOWNLOAD_FILE) < len(self._EXTINF__)*4:
                self._LIVE_DOWNLOAD_FILE.append(ts_file_name)
            else:
                count = 0
                size = int((len(self._EXTINF__)*4)/2)
                while True:
                    if count < size:
                        try:
                            os.remove(os.path.join(
                                self._save_folder, *(self._ts_save_folder, self._LIVE_DOWNLOAD_FILE[0])))
                            del self._LIVE_DOWNLOAD_FILE[0]
                        except Exception as e:
                            raise e
                    else:
                        break
                    count += 1
                self._LIVE_DOWNLOAD_FILE.append(ts_file_name)
        self._del_state = False
        self._lock__.release()

    def __create_m3u8_play_file__(self):
        '''
            建立.m3u8播放索引文件
        '''
        if self._create_m3u8_index_file:
            return None
        tmp_saveTsPath = os.path.join(self._save_folder, self._ts_save_folder)
        self.__checkfolder__(tmp_saveTsPath)
        with open(os.path.join(tmp_saveTsPath, 'index.m3u8'), mode='wb') as f:
            f.writelines([bytes('#EXTM3U\n', encoding='utf-8')])
            f.writelines([bytes('#EXT-X-ALLOW-CACHE:YES\n', encoding='utf-8')])
            for item in self.__dict__:
                if self.__dict__[item]:
                    if re.match('_EXT_X_VERSION', item, re.I | re.M):
                        if self._EXT_X_VERSION__:
                            m_line = "#EXT-X-VERSION:"+self._EXT_X_VERSION__+"\n"
                            f.writelines([bytes(m_line, encoding='utf-8')])
                    elif re.match('_EXT_X_KEY__', item, re.I | re.M):
                        m_line = "#EXT-X-KEY:METHOD={0},URI=\"key.key\"\n".format(
                            self._m3u8_key['method'])
                        f.writelines([bytes(m_line, encoding='utf-8')])
                    elif re.match('_EXT_X_MEDIA_SEQUENCE', item, re.I | re.M):
                        m_line = ''
                        if self._EXT_X_MEDIA_SEQUENCE__:
                            m_line = "#EXT-X-MEDIA-SEQUENCE:"+self._EXT_X_MEDIA_SEQUENCE__+"\n"
                        else:
                            m_line = "#EXT-X-MEDIA-SEQUENCE:0\n"
                        f.writelines([bytes(m_line, encoding='utf-8')])
                    elif re.match('_EXT_X_ALLOW_CACHE', item, re.I | re.M):
                        pass
                    elif re.match('_EXT_X_TARGETDURATION', item, re.I | re.M):
                        if self._EXT_X_TARGETDURATION__:
                            m_line = '#EXT-X-TARGETDURATION:'+self._EXT_X_TARGETDURATION__+"\n"
                            f.writelines([bytes(m_line, encoding='utf-8')])
                    elif re.match('_EXT_X_M3U8_PLAY_LIST__', item, re.I | re.M):
                        for extinf in self._EXT_X_M3U8_PLAY_LIST__:
                            m_line = '#EXTINF:' + \
                                extinf['duration']+',\n' + \
                                str(extinf['filename'].split('?')[0])+"\n"
                            f.writelines([bytes(m_line, encoding='utf-8')])
                        if self._EXT_X_ENDLIST__:
                            m_line = '#EXT-X-ENDLIST\n'
                            f.writelines([bytes(m_line, encoding='utf-8')])
            f.flush()
            f.close()
        self._create_m3u8_index_file = True
        self._event_create_m3u8_file__.set()  # 解锁M3U8文件生成成功锁

    def __get_key__(self):
        '''
            获取密钥文件
        '''
        if self._EXT_X_KEY__:
            with StringIO(self._EXT_X_KEY__.replace(',', "\r\n")) as f:
                self._m3u8_key = {'method': None, 'key': '', 'iv': None}
                while True:
                    line = f.readline()
                    if line:
                        line = re.sub('\"|\'|\\r|\\n', '',
                                      line, count=0, flags=0)
                        line = line.split('=')
                        if len(line) == 2:
                            if line[0].lower() == 'method':
                                self._m3u8_key['method'] = line[1]
                            elif line[0].lower() == 'uri':
                                key_url = ''
                                if not re.match("http://|https://", line[1], re.I | re.M):
                                    key_url = self.__spell_uri__(line[1])
                                    state = False
                                    content = ''
                                    download_count = 0
                                    while True:
                                        content, state = self.__request__(
                                            key_url)
                                        download_count += 1
                                        if download_count > self._reconnect or state == 200:
                                            break
                                        if state == 200:
                                            self._m3u8_key['key'] = content
                                            with open(os.path.join(self._save_folder, *(self._ts_save_folder, 'key.key')), mode='wb') as keyStream:
                                                keyStream.write(
                                                    bytes(content, encoding='utf-8'))
                                                keyStream.flush()
                            elif line[0].lower() == 'iv':
                                self._m3u8_key['iv'] = line[1]
                    else:
                        break
        else:
            return None

    def __spell_uri__(self, uri):
        '''
            拼接资源URI
        '''
        ret = self._prefix_url
        for item in uri.split('/'):
            if not re.search(item, self._prefix_url, re.M | re.I):
                ret += '/'+item
        return ret

    def __analysis__(self, content):
        '''
            分析M3U8文件
        '''
        m_result = self.__match__(r"^#EXTM3U", content)
        if m_result:
            content = re.sub(r",(\r\n|\r|\n)", ",", content, count=0, flags=0)
            for line in StringIO(content):
                self.__m3u8_attr__(line)
        else:
            return '', M3U8_STATUS.TYPE_ERROR

    def __m3u8_attr__(self, content):
        '''
            设置M3U8 TAG值
        '''
        content = re.sub(r"^\r\n|\r|\n", "", content, count=0, flags=0)
        m_result = re.match(r"^#.*", content)
        if m_result:
            m_data = content.split(":", 1)
            m_key = None
            if len(m_data) == 2:
                m_key = m_data[0]
            else:
                m_key = content
            m_key += "__"
            m_key = re.sub(r"#|-", "_", m_key, count=0, flags=0)
            m_key = m_key.replace(" ", "")
            if len(m_data) == 2:
                if m_key == "_EXTINF__":
                    self._EXTINF__.append(m_data[1])
                else:
                    self.__setattr__(m_key, m_data[1])
            else:
                self.__setattr__(m_key, True)
        else:
            if re.search(r"\.m3u8", content):
                self._EXT_X_M3U8_FILE_LIST__.append(content)

    def __match__(self, regex, content):
        '''
            验证字符串是否存在
        '''
        m_result = re.match(regex, content, re.I | re.M)
        if m_result:
            return True
        return False

    def __request__(self, url):
        '''
            requests
        '''
        if self._http:
            try:
                for i in range(0, self._reconnect+1):
                    response, status = self._http.getcontent(
                        url, p_timeout=self._timeout)
                    if status == 200:
                        return response, status
                if i+1 >= self._reconnect:
                    return 'timeout', 408
            except Exception as e:
                return e, -1
        else:
            return 'http is none.', -1
