#!/usr/bin/env python3
import json
import sys
import os
import re
import base64
import logging
import distutils.util
import asyncio
import tornado.web
import tornado.routing
import stream_providers
import aiohttp

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

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

nextcloud_server = os.environ.get("NEXTCLOUD_SERVER")
if nextcloud_server is not None:
    providers["nextcloud"] = nextcloud_server

seafile_server = os.environ.get("SEAFILE_SERVER")
if seafile_server is not None:
    providers["seafile"] = seafile_server

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

class ProxyElem():
    def __init__(self, proxy):
        self.proxy = proxy
    def local(self):
        timeout = aiohttp.ClientTimeout(total=5)
        return aiohttp.ClientSession(timeout=timeout)
    def __repr__(self):
        return str(self.proxy)
    async def proxy_url(self, urls):
        if not isinstance(proxy_server, str):
            return urls
        jdata = None
        data_list = []
        for url in urls:
            data = {}
            data["upstream"] = url
            data["proxy"] = self.proxy
            data["proxied"] = isinstance(self.proxy, str)
            data_list.append(data)
        try:
            async with self.local() as session:
                resp = await session.post(proxy_server, json=data_list)
                text = await resp.text()
                jdata = json.loads(text)
        except Exception as e:
            logger.info(e)
        if isinstance(jdata, list):
            ret_data = []
            for src, dst in zip(urls, jdata):
                if isinstance(src, str):
                    ret_data.append(dst)
                else:
                    ret_data.append(None)
            return ret_data
        else:
            return urls

proxies = {}
new_providers = {}
for key in providers:
    expr = re.compile(f'^{key}_proxy(_[a-z][a-z])?[0-9]+$', re.IGNORECASE)
    matches = list(filter(expr.match, os.environ.keys()))
    current = []
    current_keys = set()
    current_keys.add(key)
    countries = []
    empty = True
    for match in matches:
        country_groups = expr.match(match.lower()).groups()
        country = None
        pos = len(country_groups) - 1
        if pos >= 0:
            country_temp = country_groups[pos]
            if isinstance(country_temp, str):
                country = country_temp.strip("_")
                current_keys.add(f'{key}_{country}')
        proxy = os.environ.get(match)
        if proxy is not None:
            current.append(proxy)
            countries.append(country)
            if country is None:
                empty = False
    for elem in current_keys:
        proxies[elem] = []
        new_providers[elem] = providers[key]
    print(proxies)
    for proxy, country in zip(current, countries):
        new_key = key
        if country is not None:
            new_key = f'{key}_{country}'
        proxies[new_key].append(ProxyElem(proxy))
    for elem in current_keys:
        if len(proxies[elem]) == 0:
            proxies[elem].append(ProxyElem(None))
providers = new_providers

proxy_keys = []
for proxy_provider in proxies.values():
    for proxy in proxy_provider:
        if isinstance(proxy, ProxyElem) and isinstance(proxy.proxy, str):
            proxy_keys.append(proxy.proxy)
stream_providers.setup(proxy_keys)

