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

import os
import yaml
from aiohttp import web
from telethon import TelegramClient
from telethon.utils import _get_file_info

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')
port = os.environ.get('PORT', 8080)

async def handler(request):
    query = request.query
    if authorized_tokens:
        if 'token' not in query:
            return web.Response(status=401, text='Unauthorized')
        if query['token'] not in authorized_tokens:
            return web.Response(status=403, text='Forbidden')
    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:
        pass
    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')
    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')
    http_range = request.http_range
    offset = http_range.start or 0
    end = http_range.stop
    max_size = 0
    for i in messages:
        max_size += _get_file_info(i).size
    if end is None:
        end = max_size
    elif end > max_size:
        return web.Response(status=416, text='Range end size is bigger than file sizes', headers={'Content-Range': f'bytes */{max_size}'})
    length = end - offset

    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, msg_data=(chat_id, i.id)):
                yield chunk[:tmp_length]
                tmp_length -= len(chunk)
                if tmp_length < 1:
                    break

    return web.Response(status=206 if (end - offset != max_size) else 200,
        body=download(),
        headers={
            'Content-Range': f'bytes {offset}-{end}/{max_size}',
            'Content-Length': str(end - offset),
            'Accept-Ranges': 'bytes'
        }
    )

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, '127.0.0.1', port)
    await site.start()
    await client.run_until_disconnected()

client.loop.run_until_complete(main())