import argparse
import json
import logging
import os
import sys
import uuid

import redis

from .tasks.train_model import train_models
from .utils.aws_connectors import get_secret

logging.getLogger().setLevel(logging.INFO)


def run_training(params):
    schema = params.schema
    collection_ids = [int(collection_id) for collection_id in params.collection_ids.split(",")]
    bucket = params.bucket
    file_key = params.file_key

    """ We used this in legacy to differentiate between "RFM modelling" and "other". When feeding input fields
    from whatever is calling model training, we no longer have a need for this.
    """
    modelling_config = {
        "plot_2d": False,
        "plot_3d": False,
        "plot_silhouette": False,
        "importance": True,
        "boxplot": False,
        "save_model": False,
        "modelling_type": "clustering",
        # -1 means no pca dimensionality reduction is done. Multiple values mean for each of those values a model will be trained
        "pca_components": [4],
        # TODO! All experiments with -1 (no pca) were inferior to PCA, but there is more room to experiment
        "drop_corr_threshold": 0.99
        # Drop features that are highly correlated. That's why - https://www.trchome.com/docs/5-cluster-analysis-gets-complicated/file
    }

    return train_models(schema, collection_ids, bucket, file_key, modelling_config)


def publish_result(job_id, training_result):
    redis_config = json.loads(get_secret("fansifter-redis"))

    client = redis.Redis(**redis_config, db=0)

    client.publish(job_id, json.dumps(training_result))  # TODO Remove after Airflow enrichments are deployed to LIVE

    redis_key = f"batch-ml-result:{job_id}"
    client.set(redis_key, json.dumps(training_result))
    client.expire(redis_key, 6 * 60 * 60)  # 6 hours
    logging.info(f"Published training result to Redis (key: {redis_key}) : {json.dumps(training_result)}")


if __name__ == "__main__":
    os.environ["MLFLOW_S3_UPLOAD_EXTRA_ARGS"] = "{\"ServerSideEncryption\": \"aws:kms\"}"

    parser = argparse.ArgumentParser(description="FanSifter ML Engine CLI")
    parser.add_argument("--schema", dest="schema", required=True)
    parser.add_argument("--collection-ids", dest="collection_ids", required=True)
    parser.add_argument("--bucket", dest="bucket", required=True)
    parser.add_argument("--file-key", dest="file_key", required=True)

    args = parser.parse_args()
    logging.info(f"input data: {args}")

    aws_batch_job_id = os.environ.get("AWS_BATCH_JOB_ID") or str(uuid.uuid4())  # uuid fallback for local testing

    try:
        training_result = run_training(args)
    except Exception as e:
        logging.error("Failed to train models!", exc_info=e)
        publish_result(aws_batch_job_id, dict(error="Failed to train models!"))
        sys.exit(1)

    try:
        publish_result(aws_batch_job_id, training_result)
    except Exception as e:
        logging.error("Failed to publish training result to Redis!", exc_info=e)
        sys.exit(1)

    logging.info(f"Completed job {aws_batch_job_id} with result: {training_result}")