#!/bin/python
import zipfile
import os
import urllib.parse as urlparse
import requests
import logging

from rasa_nlu.train import do_train
from rasa_nlu import utils, config
from rasa_nlu.utils import zip_folder, EndpointConfig

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.DEBUG)


# push_model
def push_model_to_platform(platform_host, api_token, model_path):

    model_name = os.path.basename(model_path)
    filename = zip_folder(model_path)

    files = {'model': ('{}.zip'.format(model_name), open(filename, 'rb'), 'application/zip')}
    url = "{}/api/projects/default/models/nlu?api_token={}"
    url = url.format(platform_host, api_token)
    requests.post(url, files=files)

    return model_name


# eval_models
def eval_new_and_active_models(platform_host, api_token, model_name):

    # get the active model
    url = "{}/api/projects/default/models/nlu?api_token={}"
    url = url.format(platform_host, api_token)
    models = requests.get(url).json()

    active_model = None
    for model in models:
        if "active" in model["tags"]:
            active_model = model["model"]
            break
    if active_model is None:
        raise ValueError("no active model found")

    logger.debug("active model is: {}".format(active_model))

    def get_eval(model_name):
        url = "{}/api/projects/default/evaluations/{}?api_token={}"
        url = url.format(platform_host, model_name, api_token)
        _ = requests.put(url)
        response = requests.get(url)
        print(response)
        return response.json()["intent_evaluation"]

    eval_new = get_eval(model_name)
    eval_active = get_eval(active_model)

    return eval_new, eval_active


# activate_model
def set_model_as_active(platform_host, api_token, model_name):
    url = "{}/api/projects/default/models/nlu/{}/tags/active?api_token={}"
    url = url.format(platform_host, model_name, api_token)
    response = requests.put(url)
    return response.status_code == 204

# main
if __name__ == "__main__":

    ### initialize constants
    platform_host = os.environ.get("RASA_PLATFORM_HOST", "https://rasa.example.com")
    api_token = os.environ.get("RASA_API_TOKEN", "")
    path = "./projects"
    project = "default"
    config_path="nlu_config.yml"

    data_endpoint = EndpointConfig(
        "{}/api/projects/default/data.json".format(platform_host),
        token=api_token,
        token_name="api_token")

    ### train a model with the latest training data from Rasa Platform
    _, _, model_path = do_train(config.load(config_path), 
                                None,
                                path=path,
                                project=project,
                                data_endpoint=data_endpoint)
    logger.info("training finished, model path {}".format(model_path))

    ### upload newly trained model to Rasa Platform
    model_name = push_model_to_platform(platform_host, api_token, model_path)

    ### mark new model as active if performance is good
    eval_new, eval_active = eval_new_and_active_models(platform_host, api_token, model_name)
    f1_new, f1_active = eval_new["f1_score"], eval_active["f1_score"]

    logger.info("intent evaluation f1 scores:")
    logger.info("new: {}".format(f1_new))
    logger.info("active: {}".format(f1_active))

    if f1_new > f1_active:
        logger.info("Newly trained model is better! going to mark it as active")
        success = set_model_as_active(platform_host, api_token, model_name)
        if success:
            logger.info("successfully activated model {}".format(model_name))
        else:
            logger.info("failed to activate model {}".format(model_name))
    else:
        logger.info("model did not improve")