# -*- coding: utf-8 -*-

import json
import uuid
from datetime import datetime
import apache_beam as beam
from beamism.utils import *


class PubSubLambda(beam.DoFn):
    def __init__(self,
                 func,
                 input_schema,
                 pipeline_name='',
                 mongo_uri=None,
                 mongo_db=None,
                 mongo_coll=None,
                 func_version=None):
        self.func = func
        self.input_schema = input_schema
        self.pipeline_name = pipeline_name
        self.mongo_uri = mongo_uri
        self.mongo_db = mongo_db
        self.mongo_coll = mongo_coll
        self.func_version = func_version

    def save(self, output_json):
        if self.mongo_db is not None and self.mongo_coll is not None:
            is_acknowledged = save_to_mongo(self.mongo_uri, self.mongo_db, self.mongo_coll, output_json)
            return is_acknowledged
        return False

    def process(self, msg_text):
        element_id = str(uuid.uuid4())
        start_time = datetime.now()

        input_json = parse_json_message(msg_text)
        if input_json is None:
            output_json = {
                "params": { "input_text": msg_text, "error": "parse error" },
                "element_id": element_id,
                "data": None,
                "succeeded": False,
                "message": "Cannot parse this message as json",
                "start_time": str(start_time),
                "end_time": str(datetime.now()),
                "processing_time": (datetime.now() - start_time).total_seconds(),
                "func_version": self.func_version
            }
            self.save(output_json)
            return pubsub_packing(output_json)

        if self.input_schema is not None:
            input_validation_err_message = validate_json_based_on_schema(input_json, self.input_schema)

            if input_validation_err_message is not None:
                output_json = {
                    "params": { "input_text": msg_text, "error": "validation error" },
                    "element_id": element_id,
                    "data": None,
                    "succeeded": False,
                    "message": input_validation_err_message,
                    "start_time": str(start_time),
                    "end_time": str(datetime.now()),
                    "processing_time": (datetime.now() - start_time).total_seconds(),
                    "func_version": self.func_version
                }
                self.save(output_json)
                return pubsub_packing(output_json)

        if not input_json["succeeded"] or input_json["data"] is None:
            output_jsons = {
                "params": input_json["data"],
                "sender_element_id": input_json["element_id"] if "element_id" in input_json else None,
                "element_id": element_id,
                "data": None,
                "succeeded": False,
                "message": input_json["message"],
                "start_time": str(start_time),
                "end_time": str(datetime.now()),
                "processing_time": (datetime.now() - start_time).total_seconds(),
                "func_version": self.func_version
            }
            self.save(output_json)
            return pubsub_packing(output_json)

        indata = input_json["data"]
        results, is_func_success, func_msg = self.func(indata)

        if type(results) != list:
            results = [{}] if results is None else [results]

        output_json_for_mongo = { 
            "params": input_json["data"],
            "sender_element_id": input_json["element_id"] if "element_id" in input_json else None,
            "element_id": element_id,
            "data": results,
            "succeeded": is_func_success,
            "message": func_msg,
            "start_time": str(start_time),
            "end_time": str(datetime.now()),
            "processing_time": (datetime.now() - start_time).total_seconds(),
            "func_version": self.func_version 
        }
        self.save(output_json_for_mongo)

        outdatas = [merge_jsons(indata, result) for result in results]
        output_jsons = [{ "params": input_json["data"],
                          "sender_element_id": input_json["element_id"] if "element_id" in input_json else None,
                          "element_id": element_id,
                          "data": outdata,
                          "succeeded": is_func_success,
                          "message": func_msg,
                          "start_time": str(start_time),
                          "end_time": str(datetime.now()),
                          "processing_time": (datetime.now() - start_time).total_seconds(),
                          "func_version": self.func_version } for outdata in outdatas]
        return pubsub_packing(output_jsons)
