import logging
logging.basicConfig(level=logging.INFO)

import os
import re
import yaml
import hmac
from aiohttp import web
from telethon import TelegramClient
from telethon.utils import _get_file_info
from telethon.client.downloads import _GenericDownloadIter

with open('config.yaml') as file:
    config = yaml.safe_load(file)

client = TelegramClient('streamtg', config['telegram']['api_id'], config['telegram']['api_hash'])
authorized_tokens = config.get('authorized_tokens')
hmacs = [hmac.new(i['key'].encode(), digestmod=i['digest']) for i in config.get('hmac', ())]
port = os.environ.get('PORT', 8080)
bytes_regex = re.compile(r'^(\d*)-(\d*)$')

def verify_token(token):
    return token in authorized_tokens

def verify_hmac(hexdigest, chat_id, message_ids):
    text = f'{chat_id}|{"|".join(message_ids)}'.encode()
    for i in hmacs:
        i = i.copy()
        i.update(text)
        if hmac.compare_digest(hexdigest, i.hexdigest()):
            return True
    return False

# saftea
def string_is_printable_ascii(string):
    return not any(True for i in string.encode('ascii') if i < b'!' or i > b'~')

async def handler(request):
    query = request.query
    content_type = query.get('Content-Type')
    if content_type and not string_is_printable_ascii(content_type):
        return web.Response(status=400, text='Content-Type has a blacklisted character')
    content_disposition = query.get('Content-Disposition')
    if content_disposition and not string_is_printable_ascii(content_disposition):
        return web.Response(status=400, text='Content-Disposition has a blacklisted character')
    token = query.get('token')
    hexdigest = query.get('hmac')
    if not token and not hexdigest and (authorized_tokens or hmacs):
        return web.Response(status=401, text='Missing token or hmac')
    if 'chat_id' not in query:
        return web.Response(status=400, text='Missing chat_id')
    chat_id = query['chat_id']
    try:
        chat_id = int(chat_id)
    except ValueError:
        try:
            chat_id = await client.get_peer_id(chat_id)
        except BaseException:
            if authorized_tokens or hmacs:
                logging.exception('Exception occured while getting chat id of %s, returning 403 to hide known chats', chat_id)
                return web.Response(status=403, text='Forbidden')
            raise
    if 'message_id' not in query:
        return web.Response(status=400, text='Missing message_id')
    message_ids = query.getall('message_id')
    if any(True for i in message_ids if not i.isnumeric() or i == '0'):
        return web.Response(status=400, text='Invalid message_id')
    if authorized_tokens or hmacs:
        if not token or not verify_token(token):
            if hexdigest:
                if not verify_hmac(hexdigest, chat_id, message_ids):
                    return web.Response(status=403, text='Forbidden')
            else:
                return web.Response(status=403, text='Forbidden')
    message_ids = list(map(int, message_ids))
    messages = await client.get_messages(chat_id, ids=message_ids)
    if any(True for i in messages if i is None):
        return web.Response(status=400, text='At least one of the messages does not exist')
    if any(True for i in messages if not i.media):
        return web.Response(status=400, text='At least one of the messages do not contain media')
    max_size = 0
    for i in messages:
        max_size += _get_file_info(i).size
    status = 200
    offset = 0
    end = max_size - 1
    if range := request.headers.get('Range'):
        if range.startswith('bytes='):
            match = bytes_regex.match(range[6:])
            if not match:
                return web.Response(status=416,
                    text='Failed to match range with regex',
                    headers={
                        'Content-Range': f'bytes */{max_size}',
                        'Accept-Ranges': 'bytes'
                    }
                )
            tmp_offset = match.group(1)
            tmp_end = match.group(2)
            if tmp_offset:
                offset = int(tmp_offset)
                if tmp_end:
                    end = int(tmp_end)
            elif tmp_end:
                offset = max_size - int(tmp_end)
            else:
                return web.Response(status=400,
                    text='Both range sides missing',
                    headers={
                        'Content-Range': f'bytes */{max_size}',
                        'Accept-Ranges': 'bytes'
                    }
                )
            if offset < 0 or offset > max_size or end >= max_size or offset > end:
                return web.Response(status=416,
                    text='Range not satisfiable',
                    headers={
                        'Content-Range': f'bytes */{max_size}',
                        'Accept-Ranges': 'bytes'
                    }
                )
            status = 206
        else:
            return web.Response(status=416,
                text='Unknown range type',
                headers={
                    'Content-Range': f'bytes */{max_size}',
                    'Accept-Ranges': 'bytes'
                }
            )
    length = end - offset + 1
    headers = {
        'Content-Range': f'bytes {offset}-{end}/{max_size}',
        'Content-Length': str(length),
        'Accept-Ranges': 'bytes'
    }
    if content_type:
        headers['Content-Type'] = content_type
    if content_disposition:
        headers['Content-Disposition'] = content_disposition

    async def download():
        tmp_offset = offset
        tmp_length = length
        for i in messages:
            if tmp_length < 1:
                break
            size = _get_file_info(i).size
            if tmp_offset > size:
                tmp_offset -= size
                continue
            async for chunk in client._iter_download(i, offset=tmp_offset, limit=131072, msg_data=(chat_id, i.id)):
                yield chunk[:tmp_length]
                tmp_length -= len(chunk)
                if tmp_length < 1:
                    break
            tmp_offset = 0

    return web.Response(status=status,
        body=download(),
        headers=headers
    )

app = web.Application()
app.add_routes([web.get('/', handler)])

async def main():
    await client.start(bot_token=config['telegram'].get('bot_token'))
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, '0.0.0.0', port)
    await site.start()
    await client.run_until_disconnected()

client.loop.run_until_complete(main())