#!/usr/bin/env python3
import json
import urllib.parse
import re
import os
import streamlink
import tornado.web
import tornado.routing
import requests
import base64
import logging

logging.basicConfig(format='[%(filename)s:%(lineno)d] %(message)s', level=logging.INFO)
logger = logging.getLogger(__name__)

providers = {}
providers["nrk"] = "https://tv.nrk.no"
providers["svt"] = "https://svtplay.se"
providers["youtube"] = "https://www.youtube.com/watch?v="
providers["twitch"] = "https://twitch.tv"

class ProxyElem():
    def __init__(self, proxy):
        self.proxy = proxy
        self.req = {}
        if proxy is not None:
            self.req["http"] = "socks5://" + proxy
            self.req["https"] = "socks5://" + proxy
    def stream(self):
        session = streamlink.Streamlink()
        session.set_option("http-timeout", 2.0)
        if self.proxy is not None:
            session.set_option("https-proxy", "socks5://" + self.proxy)
            session.set_option("http-proxy", "socks5://" + self.proxy)
        return session
    def __repr__(self):
        return str(self.proxy)
    
proxies = {}
for key in providers:
    proxies[key] = []
    current = []
    for i in range(0,9):
        proxy = os.environ.get(f'{key}_proxy{i}'.upper())
        if proxy is not None:
            current.append(proxy)
    if len(current) == 0:
        proxies[key].append(ProxyElem(None))
    else:
        for proxy in current:
            proxies[key].append(ProxyElem(proxy))

playlist = None
icecast_server = os.environ.get("ICECAST_SERVER")
stream_server = os.environ.get("STREAM_SERVER")
proxy_server = os.environ.get("PROXY_SERVER")

if icecast_server is not None and stream_server is not None:
    try:
        with open("/app/sources.json", "r") as f:
            data = json.loads(f.read())
            playlist = "#EXTM3U\n"
            for key in data:
                current = data[key]
                name = current["name"]
                radio = current["radio"]
                if radio:
                    playlist += f'#EXTINF:0 radio="true", {name}\n'
                    playlist += icecast_server + key + "\n"
                else:
                    playlist += f'#EXTINF:0 radio="false", {name}\n'
                    playlist += stream_server + key + "\n"
    except Exception as e:
        logger.info(e)

template_html = None
template_js = None
template_embed = tornado.template.Template('<iframe src="{{ origin }}?provider={{ provider }}&render=true"></iframe>')

videojs_version = None
castjs_version = None
custom_style = None
try:
    with open("/app/index.html", "r") as f:
        template_html = tornado.template.Template(f.read().strip())
    with open("/app/script.js", "r") as f:
        template_js = tornado.template.Template(f.read().strip())
    with open("/app/videojs-version.txt", "r") as f:
        videojs_version = f.read().strip()
    with open("/app/chromecast-version.txt", "r") as f:
        chromecast_version = f.read().strip()
    with open("/app/style.css", "r") as f:
        custom_style_raw = bytes(f.read().strip(), "utf-8")
        b64 = str(base64.b64encode(custom_style_raw), "ascii")
        custom_style = f'data:text/css;charset=utf-8;base64,{b64}'
except Exception as e:
    logger.info(e)

def get_proxy_url(proxy, current, path):
    data = {}
    data["upstream"] = urllib.parse.urljoin(current, path)
    data["proxied"] = True
    ret = None
    if proxy is None:
        data["proxied"] = False
    else:
        data["proxy"] = proxy
    if proxy_server is None:
        return data["upstream"]
    presp = requests.post(proxy_server, json=data)
    return presp.text

def upstream_type(current, proxy):
    if proxy is None:
        resp = requests.head(current)
    else:
        resp = requests.head(current, proxies=proxy.req)
    return resp.headers.get("Content-Type", "binary/octet-stream")
 
def rewrite(current, provider, proxy):
    resp = requests.get(current, proxies=proxy.req)
    ndata = None
    if resp.text is not None:
        links = []
        for line in resp.text.splitlines():
            if line.startswith("#EXT-X-KEY:METHOD="):
                matches = re.findall(r'(?<=URI=").+(?=")', line)
                if len(matches) == 1:
                    ldata = {}
                    ldata["upstream"] = urllib.parse.urljoin(current, matches[0])
                    ldata["proxy"] = proxy.proxy
                    ldata["proxied"] = isinstance(proxy.proxy, str)
                    links.append(ldata)
            elif line.startswith("#"):
                pass
            else:
                ldata = {}
                ldata["upstream"] = urllib.parse.urljoin(current, line)
                ldata["proxy"] = proxy.proxy
                ldata["proxied"] = isinstance(proxy.proxy, str)
                links.append(ldata)
        if isinstance(proxy_server, str):
            ndata = ""
            presp = requests.post(proxy_server, json=links)
            if isinstance(presp.text, str):
                links = json.loads(presp.text)
            for line in resp.text.splitlines():
                if line.startswith("#EXT-X-KEY:METHOD="):
                    matches = re.findall(r'(?<=URI=").+(?=")', line)
                    if len(matches) == 1:
                        new_url = links.pop(0)
                        ndata += re.sub(r'URI=".+"', f'URI="{new_url}"', line)
                elif line.startswith("#"):
                    ndata += line
                else:
                    ndata += links.pop(0)
                ndata += "\n"
        return ndata

