from rpython.rtyper.lltypesystem import rffi, lltype
from pypy.module.cpyext.api import (
    cpython_api, CANNOT_FAIL, build_type_checkers, Py_ssize_t,
    Py_ssize_tP, CONST_STRING)
from pypy.module.cpyext.pyobject import PyObject, PyObjectP, as_pyobj
from pypy.module.cpyext.pyerrors import PyErr_BadInternalCall
from pypy.interpreter.error import OperationError
from rpython.rlib.objectmodel import specialize

@cpython_api([], PyObject)
def PyDict_New(space):
    return space.newdict()

PyDict_Check, PyDict_CheckExact = build_type_checkers("Dict")

@cpython_api([PyObject, PyObject], PyObject, error=CANNOT_FAIL,
             result_borrowed=True)
def PyDict_GetItem(space, w_dict, w_key):
    try:
        w_res = space.getitem(w_dict, w_key)
    except:
        return None
    # NOTE: this works so far because all our dict strategies store
    # *values* as full objects, which stay alive as long as the dict is
    # alive and not modified.  So we can return a borrowed ref.
    return w_res

@cpython_api([PyObject, PyObject, PyObject], rffi.INT_real, error=-1)
def PyDict_SetItem(space, w_dict, w_key, w_obj):
    if PyDict_Check(space, w_dict):
        space.setitem(w_dict, w_key, w_obj)
        return 0
    else:
        PyErr_BadInternalCall(space)

@cpython_api([PyObject, PyObject], rffi.INT_real, error=-1)
def PyDict_DelItem(space, w_dict, w_key):
    if PyDict_Check(space, w_dict):
        space.delitem(w_dict, w_key)
        return 0
    else:
        PyErr_BadInternalCall(space)

@cpython_api([PyObject, CONST_STRING, PyObject], rffi.INT_real, error=-1)
def PyDict_SetItemString(space, w_dict, key_ptr, w_obj):
    if PyDict_Check(space, w_dict):
        key = rffi.charp2str(key_ptr)
        space.setitem_str(w_dict, key, w_obj)
        return 0
    else:
        PyErr_BadInternalCall(space)

@cpython_api([PyObject, CONST_STRING], PyObject, error=CANNOT_FAIL,
             result_borrowed=True)
def PyDict_GetItemString(space, w_dict, key):
    """This is the same as PyDict_GetItem(), but key is specified as a
    char*, rather than a PyObject*."""
    try:
        w_res = space.finditem_str(w_dict, rffi.charp2str(key))
    except:
        w_res = None
    # NOTE: this works so far because all our dict strategies store
    # *values* as full objects, which stay alive as long as the dict is
    # alive and not modified.  So we can return a borrowed ref.
    return w_res

@cpython_api([PyObject, CONST_STRING], rffi.INT_real, error=-1)
def PyDict_DelItemString(space, w_dict, key_ptr):
    """Remove the entry in dictionary p which has a key specified by the string
    key.  Return 0 on success or -1 on failure."""
    if PyDict_Check(space, w_dict):
        key = rffi.charp2str(key_ptr)
        # our dicts dont have a standardized interface, so we need
        # to go through the space
        space.delitem(w_dict, space.wrap(key))
        return 0
    else:
        PyErr_BadInternalCall(space)

@cpython_api([PyObject], Py_ssize_t, error=-1)
def PyDict_Size(space, w_obj):
    """
    Return the number of items in the dictionary.  This is equivalent to
    len(p) on a dictionary."""
    return space.len_w(w_obj)

@cpython_api([PyObject, PyObject], rffi.INT_real, error=-1)
def PyDict_Contains(space, w_obj, w_value):
    """Determine if dictionary p contains key.  If an item in p is matches
    key, return 1, otherwise return 0.  On error, return -1.
    This is equivalent to the Python expression key in p.
    """
    w_res = space.contains(w_obj, w_value)
    return space.int_w(w_res)

@cpython_api([PyObject], lltype.Void)
def PyDict_Clear(space, w_obj):
    """Empty an existing dictionary of all key-value pairs."""
    space.call_method(space.w_dict, "clear", w_obj)

@cpython_api([PyObject], PyObject)
def PyDict_Copy(space, w_obj):
    """Return a new dictionary that contains the same key-value pairs as p.
    """
    return space.call_method(space.w_dict, "copy", w_obj)