class UpstreamHandler():
    def __init__(self):
        self.provider = None
        self.raw = False
        self.valid = False
        self.proxy = ProxyElem(None)
        self.upstream = None
    async def test_socks(self, proxy):
        if not hasattr(proxy, "proxy") or not isinstance(proxy.proxy, str):
            return (True, ProxyElem(None))
        try:
            splitted = proxy.proxy.rsplit(":", 1)
            host = proxy
            port = 1080
            if len(splitted) == 2:
                host = splitted[0]
                port = splitted[1]
            future = asyncio.open_connection(host=host, port=port)
            await asyncio.wait_for(future, timeout=1)
        except Exception as e:
            return (False, proxy)
        else:
            return (True, proxy)
    async def setup(self, handler):
        self.provider = handler.get_query_argument("provider", None)
        raw_str = handler.get_query_argument("raw", None)
        if isinstance(raw_str, str):
            try:
                self.raw = bool(distutils.util.strtobool(raw_str))
            except ValueError as e:
                logger.info(e)
        if self.provider in providers.keys():
            self.valid = True
            path = handler.request.path
            if self.provider.startswith("nextcloud"):
                path = path.removesuffix("/").removesuffix("download").removesuffix("/")
            elif self.provider.startswith("youtube"):
                path = path.removeprefix("/")
            self.upstream = providers[self.provider] + path
            proxy_list = proxies.get(self.provider)
            if isinstance(proxy_list, list):
                futures = []
                for current in proxy_list:
                    future = asyncio.create_task(self.test_socks(current))
                    futures.append(future)
                for future in asyncio.as_completed(futures):
                    success, current = await future
                    if success:
                        self.proxy = current
                        break
                for future in futures:
                    if not future.done():
                        future.cancel()

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_script = None
videojs_version = None
font_awesome_version = None
custom_style = None
favicon = 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_script = tornado.template.Template(f.read().strip())
    with open("/app/version/video.js.txt", "r") as f:
        videojs_version = f.read().strip()
    with open("/app/version/chromecast.txt", "r") as f:
        chromecast_version = f.read().strip()
    with open("/app/version/font-awesome.txt", "r") as f:
        font_awesome_version = f.read().strip()
    with open("/app/favicon.png", "rb") as f:
        favicon = f.read()
    with open("/app/style.css", "r") as f:
        custom_style = f.read()
except Exception as e:
    logger.info(e)

class MainHandler(tornado.web.RequestHandler):
    async def handle_any(self, redir):
        handler = UpstreamHandler()
        await handler.setup(self)
        if handler.valid:
            if handler.raw:
                await self.handle_raw(handler)
            else:
                await self.handle_render(handler)
        else:
            logger.info(f'provider missing {self.request.uri}')
            self.set_status(404)
            self.write("Stream not found. (provider missing)")
    async def get_data(self, handler):
        video_info = None
        meta = None
        title = None
        if template_script is not None and template_html is not None:
            provider_data = None
            if handler.provider.startswith("nextcloud"):
                provider_data = await stream_providers.get_nextcloud(handler.upstream, handler.proxy, logger)
            elif handler.provider.startswith("seafile"):
                provider_data = await stream_providers.get_seafile(handler.upstream, handler.proxy, logger)
            else:
                provider_data = await stream_providers.get_any(handler.upstream, handler.proxy, logger)
            proxied = await handler.proxy.proxy_url([provider_data.upstream(), provider_data.thumbnail()])
            video_info = {}
            video_info["upstream"] = proxied[0]
            video_info["poster"] = proxied[1]
            video_info["ctype"] = provider_data.ctype()
            meta = provider_data.meta()
            title = provider_data.title()
        return (video_info, meta, title)
    async def handle_raw(self, handler):
        video_info, meta, title = await self.get_data(handler)
        if video_info is not None:
            self.redirect(url=video_info["upstream"], status=303)
        else:
            self.set_status(404)
            self.write("HTML template missing.")
    async def handle_render(self, handler):
        video_info, meta, title = await self.get_data(handler)
        if video_info is not None:
            script = template_script.generate(info=json.dumps(video_info))
            b64 = str(base64.b64encode(script), "ascii")
            script_file = f'data:text/javascript;charset=utf-8;base64,{b64}'
            data["script"] = script_file
            data["videojs_version"] = videojs_version
            data["chromecast_version"] = chromecast_version
            data["font_awesome_version"] = font_awesome_version
            rendered_html = template_html.generate(data=data, meta=meta, title=title)
            self.write(rendered_html)
        else:
            self.set_status(404)
            self.write("HTML template missing.")
    async def get(self):
        await self.handle_any(True)
    async def head(self):
        await self.handle_any(False)

class PlaylistHandler(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", "image/png")
        self.write(favicon)

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