class MainHandler(tornado.web.RequestHandler):
    def handle_any(self, write):
        provider = self.get_query_argument("provider", None)
        render = self.get_query_argument("render", "false")
        embed = self.get_query_argument("embed", "false")
        if isinstance(provider, str):
            if render.lower() == "true":
                self.handle_render(provider)
            elif embed.lower() == "true":
                self.handle_embed(provider)
            else:
                self.handle_stream(provider, write)
        else:
            logger.info(f'provider missing {self.request.uri}')
            self.set_status(404)
            if write:
                self.write("Stream not found. (provider missing)")
    
    def handle_render(self, provider):
        if template_js is not None and template_html is not None:
            origin = self.request.path
            if stream_server is not None:
                origin = f'{stream_server}{self.request.path}'
            stream_path = f'{origin}?provider={provider}'
            rendered_js = template_js.generate(stream=stream_path);
            b64_js = str(base64.b64encode(rendered_js), "ascii")
            script = f'data:text/javascript;charset=utf-8;base64,{b64_js}'
            rendered_html = template_html.generate(script=script, videojs_version=videojs_version, chromecast_version=chromecast_version, custom_style=custom_style, provider=provider, origin=origin)
            self.write(rendered_html)
        else:
            self.set_status(404)
            self.write("HTML template missing.")
    def handle_embed(self, provider):
        max_width = self.get_query_argument("maxwidth", "320")
        max_height = self.get_query_argument("maxheight", "180")
        origin = self.request.path
        if stream_server is not None:
            origin = f'{stream_server}{self.request.path}'
        embed_json = {}
        embed_json["version"] = "1.0"
        embed_json["type"] = "video"
        embed_json["width"] = max_width
        embed_json["height"] = max_height
        embed_json["html"] = str(template_embed.generate(origin=origin, provider=provider), "utf-8")
        self.set_header("Content-Type", "application/json; charset=utf-8")
        self.write(json.dumps(embed_json))

    def handle_stream(self, provider, write):
        upstream = None
        proxy = None
        if provider in providers.keys():
            path = self.request.path
            if provider == "youtube":
                path = path.strip("/")
            src = providers[provider] + path
            proxy_list = None
            proxy_iter = None
            proxy_list_orig = proxies.get(provider)
            if isinstance(proxy_list_orig, list):
                proxy_list = proxy_list_orig.copy()
                proxy_iter = proxy_list_orig.copy()
            if isinstance(proxy_list, list):
                for i in proxy_iter:
                    current_list = proxy_list.copy()
                    current = proxy_list.pop()
                    proxy_list = [current] + proxy_list
                    try:
                        resp = requests.head(src, allow_redirects=True, proxies=current.req, timeout=2)
                        if resp is not None:
                            logger.info(src)
                            src = resp.url
                    except Exception as e:
                        logger.info(e)
                    else:
                        proxies[provider] = current_list
                        proxy = current
                        break
            if proxy is not None:
                try:
                    logger.info(proxy)
                    streams = proxy.stream().streams(src)
                except Exception as e:
                    logger.info(e)
                else:
                    for key in reversed(streams):
                        stream = streams.get(key)
                        logger.info(stream)
                        if hasattr(stream, "url"):
                            upstream = stream.url
                            break
        else:
            logger.info(f'invalid provider ({provider})')
            self.set_status(404)
            if write:
                self.write("Stream not found. (invalid provider)")
            return
        if upstream is None:
            logger.info(f'invalid upstream ({provider})')
            self.set_status(404)
            if write:
                self.write("Stream not found. (invalid upstream)")
        else:
            ctype = upstream_type(upstream, proxy)
            data = None
            if "mpegurl" in ctype.lower():
                data = rewrite(upstream, provider, proxy)
            else:
                ldata = {}
                ldata["upstream"] = upstream
                ldata["proxy"] = proxy.proxy
                ldata["proxied"] = isinstance(proxy.proxy, str)
                links = [ldata]
                if isinstance(proxy_server, str):
                    try:
                        resp = requests.post(proxy_server, json=links)
                        if isinstance(resp.text, str):
                            new_links = json.loads(resp.text)
                            if isinstance(new_links, list) and len(new_links) == 1:
                                upstream = new_links.pop()
                    except Exception as e:
                        logger.info(e)
            if data is None:
                self.redirect(upstream, status=303)
            else:
                self.set_header("Content-Type", "application/vnd.apple.mpegurl")
                self.write(data)
    def get(self):
        self.handle_any(True)
    def head(self):
        self.handle_any(False)

class FileHandler(tornado.web.RequestHandler):
    def get(self):
        self.set_header("Content-Type", "text/plain; charset=utf-8")
        self.write(playlist)

class IconHandler(tornado.web.RequestHandler):
    def get(self):
        self.set_header("Content-Type", "text/plain; charset=utf-8")
        self.set_status(204)
try:
    handlers = []
    handlers.append((tornado.routing.PathMatches("/sources.m3u8"), FileHandler))
    handlers.append((tornado.routing.PathMatches("/favicon.ico"), IconHandler))
    handlers.append((tornado.routing.AnyMatches(), MainHandler))
    app_web = tornado.web.Application(handlers)
    app_web.listen(8080)
    tornado.ioloop.IOLoop.current().start()
except KeyboardInterrupt:
    print()