@cpython_api([PyObject, PyObject], rffi.INT_real, error=-1)
def PyDict_Update(space, w_obj, w_other):
    """This is the same as PyDict_Merge(a, b, 1) in C, or a.update(b) in
    Python.  Return 0 on success or -1 if an exception was raised.
    """
    space.call_method(space.w_dict, "update", w_obj, w_other)
    return 0

@cpython_api([PyObject], PyObject)
def PyDict_Keys(space, w_obj):
    """Return a PyListObject containing all the keys from the dictionary,
    as in the dictionary method dict.keys()."""
    return space.call_method(space.w_dict, "keys", w_obj)

@cpython_api([PyObject], PyObject)
def PyDict_Values(space, w_obj):
    """Return a PyListObject containing all the values from the
    dictionary p, as in the dictionary method dict.values()."""
    return space.call_method(space.w_dict, "values", w_obj)

@cpython_api([PyObject], PyObject)
def PyDict_Items(space, w_obj):
    """Return a PyListObject containing all the items from the
    dictionary, as in the dictionary method dict.items()."""
    return space.call_method(space.w_dict, "items", w_obj)

@cpython_api([PyObject, Py_ssize_tP, PyObjectP, PyObjectP], rffi.INT_real, error=CANNOT_FAIL)
def PyDict_Next(space, w_dict, ppos, pkey, pvalue):
    """Iterate over all key-value pairs in the dictionary p.  The
    Py_ssize_t referred to by ppos must be initialized to 0
    prior to the first call to this function to start the iteration; the
    function returns true for each pair in the dictionary, and false once all
    pairs have been reported.  The parameters pkey and pvalue should either
    point to PyObject* variables that will be filled in with each key
    and value, respectively, or may be NULL.  Any references returned through
    them are borrowed.  ppos should not be altered during iteration. Its
    value represents offsets within the internal dictionary structure, and
    since the structure is sparse, the offsets are not consecutive.

    For example:

    PyObject *key, *value;
    Py_ssize_t pos = 0;

    while (PyDict_Next(self->dict, &pos, &key, &value)) {
        /* do something interesting with the values... */
        ...
    }

    The dictionary p should not be mutated during iteration.  It is safe
    (since Python 2.1) to modify the values of the keys as you iterate over the
    dictionary, but only so long as the set of keys does not change.  For
    example:

    PyObject *key, *value;
    Py_ssize_t pos = 0;

    while (PyDict_Next(self->dict, &pos, &key, &value)) {
        int i = PyInt_AS_LONG(value) + 1;
        PyObject *o = PyInt_FromLong(i);
        if (o == NULL)
            return -1;
        if (PyDict_SetItem(self->dict, key, o) < 0) {
            Py_DECREF(o);
            return -1;
        }
        Py_DECREF(o);
    }"""
    if w_dict is None:
        return 0

    # XXX XXX PyDict_Next is not efficient. Storing an iterator would probably
    # work, but we can't work out how to not leak it if iteration does
    # not complete.  Alternatively, we could add some RPython-only
    # dict-iterator method to move forward by N steps.

    w_dict.ensure_object_strategy()     # make sure both keys and values can
                                        # be borrwed
    try:
        w_iter = space.call_method(space.w_dict, "iteritems", w_dict)
        pos = ppos[0]
        while pos:
            space.call_method(w_iter, "next")
            pos -= 1

        w_item = space.call_method(w_iter, "next")
        w_key, w_value = space.fixedview(w_item, 2)
        if pkey:
            pkey[0]   = as_pyobj(space, w_key)
        if pvalue:
            pvalue[0] = as_pyobj(space, w_value)
        ppos[0] += 1
    except OperationError, e:
        if not e.match(space, space.w_StopIteration):
            raise
        return 0
    return 1

@specialize.memo()
def make_frozendict(space):
    return space.appexec([], '''():
    import _abcoll
    class FrozenDict(_abcoll.Mapping):
        def __init__(self, *args, **kwargs):
            self._d = dict(*args, **kwargs)
        def __iter__(self):
            return iter(self._d)
        def __len__(self):
            return len(self._d)
        def __getitem__(self, key):
            return self._d[key]
    return FrozenDict''')

@cpython_api([PyObject], PyObject)
def PyDictProxy_New(space, w_dict):
    w_frozendict = make_frozendict(space)
    return space.call_function(w_frozendict, w_dict)

