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

providers = {}
providers["nrk"] = "https://tv.nrk.no"
providers["svt"] = "https://svtplay.se"
providers["youtube"] = "https://youtu.be"
providers["twitch"] = "https://twitch.tv"

proxies = {}
proxies_raw = {}
proxies_req = {}

for key in providers:
    proxies[key] = streamlink.Streamlink()
    proxy = os.environ.get(key.upper() + "_PROXY")
    proxies[key].set_option("http-timeout", 2.0)
    if proxy is not None:
        proxies[key].set_option("https-proxy", "socks5://" + proxy)
        proxies_raw[key] = proxy
        proxies_req[key] = {}
        proxies_req[key]["http"] = "socks5://" + proxy
        proxies_req[key]["https"] = "socks5://" + 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:
    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"

template = None
try:
    with open("/app/index.html", "r") as f:
        template = tornado.template.Template(f.read().strip())
except Exception as e:
    print(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 rewrite(current, provider):
    proxy_req = proxies_req.get(provider)
    proxy = proxies_raw.get(provider)
    resp = requests.head(current, proxies=proxy_req)
    ctype = resp.headers.get("Content-Type")
    if ctype is None:
        return None
    else:
        if "mpegurl" not in ctype.lower():
            return None
    resp = requests.get(current, proxies=proxy_req)
    ndata = None
    if resp.text is not None:
        ndata = ""
        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 = get_proxy_url(proxy, current, matches[0])
                    ndata += re.sub(r'URI=".+"', f'URI="{new_url}"', line)
            elif line.startswith("#"):
                ndata += line
            else:
                ndata += get_proxy_url(proxy, current, line)
            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)
        if isinstance(provider, str):
            if isinstance(render, str):
                if render.lower() == "true":
                    self.handle_render(provider, write)
                else:
                    self.handle_stream(provider, write)
            else:
                self.handle_stream(provider, write)
        else:
            self.set_status(404)
            if write:
                self.write("Stream not found.")
    
    def handle_render(self, provider, write):
        if template is not None:
            stream_path = f'{self.request.path}?provider={provider}'
            rendered = template.generate(stream=stream_path)
            self.write(rendered)
        else:
            self.set_status(404)
            self.write("HTML template missing.")

    def handle_stream(self, provider, write):
        upstream = None
        proxy = None
        if provider in providers.keys():
            proxy = proxies.get(provider)
            src = providers[provider] + self.request.path
            try:
                resp = requests.head(src, allow_redirects=True)
                if resp is not None:
                    src = resp.url
                    if isinstance(src, str) and "consent.youtube.com" in src:
                        video_id = self.request.path.strip("/")
                        src = f'https://www.youtube.com/watch?v={video_id}'
            except Exception as e:
                print(e)
            else:
                try:
                    streams = proxy.streams(src)
                    for key in reversed(streams):
                        stream = streams.get(key)
                        if hasattr(stream, "url"):
                            upstream = stream.url
                            break
                except Exception as e:
                    print(str(e))
                    upstream = None
        if upstream is None:
            self.set_status(404)
            if write:
                self.write("Stream not found.")
        else:
            data = rewrite(upstream, provider)
            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)

try:
    handlers = []
    handlers.append((tornado.routing.PathMatches("/sources.m3u8"), FileHandler))
    handlers.append((tornado.routing.AnyMatches(), MainHandler))
    app_web = tornado.web.Application(handlers)
    app_web.listen(8080)
    tornado.ioloop.IOLoop.current().start()
except KeyboardInterrupt:
    print()