import logging
import os
import threading
import traceback
from pathlib import Path
from typing import List

from flask import Flask, render_template
from flask_bootstrap import Bootstrap
from flask_wtf import FlaskForm
from torch.utils.tensorboard import SummaryWriter

from flair.models import SequenceTagger
from flair.trainers import ModelTrainer
from flair.datasets import ColumnCorpus
from flair.embeddings import (
    WordEmbeddings,
    StackedEmbeddings,
    FlairEmbeddings,
    TransformerWordEmbeddings,
    TokenEmbeddings,
    CharacterEmbeddings,
    BytePairEmbeddings
)

import grpc
from wtforms import TextAreaField, SubmitField
from wtforms.validators import DataRequired

import model_pb2
import model_pb2_grpc
from concurrent import futures
import predict
import sys


def my_except_hook(exctype, value, traceback):
    if exctype == KeyboardInterrupt:
        print("Handler code goes here")
    else:
        sys.__excepthook__(exctype, value, traceback)


sys.excepthook = my_except_hook

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

port = 8061
results = []

shared_folder = os.getenv("SHARED_FOLDER_PATH")
print(f'shared_folder: {shared_folder}')
writer = SummaryWriter(log_dir=shared_folder)
app = Flask(__name__)


class NerInputForm(FlaskForm):
    Text = TextAreaField('Text to be processed', validators=[DataRequired()])

    predict = SubmitField('Submit Form')


@app.route('/', methods=['GET', 'POST'])
def ner_input():
    form = NerInputForm()
    print("classifier_input()")

    if form.predict.data and form.validate_on_submit():
        logger.debug("Processing user inputs")
        request = model_pb2.Text()
        request.text = form.Text.data
        result = flair_model_servicer.extract_entities_from_text(request, None)
    else:
        result = model_pb2.Text()
        result.text = ""
    return render_template("index.html", example_form=form, result=result)


def app_run():
    app.secret_key = "ner"
    bootstrap = Bootstrap(app)
    app.run(host="0.0.0.0", port=8062)
    # app.run()


class FlairModelServicer(model_pb2_grpc.FlairModelServicer):
    def startTraining(self, request, context):
        print(f"start training with request:\n{request}")

        training_data_filename = request.training_data_filename
        validation_data_filename = request.validation_data_filename
        test_data_filename = request.test_data_filename
        embeddings = request.embeddings
        epochs = request.epochs
        learning_rate = request.learning_rate
        batch_size = request.batch_size
        model_filename = request.model_filename

        columns = {0: "text", 1: "ner"}

        corpus: ColumnCorpus = ColumnCorpus(
            Path(shared_folder),
            columns,
            train_file=training_data_filename,
            dev_file=validation_data_filename,
            test_file=test_data_filename,
        )

        tag_type = 'ner'

        tag_dictionary = corpus.make_tag_dictionary(tag_type=tag_type)

        embedding_types: List[TokenEmbeddings] = []

        try:
            for element in embeddings:
                if element.embedding == "wikipedia":
                    print("Using Wikipedia FastText embeddings")
                    embedding_types.append(WordEmbeddings("de"))
                elif element.embedding == "character":
                    print("Using character embeddings")
                    embedding_types.append(CharacterEmbeddings())
                elif element.embedding == "crawl":
                    logger.info("Using CommonCrawl FastText embeddings")
                    embedding_types.append(WordEmbeddings("de-crawl"))
                elif element.embedding == "bpe":
                    logger.info("Using German BPEmbeddings with 100k vocab and 300 dims")
                    embedding_types.append(BytePairEmbeddings(language="de", dim=300))
                elif element.embedding == "flair":
                    logger.info("Using Flair embeddings")
                    embedding_types.append(FlairEmbeddings('de-forward'))
                    embedding_types.append(FlairEmbeddings('de-backward'))
                elif element.embedding.startswith("bert"):
                    bert_model = 'bert-base-multilingual-cased'
                    logger.info("Using BERT multi-language model with the 4 last layers:")
                    embedding_types.append(
                        TransformerWordEmbeddings(
                            model=bert_model, layers='-1,-2,-3,-4'
                        )
                    )
                else:
                    logger.error(
                        f"Embedding name {element.embedding} not recognized! Please check your input!"
                    )
                    exit(1)

            print("Downloaded all embeddings...")

            stacked_embeddings: StackedEmbeddings = StackedEmbeddings(embeddings=embedding_types)
        except Exception as e:
            logging.error(traceback.format_exc())

        tagger: SequenceTagger = SequenceTagger(
            hidden_size=512,
            embeddings=stacked_embeddings,
            tag_dictionary=tag_dictionary,
            tag_type=tag_type,
            use_crf=True,
        )

        trainer: ModelTrainer = ModelTrainer(model=tagger, corpus=corpus, use_tensorboard=True)

        result = trainer.train(
            f"{shared_folder}/trained_model/{model_filename}",
            mini_batch_size=batch_size,
            learning_rate=learning_rate,
            max_epochs=epochs,
        )

        print(result)

        for i in range(len(result['train_loss_history'])):
            writer.add_scalar("Loss/train", result['train_loss_history'][i], i + 1)
            writer.add_scalar("Loss/dev", result['dev_loss_history'][i], i + 1)
            writer.add_scalar("Score/dev", result['dev_score_history'][i], i + 1)

        writer.flush()

        response = model_pb2.TrainingStatus()
        response.status_text = 'SUCCEEDED'

        writer.close()
        return response

    def extract_entities_from_text(self, request, context):
        response = model_pb2.Text()
        response.text = predict.predict_entity_from_text(request.text) + "\n"
        return response


server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
flair_model_servicer = FlairModelServicer()
model_pb2_grpc.add_FlairModelServicer_to_server(
    flair_model_servicer, server)
server.add_insecure_port('[::]:{}'.format(port))
print("Starting server. Listening on port : " + str(port))
logger.debug("[ner_model_server.py] LOG: Server about to start")
server.start()
logger.debug("[ner_model_server.py] LOG: Server successfully started")
threading.Thread(target=app_run()).start()
server.wait_for_termination()
