from gensim.models import KeyedVectors
from colorama import Fore, Back, Style
import argparse
import readline
import random
import time

import numpy as np
np.seterr(divide='ignore', invalid='ignore')

class Server():
    def __init__(self):
        pass

    def set_random_word(self):
        pass

    def get_rank(self, guess):
        pass

    def get_temp(self, guess):
        pass

    def _help(self):
        raise NotImplementedError

    def _reveal_word(self):
        raise NotImplementedError

class LocalServer(Server):
    def __init__(self, word=None, file="models/selected_word2vec_model.bin"):
        self.model = KeyedVectors.load_word2vec_format(
            file,
            binary=True,
            unicode_errors="ignore"
        )
        self.word = word
        self.nearest = []

    def init_word(self, k=1, dist=100):
        while (self.word is None or len(self.word) < 5
               or '-' in self.word or '_' in self.word):
            base_words = [
                self.model.index_to_key[random.randint(0, len(self.model))]
                for _ in range(k)
            ]

            complete_list = base_words.copy()
            for word in base_words:
                complete_list += [i[0] for i in self.model.most_similar(word, topn=dist)]

            rk_words = self.model.rank_by_centrality(complete_list)

            self.word = rk_words[random.randint(0,5)%len(rk_words)][1]
            self.nearest = [word]+[i[0] for i in self.model.most_similar(self.word, topn=1000)]

    def get_rank(self, guess):
        if guess not in self.nearest:
            return None
        return 1000 - self.nearest.index(guess)

    def get_temp(self, guess):
        return round(self.model.distance(self.word, guess)*100, 2)

    def _help(self, rk):
        return self.nearest[rk]

    def _reveal_word(self):
        return self.word


def cemantix(server: Server):
    server.init_word()

    guesses = [] # guess, temp, rank
    def formatted_status(guesses, last=None):
        text = ""
        for w, temp, rank in guesses:
            if rank is not None:
                text += Back.RED+Fore.BLACK+Style.BRIGHT+str(rank)+Style.RESET_ALL
            text += "\t"
            if temp >= 75:
                text += Fore.BLUE
            text += str(temp) + Style.RESET_ALL + "\t"
            if last == w:
                text += Style.BRIGHT+Back.WHITE+Fore.BLACK
            text += w + Style.RESET_ALL+"\n"
        return text[:-1]

    def tried(word, guessed):
        return (word in [i[0] for i in guessed])

    def interpret_command(cmd, guesses):
        match cmd:
            case "clear":
                guesses = [g for g in guesses if g[1] <= 75.]
            case "help":
                try:
                    best_rk = max([rk for _, _, rk in guesses if rk is not None]+[749])
                    print("Maybe try "+Back.YELLOW+Fore.BLACK+server._help(999-best_rk)
                        +Style.RESET_ALL)
                except NotImplementedError:
                    print(Fore.RED+"No help available"+Style.RESET_ALL)
            case _:
                print(Fore.RED+"Unknown command"+Style.RESET_ALL)

        return guesses


    while True:
        try:
            guess = input(Style.BRIGHT+"Your guess > "+Style.RESET_ALL).strip()
            readline.add_history(guess)
            if guess.endswith("()"):
                guesses = interpret_command(guess[:-2], guesses)
                continue
        except (EOFError, KeyboardInterrupt):
            try:
                print("The word was "+Style.BRIGHT+server._reveal_word()+Style.RESET_ALL)
            except NotImplementedError:
                pass
            print("Goodbye!")
            return -1
        try:
            dist = server.get_temp(guess)
        except KeyError:
            print(Fore.RED+"Key not present"+Style.RESET_ALL)
            continue
        
        if not tried(guess, guesses):
            guesses.append((guess, dist, server.get_rank(guess)))
            guesses.sort(key=lambda x:-x[1])
        print(chr(27) + "[2J")
        print(formatted_status(guesses, last=guess))
        if server.get_rank(guess) == 1000:
            time.sleep(1)
            print(Fore.GREEN+"Correct!"+Style.RESET_ALL+f" {len(guesses)} tries.")
            return len(guesses)




def main():
    parser = argparse.ArgumentParser(description="local infinite cemantix")
    parser.add_argument("-w", "--word", dest="word", action="store",
                        help="Specify goal word")
    parser.add_argument("-m", "--model", dest="model", action="store",
                        default="models/selected_word2vec_model.bin",
                        help="Specify model to use")
    args = parser.parse_args()

    cemantix(LocalServer(word=args.word, file=args.model))

if __name__ == "__main__":
    main()