#!/usr/bin/env python
# -*- mode: python; encoding: utf-8 -*-

"""An implementation of a data store based on mongo."""


import hashlib
import threading
import time
from bson import binary
from bson import objectid
import pymongo
from pymongo import errors
import logging

from grr.lib import config_lib
from grr.lib import data_store
from grr.lib import utils


class MongoDataStore(data_store.DataStore):
  """A Mongo based data store."""

  def __init__(self):
    # Support various versions on the pymongo connection object.
    try:
      connector = pymongo.MongoClient
    except AttributeError:
      connector = pymongo.Connection

    if config_lib.CONFIG["Mongo.server"]:
      mongo_client = connector(
          host=config_lib.CONFIG["Mongo.server"],
          port=int(config_lib.CONFIG["Mongo.port"]))

    else:
      mongo_client = connector()

    # For now use a single "data" collection
    self.db_handle = mongo_client[config_lib.CONFIG["Mongo.db_name"]]

    # We have two collections - the latest collection maintains the latest data
    # and the versioned collection maintains versioned data.
    self.latest_collection = self.db_handle.latest
    self.versioned_collection = self.db_handle.versioned

    # Ensure we have the correct indexes.
    for idx in ["subject", "predicate", "timestamp"]:
      self.latest_collection.ensure_index(idx)
      self.versioned_collection.ensure_index(idx)

    super(MongoDataStore, self).__init__()

  def _GetCursor(self, spec, timestamp, limit):
    """Create a mongo cursor based on the timestamp restriction."""

    if timestamp == self.NEWEST_TIMESTAMP or timestamp is None:
      collection = self.latest_collection
    elif timestamp == self.ALL_TIMESTAMPS:
      collection = self.versioned_collection
    elif isinstance(timestamp, tuple):
      collection = self.versioned_collection
      start, end = timestamp
      spec = {"$and": [dict(timestamp={"$gte": int(start)}),
                       dict(timestamp={"$lte": int(end)}),
                       spec]}
    else:
      raise data_store.Error("Undefined timestamp specification.")

    cursor = collection.find(spec).sort("timestamp", pymongo.DESCENDING)

    if limit:
      cursor = cursor.limit(limit)

    return cursor

  def ResolveMulti(self, subject, attributes, timestamp=None, limit=None,
                   token=None):
    """Resolves multiple attributes at once for one subject."""
    self.security_manager.CheckDataStoreAccess(
        token, [subject], self.GetRequiredResolveAccess(attributes))

    # Build a query spec.
    spec = {"$and": [
        # Subject matches any of the requested subjects.
        dict(subject=utils.SmartUnicode(subject)),
        {"$or": [dict(predicate=utils.SmartUnicode(x)) for x in attributes]},
    ]}

    results_returned = 0
    for document in self._GetCursor(spec, timestamp, 0):
      subject = document["subject"]
      value = Decode(document)
      if limit:
        if results_returned >= limit:
          return
        results_returned += 1

      yield (document["predicate"], value, document["timestamp"])

  def DeleteSubject(self, subject, sync=False, token=None):
    """Completely deletes all information about the subject."""
    _ = sync
    subject = utils.SmartUnicode(subject)
    self.security_manager.CheckDataStoreAccess(token, [subject], "w")
    self.latest_collection.remove(dict(subject=subject))
    self.versioned_collection.remove(dict(subject=subject))

  def MultiSet(self, subject, values, timestamp=None, replace=True,
               sync=True, to_delete=None, token=None):
    """Set multiple attributes' values for this subject in one operation."""
    self.security_manager.CheckDataStoreAccess(token, [subject], "w")

    if timestamp is None:
      timestamp = time.time() * 1e6

    # Prepare a mongo bulk insert for all the values.
    documents = []
    subject = utils.SmartUnicode(subject)
    to_delete = set(to_delete or [])

    latest = {}

    # Build a document for each unique timestamp.
    for attribute, sequence in values.items():
      for value in sequence:
        if isinstance(value, tuple):
          value, entry_timestamp = value
        else:
          entry_timestamp = timestamp

        if entry_timestamp is None:
          entry_timestamp = timestamp

        attribute = utils.SmartUnicode(attribute)
        prefix = attribute.split(":", 1)[0]