# maubot - A plugin-based Matrix bot system.
# Copyright (C) 2022 Tulir Asokan
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
from datetime import datetime
import asyncio

from aiohttp import ClientSession, WSMessage, WSMsgType
from colorama import Fore
import click

from mautrix.types import Obj

from ..base import app
from ..config import get_token

history_count: int = 10


@app.command(help="View the logs of a server")
@click.argument("server", required=False)
@click.option("-t", "--tail", default=10, help="Maximum number of old log lines to display")
def logs(server: str, tail: int) -> None:
    server, token = get_token(server)
    if not token:
        return
    global history_count
    history_count = tail
    loop = asyncio.get_event_loop()
    loop.run_until_complete(view_logs(server, token))


def parsedate(entry: Obj) -> None:
    i = entry.time.index("+")
    i = entry.time.index(":", i)
    entry.time = entry.time[:i] + entry.time[i + 1 :]
    entry.time = datetime.strptime(entry.time, "%Y-%m-%dT%H:%M:%S.%f%z")


levelcolors = {
    "DEBUG": "",
    "INFO": Fore.CYAN,
    "WARNING": Fore.YELLOW,
    "ERROR": Fore.RED,
    "FATAL": Fore.MAGENTA,
}


def print_entry(entry: dict) -> None:
    entry = Obj(**entry)
    parsedate(entry)
    print(
        "{levelcolor}[{date}] [{level}@{logger}] {message}{resetcolor}".format(
            date=entry.time.strftime("%Y-%m-%d %H:%M:%S"),
            level=entry.levelname,
            levelcolor=levelcolors.get(entry.levelname, ""),
            resetcolor=Fore.RESET,
            logger=entry.name,
            message=entry.msg,
        )
    )
    if entry.exc_info:
        print(entry.exc_info)


def handle_msg(data: dict) -> bool:
    if "auth_success" in data:
        if data["auth_success"]:
            print(Fore.GREEN + "Connected to log websocket" + Fore.RESET)
        else:
            print(Fore.RED + "Failed to authenticate to log websocket" + Fore.RESET)
            return False
    elif "history" in data:
        for entry in data["history"][-history_count:]:
            print_entry(entry)
    else:
        print_entry(data)
    return True


async def view_logs(server: str, token: str) -> None:
    async with ClientSession() as session:
        async with session.ws_connect(f"{server}/_matrix/maubot/v1/logs") as ws:
            await ws.send_str(token)
            try:
                msg: WSMessage
                async for msg in ws:
                    if msg.type == WSMsgType.TEXT:
                        if not handle_msg(msg.json()):
                            break
                    elif msg.type == WSMsgType.ERROR:
                        print(Fore.YELLOW + "Connection error: " + msg.data + Fore.RESET)
                    elif msg.type == WSMsgType.CLOSE:
                        print(Fore.YELLOW + "Server closed connection" + Fore.RESET)
            except asyncio.CancelledError:
                pass