import os
import datetime

import pandas._libs.json as ujson
from bson import ObjectId
from pymongo import MongoClient

from cirrocumulus.abstract_db import AbstractDB
from cirrocumulus.util import get_email_domain, get_fs

from .envir import (
    CIRRO_AUTH_CLIENT_ID,
    CIRRO_DB_URI,
    CIRRO_JOB_RESULTS,
    CIRRO_TEST,
    SERVER_CAPABILITY_ADD_DATASET,
    SERVER_CAPABILITY_DELETE_DATASET,
    SERVER_CAPABILITY_EDIT_DATASET,
    SERVER_CAPABILITY_FEATURE_SETS,
    SERVER_CAPABILITY_JOBS,
    SERVER_CAPABILITY_LINKS,
    SERVER_CAPABILITY_RENAME_CATEGORIES,
)
from .invalid_usage import InvalidUsage
from .job_api import save_job_result_to_file


def format_doc(d):
    d["id"] = str(d.pop("_id"))
    return d


class MongoDb(AbstractDB):
    def __init__(self):
        super().__init__()
        if bool(os.environ.get(CIRRO_TEST)):
            import mongomock

            self.client = mongomock.MongoClient(os.environ[CIRRO_DB_URI])
        else:
            self.client = MongoClient(os.environ[CIRRO_DB_URI])
        self.db = self.client.get_default_database()
        self.db.categories.create_index("dataset_id")
        self.db.views.create_index("dataset_id")
        self.db.feature_sets.create_index("dataset_id")
        self.db.jobs.create_index("dataset_id")
        self.db.datasets.create_index("readers")
        self.fs = None

    def category_names(self, email, dataset_id):

        self.get_dataset(email, dataset_id)
        collection = self.db.categories
        results = {}
        for doc in collection.find(dict(dataset_id=dataset_id)):
            category = doc["category"]
            cat_results = results.get(category)
            if cat_results is None:
                cat_results = {}
                results[category] = cat_results
            cat_results[doc["original"]] = {
                "newValue": doc.get("new"),
                "negativeMarkers": doc.get("negativeMarkers"),
                "positiveMarkers": doc.get("positiveMarkers"),
                "color": doc.get("color"),
            }
        return results

    def upsert_category_name(self, email, dataset_id, category, original_value, update):
        if not self.capabilities()[SERVER_CAPABILITY_RENAME_CATEGORIES]:
            return
        self.get_dataset(email, dataset_id)
        collection = self.db.categories
        key = str(dataset_id) + "-" + str(category) + "-" + str(original_value)
        update["category"] = category
        update["dataset_id"] = dataset_id
        update["original"] = original_value
        update["last_updated"] = datetime.datetime.utcnow()
        if "newValue" in update:  # backwards compatibility
            update["new"] = update.pop("newValue")
        collection.update_one(dict(cat_id=key), {"$set": update}, upsert=True)

    def user(self, email):
        collection = self.db.users
        doc = collection.find_one(dict(email=email))
        if doc is None:
            collection.insert_one(dict(email=email))
            return {"id": email}
        else:
            return {"id": email, "importer": doc.get("importer", False)}

    def get_dataset(self, email, dataset_id, ensure_owner=False):
        collection = self.db.datasets
        auth_client_id = os.environ.get(CIRRO_AUTH_CLIENT_ID)

        if auth_client_id is None:  # allow unregistered URL
            try:
                dataset_id.index("://")
                return {"id": dataset_id, "name": dataset_id, "url": dataset_id}
            except ValueError:
                pass
        doc = collection.find_one(dict(_id=ObjectId(dataset_id)))
        if doc is None:
            raise InvalidUsage("Please provide a valid id", 400)
        readers = doc.get("readers")
        domain = get_email_domain(email)
        if email not in readers and domain not in readers and "" not in readers:
            raise InvalidUsage("Not authorized", 403)
        if ensure_owner and email not in doc["owners"]:
            raise InvalidUsage("Not authorized", 403)
        doc["owner"] = "owners" in doc and email in doc.pop("owners")
        doc["id"] = str(doc.pop("_id"))
        return doc

    def datasets(self, email):
        collection = self.db.datasets
        results = []
        domain = get_email_domain(email)
        if domain is None:
            query = dict(readers={"$in": [email, ""]})
        else:
            query = dict(readers={"$in": [email, domain]})

        for doc in collection.find(query):
            doc["owner"] = "owners" in doc and email in doc.pop("owners")
            doc["id"] = str(doc.pop("_id"))
            results.append(doc)
        return results

    # views
    def dataset_views(self, email, dataset_id):
        self.get_dataset(email, dataset_id)
        collection = self.db.views
        results = []
        for doc in collection.find(dict(dataset_id=dataset_id), dict(value=0)):
            results.append(format_doc(doc))
        return results

    def delete_dataset_view(self, email, view_id):
        if not self.capabilities()[SERVER_CAPABILITY_LINKS]:
            return
        collection = self.db.views
        doc = collection.find_one(dict(_id=ObjectId(view_id)))
        self.get_dataset(email, doc["dataset_id"])
        collection.delete_one(dict(_id=ObjectId(view_id)))

    def get_dataset_view(self, email, view_id):
        collection = self.db.views
        doc = collection.find_one(dict(_id=ObjectId(view_id)))
        self.get_dataset(email, doc["dataset_id"])
        return format_doc(doc)

    def upsert_dataset_view(self, email, dataset_id, view):
        if not self.capabilities()[SERVER_CAPABILITY_LINKS]:
            return
        self.get_dataset(email, dataset_id)
        collection = self.db.views
        view["last_updated"] = datetime.datetime.utcnow()

        if email is not None:
            view["email"] = email
        if dataset_id is not None:
            view["dataset_id"] = dataset_id
        if view.get("id") is None:
            return dict(
                id=str(collection.insert_one(view).inserted_id), last_updated=view["last_updated"]
            )
        else:
            view_id = view.pop("id")
            collection.update_one(dict(_id=ObjectId(view_id)), {"$set": view})
            return dict(id=view_id, last_updated=view["last_updated"])

    def is_importer(self, email):
        # TODO check if user can modify dataset
        user = self.db.users.find_one(dict(email=email))
        if "importer" not in user:
            raise False
        return user["importer"]

    def delete_dataset(self, email, dataset_id):
        if not self.capabilities()[SERVER_CAPABILITY_DELETE_DATASET]:
            return

        self.get_dataset(email, dataset_id, True)
        self.db.datasets.delete_one(dict(_id=ObjectId(dataset_id)))
        self.db.filters.delete_many(dict(dataset_id=dataset_id))
        self.db.views.delete_many(dict(dataset_id=dataset_id))
        self.db.categories.delete_many(dict(dataset_id=dataset_id))
        self.db.feature_sets.delete_many(dict(dataset_id=dataset_id))

        for doc in self.db.jobs.find(dict(dataset_id=dataset_id), dict(result=1)):
            result = doc.get("result")
            self.delete_job_result(result)
        self.db.jobs.delete_many(dict(dataset_id=dataset_id))

    def upsert_dataset(self, email, readers, dataset):

        if dataset.get("id") is None and not self.capabilities()[SERVER_CAPABILITY_ADD_DATASET]:
            return
        if (
            dataset.get("id") is not None
            and not self.capabilities()[SERVER_CAPABILITY_EDIT_DATASET]
        ):
            return
        collection = self.db.datasets
        if readers is not None:
            readers = set(readers)
            if email in readers:
                readers.remove(email)
            readers.add(email)
            dataset["readers"] = list(readers)
        dataset["last_updated"] = datetime.datetime.utcnow()
        if dataset.get("id") is None:  # new dataset
            if email != "":
                user = self.db.users.find_one(dict(email=email))
                if "importer" not in user or not user["importer"]:
                    raise InvalidUsage("Not authorized", 403)
            dataset["owners"] = [email]
            if "readers" not in dataset:
                dataset["readers"] = [email]
            return str(collection.insert_one(dataset).inserted_id)
        else:  # update
            self.get_dataset(email, dataset["id"], True)
            dataset_id = dataset.pop("id")
            collection.update_one(dict(_id=ObjectId(dataset_id)), {"$set": dataset})
            return dataset_id

    def get_feature_sets(self, email, dataset_id):
        self.get_dataset(email, dataset_id)
        collection = self.db.feature_sets
        results = []
        for doc in collection.find(dict(dataset_id=dataset_id)):
            results.append(format_doc(doc))
        return results

    def delete_feature_set(self, email, dataset_id, set_id):
        if not self.capabilities()[SERVER_CAPABILITY_FEATURE_SETS]:
            return
        collection = self.db.feature_sets
        doc = collection.find_one(dict(_id=ObjectId(set_id)))
        self.get_dataset(email, doc["dataset_id"])
        collection.delete_one(dict(_id=ObjectId(set_id)))

    def upsert_feature_set(self, email, dataset_id, set_id, category, name, features):
        if not self.capabilities()[SERVER_CAPABILITY_FEATURE_SETS]:
            return
        self.get_dataset(email, dataset_id)
        collection = self.db.feature_sets
        entity_update = {}
        entity_update["last_updated"] = datetime.datetime.utcnow()
        if name is not None:
            entity_update["name"] = name
        if features is not None:
            entity_update["features"] = features
        if email is not None:
            entity_update["email"] = email
        if dataset_id is not None:
            entity_update["dataset_id"] = dataset_id
        if category is not None:
            entity_update["category"] = category
        if set_id is None:
            return str(collection.insert_one(entity_update).inserted_id)
        else:
            collection.update_one(dict(_id=ObjectId(set_id)), {"$set": entity_update})
            return set_id

    def get_gridfs(self):
        import gridfs

        if self.fs is None:
            self.fs = gridfs.GridFS(self.db)
        return self.fs

    def create_job(self, email, dataset_id, job_name, job_type, params):
        if not self.capabilities()[SERVER_CAPABILITY_JOBS]:
            return
        self.get_dataset(email, dataset_id)
        collection = self.db.jobs
        job_id = str(
            collection.insert_one(
                dict(
                    dataset_id=dataset_id,
                    status="pending",
                    name=job_name,
                    email=email,
                    type=job_type,
                    params=params,
                    submitted=datetime.datetime.utcnow(),
                )
            ).inserted_id
        )
        return job_id

    def get_job(self, email, job_id, return_type):
        collection = self.db.jobs
        doc = collection.find_one(dict(_id=ObjectId(job_id)))
        if doc is None:
            return None
        self.get_dataset(email, doc["dataset_id"])
        if return_type == "result":
            result = doc.get("result")
            if isinstance(result, dict) and result.get("url") is not None:
                return result
            else:
                fs = self.get_gridfs()
                return fs.get(ObjectId(doc["result"])).read()
        elif return_type == "status":
            return dict(status=doc["status"])
        elif return_type == "params":
            return dict(params=doc.get("params"))

    def get_jobs(self, email, dataset_id):
        self.get_dataset(email, dataset_id)
        collection = self.db.jobs
        results = []
        for doc in collection.find(
            dict(dataset_id=dataset_id),
            dict(name=1, status=1, email=1, type=1, submitted=1, readonly=1),
        ):
            results.append(
                dict(
                    id=str(doc["_id"]),
                    name=doc["name"],
                    status=doc["status"],
                    type=doc["type"],
                    email=doc["email"],
                    submitted=doc.get("submitted"),
                    readonly=doc.get("readonly", False),
                )
            )
        return results

    def annotate_job(self, email, job_id, annotations):
        collection = self.db.jobs
        doc = collection.find_one(dict(_id=ObjectId(job_id)))
        self.get_dataset(email, doc["dataset_id"])
        collection.update_one(
            dict(_id=ObjectId(job_id)),
            {"$set": dict(annotations=annotations, last_updated=datetime.datetime.utcnow())},
        )

    def update_job(self, email, job_id, status, result):
        if not self.capabilities()[SERVER_CAPABILITY_JOBS]:
            return
        collection = self.db.jobs
        doc = collection.find_one(dict(_id=ObjectId(job_id)))
        self.get_dataset(email, doc["dataset_id"])
        if doc.get("readonly", False):
            raise InvalidUsage("Not authorized", 403)
        if result is not None:
            if os.environ.get(CIRRO_JOB_RESULTS) is not None:  # save to directory
                result = save_job_result_to_file(result, job_id)
            else:
                result = ujson.dumps(result, double_precision=2, orient="values")
                result = str(self.get_gridfs().put(result, encoding="ascii"))

        collection.update_one(
            dict(_id=ObjectId(job_id)), {"$set": dict(status=status, result=result)}
        )

    def delete_job_result(self, result):
        if result is not None:
            if isinstance(result, dict) and result.get("url") is not None:
                fs = get_fs(result["url"])
                if fs.exists(result["url"]):
                    fs.rm(result["url"], recursive=True)
            else:
                self.get_gridfs().delete(ObjectId(result))

    def delete_job(self, email, job_id):
        if not self.capabilities()[SERVER_CAPABILITY_JOBS]:
            return
        collection = self.db.jobs
        doc = collection.find_one(dict(_id=ObjectId(job_id)), dict(email=1))
        if doc["email"] == email and not doc.get("readonly", False):
            collection.delete_one(dict(_id=ObjectId(job_id)))
            result = doc.get("result")
            self.delete_job_result(result)
        else:
            raise InvalidUsage("Not authorized", 403)
