#!/usr/bin/env python3
import os
import sys
import glob
import json
import time
import datetime
import random
import logging
import numpy as np
import configparser
import argparse

from lib import game, model, mcts

MCTS_SEARCHES = 20
MCTS_BATCH_SIZE = 4

try:
    import telegram.ext
    from telegram.error import TimedOut
except ImportError:
    print("Wymagany pakiet python-telegram-bot")
    sys.exit()

import torch

# Plik konfiguracyjny
# [telegram]
# api=API_KEY
CONFIG_DEFAULT = "~/.config/rl_ch18_bot.ini"

log = logging.getLogger("telegram")


class Session:
    BOT_PLAYER = game.PLAYER_BLACK
    USER_PLAYER = game.PLAYER_WHITE

    def __init__(self, model_file, player_moves_first, player_id):
        self.model_file = model_file
        self.model = model.Net(input_shape=model.OBS_SHAPE, actions_n=game.GAME_COLS)
        self.model.load_state_dict(torch.load(model_file, map_location=lambda storage, loc: storage))
        self.state = game.INITIAL_STATE
        self.value = None
        self.player_moves_first = player_moves_first
        self.player_id = player_id
        self.moves = []
        self.mcts_store = mcts.MCTS()

    def move_player(self, col):
        self.moves.append(col)
        self.state, won = game.move(self.state, col, self.USER_PLAYER)
        return won

    def move_bot(self):
        self.mcts_store.search_batch(MCTS_SEARCHES, MCTS_BATCH_SIZE, self.state, self.BOT_PLAYER, self.model)
        probs, values = self.mcts_store.get_policy_value(self.state, tau=0)
        action = np.random.choice(game.GAME_COLS, p=probs)
        self.value = values[action]
        self.moves.append(action)
        self.state, won = game.move(self.state, action, self.BOT_PLAYER)
        return won

    def is_valid_move(self, move_col):
        return move_col in game.possible_moves(self.state)

    def is_draw(self):
        return len(game.possible_moves(self.state)) == 0

    def render(self):
        l = game.render(self.state)
        l = "\n".join(l)
        l = l.replace("0", 'O').replace("1", "X")
        board = "0123456\n-------\n" + l + "\n-------\n0123456"
        extra = ""
        if self.value is not None:
            extra = "Ewaluacja pozycji: %.2f\n" % float(self.value)
        return extra + "<pre>%s</pre>" % board


class PlayerBot:
    def __init__(self, models_dir, log_file):
        self.sessions = {}
        self.models_dir = models_dir
        self.models = self._read_models(models_dir)
        self.log_file = log_file
        self.leaderboard = {}
        self._read_leaderboard(log_file)

    def _read_models(self, models_dir):
        result = {}
        for idx, name in enumerate(sorted(glob.glob(os.path.join(models_dir, "*.dat")))):
            result[idx] = name
        return result

    def _read_leaderboard(self, log_file):
        if not os.path.exists(log_file):
            return 
        with open(log_file, 'rt', encoding='utf-8') as fd:
            for l in fd:
                data = json.loads(l)
                bot_name = os.path.basename(data['model_file'])
                user_name = data['player_id'].split(':')[0]
                bot_score = data['bot_score']
                self._update_leaderboard(bot_score, bot_name, user_name)

    def _update_leaderboard(self, bot_score, bot_name, user_name):
        if bot_score > 0.5:
            game.update_counts(self.leaderboard, bot_name, (1, 0, 0))
            game.update_counts(self.leaderboard, user_name, (0, 1, 0))
        elif bot_score < -0.5:
            game.update_counts(self.leaderboard, bot_name, (0, 1, 0))
            game.update_counts(self.leaderboard, user_name, (1, 0, 0))
        else:
            game.update_counts(self.leaderboard, bot_name, (0, 0, 1))
            game.update_counts(self.leaderboard, user_name, (0, 0, 1))

    def _save_log(self, session, bot_score):
        self._update_leaderboard(bot_score, os.path.basename(session.model_file),
                                 session.player_id.split(':')[0])
        data = {
            "ts": time.time(),
            "time": datetime.datetime.utcnow().isoformat(),
            "bot_score": bot_score,
            "model_file": session.model_file,
            "player_id": session.player_id,
            "player_moves_first": session.player_moves_first,
            "moves": session.moves,
            "state": session.state
        }
        with open(self.log_file, "a+t", encoding='utf-8') as f:
            f.write(json.dumps(data, sort_keys=True) + '\n')

    def command_help(self, bot, update):
        bot.send_message(chat_id=update.message.chat_id, parse_mode="HTML", disable_web_page_preview=True,
                         text="""
This a <a href="https://en.wikipedia.org/wiki/Connect_Four">4-in-a-row</a> game bot trained with AlphaGo Zero method for the <a href="https://www.packtpub.com/big-data-and-business-intelligence/practical-deep-reinforcement-learning">Practical Deep Reinforcement Learning</a> book. 

<b>Witamy!</b>

Dostępne następujące polecenia:
<b>/list</b> - lista wstępnie wytrenowanych modeli (im większy ID, tym lepiej gra)
<b>/play MODEL_ID</b> - rozpoczęcie gry z określonym modelem
<b>/top</b> - tabela wyników

Podczas gry należy podawać liczby, oznaczające kolumny, w których można umieszczać żetony.
""")


    def command_list(self, bot, update):
        if len(self.models) == 0:
            reply = ["Nie ma odpowiednich modeli!"]
        else:
            reply = ["Lista dostępnych modeli i ich identyfikatory"]
            for idx, name in sorted(self.models.items()):
                reply.append("<b>%d</b>: %s" % (idx, os.path.basename(name)))

        bot.send_message(chat_id=update.message.chat_id, text="\n".join(reply), parse_mode="HTML")

    def command_play(self, bot, update, args):
        chat_id = update.message.chat_id
        player_id = "%s:%s" % (update.message.from_user.username, update.message.from_user.id)
        try:
            model_id = int(args[0])
        except ValueError:
            bot.send_message(chat_id=chat_id, text="Błędne argumenty! Użyj '/play <MODEL_ID>, by rozpocząć grę")
            return

        if model_id not in self.models:
            bot.send_message(chat_id=chat_id, text="Nie ma takiego modelu, użyj /list, by wyświetlić listę dostępnych modeli")
            return

        if chat_id in self.sessions:
            bot.send_message(chat_id=chat_id, text="Jeśli przerwiesz grę, stracisz wyniki")
            del self.sessions[chat_id]

        player_moves = random.choice([False, True])
        session = Session(self.models[model_id], player_moves, player_id)
        self.sessions[chat_id] = session
        if player_moves:
            bot.send_message(chat_id=chat_id, text="Ty rozpoczynasz (używam O), podaj nr kolumny, w której należy umieścić żeton - liczbę od 0 do 6")
        else:
            bot.send_message(chat_id=chat_id, text="Ja rozpoczynam (używam X), wykonuję ruch...")
            session.move_bot()
        bot.send_message(chat_id=chat_id, text=session.render(), parse_mode="HTML")

    def text(self, bot, update):
        chat_id = update.message.chat_id

        if chat_id not in self.sessions:
            bot.send_message(chat_id=chat_id, text="Gra nieaktywna. Rozpocznij ją za pomocą polecenia<b>/play MODEL_ID</b> "
                                                   "(lub użyj <b>/help</b>, aby wyświetlić listę dostępnych poleceń)",
                             parse_mode='HTML')
            return
        session = self.sessions[chat_id]

        try:
            move_col = int(update.message.text)
        except ValueError:
            bot.send_message(chat_id=chat_id, text="Nie rozumiem. W trybie gry musisz podać liczbę "
                                                   "z zakresu od 0 do 6, aby wykonać ruch.")
            return

        if move_col < 0 or move_col > game.GAME_COLS:
            bot.send_message(chat_id=chat_id, text="Zły numer kolumny! Musisz podać wartość 0-6")
            return

        if not session.is_valid_move(move_col):
            bot.send_message(chat_id=chat_id, text="Ruch %d jest błędny!" % move_col)
            return

        won = session.move_player(move_col)
        if won:
            bot.send_message(chat_id=chat_id, text="Wygrałeś! Gratulacje!")
            self._save_log(session, bot_score=-1)
            del self.sessions[chat_id]
            return

        won = session.move_bot()
        bot.send_message(chat_id=chat_id, text=session.render(), parse_mode="HTML")

        if won:
            bot.send_message(chat_id=chat_id, text="Wygrałem! Hura!")
            self._save_log(session, bot_score=1)
            del self.sessions[chat_id]
        if session.is_draw():
            bot.send_message(chat_id=chat_id, text="Remis. Zdarza się to rzadko. 1:1, do zobaczenia!")
            self._save_log(session, bot_score=0)
            del self.sessions[chat_id]

    def error(self, bot, update, error):
        try:
            raise error
        except TimedOut:
            log.info("Błąd przekroczenia dopuszczalnego czasu")

    def command_top(self, bot, update):
        res = ["Tabela wyników"]
        items = list(self.leaderboard.items())
        items.sort(reverse=True, key=lambda p: p[1][0])
        for user, (wins, losses, draws) in items:
            res.append("%20s: wygrana=%d, przegrana=%d, remis=%d" % (user[:20], wins, losses, draws))
        l = "\n".join(res)
        bot.send_message(chat_id=update.message.chat_id, text="<pre>" + l + "</pre>", parse_mode="HTML")

    def command_refresh(self, bot, update):
        self.models = self._read_models(self.models_dir)
        bot.send_message(chat_id=update.message.chat_id, text="Wczytano modele, znaleziono %d plików" % len(self.models))


if __name__ == "__main__":
    logging.basicConfig(format="%(asctime)-15s %(levelname)s %(message)s", level=logging.INFO)
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default=CONFIG_DEFAULT,
                        help="Plik konfiguracyjny bota, domyślnie=" + CONFIG_DEFAULT)
    parser.add_argument("-m", "--models", required=True, help="Nazwa katalogu z modelami")
    parser.add_argument("-l", "--log", required=True, help="Nazwa pliku dziennika przechowującego gry i tabelę wyników")
    prog_args = parser.parse_args()

    conf = configparser.ConfigParser()
    if not conf.read(os.path.expanduser(prog_args.config)):
        log.error("Plik %s nie został znaleziony", prog_args.config)
        sys.exit()

    player_bot = PlayerBot(prog_args.models, prog_args.log)

    updater = telegram.ext.Updater(conf['telegram']['api'])
    updater.dispatcher.add_handler(telegram.ext.CommandHandler('help', player_bot.command_help))
    updater.dispatcher.add_handler(telegram.ext.CommandHandler('list', player_bot.command_list))
    updater.dispatcher.add_handler(telegram.ext.CommandHandler('top', player_bot.command_top))
    updater.dispatcher.add_handler(telegram.ext.CommandHandler('play', player_bot.command_play, pass_args=True))
    updater.dispatcher.add_handler(telegram.ext.CommandHandler('refresh', player_bot.command_refresh))
    updater.dispatcher.add_handler(telegram.ext.MessageHandler(telegram.ext.Filters.text, player_bot.text))
    updater.dispatcher.add_error_handler(player_bot.error)

    log.info("Bot zainicjalizowany, rozpoczęcie działania")
    updater.start_polling()
    updater.idle()

    pass
