#include "vm/shadow.h"
#include "mm/page.h"
#include "mm/pframe.h"
#include "mm/slab.h"
#include "util/debug.h"
#include "util/string.h"

#define SHADOW_SINGLETON_THRESHOLD 5

typedef struct mobj_shadow
{
    // the mobj parts of this shadow object
    mobj_t mobj;
    // a reference to the mobj that is the data source for this shadow object
    // This should be a reference to a shadow object of some ancestor process.
    // This is used to traverse the shadow object chain.
    mobj_t *shadowed;
    // a reference to the mobj at the bottom of this shadow object's chain
    // this should NEVER be a shadow object (i.e. it should have some type other
    // than MOBJ_SHADOW)
    mobj_t *bottom_mobj;
} mobj_shadow_t;

#define MOBJ_TO_SO(o) CONTAINER_OF(o, mobj_shadow_t, mobj)

static slab_allocator_t *shadow_allocator;

static long shadow_get_pframe(mobj_t *o, size_t pagenum, long forwrite,
                              pframe_t **pfp);
static long shadow_fill_pframe(mobj_t *o, pframe_t *pf);
static long shadow_flush_pframe(mobj_t *o, pframe_t *pf);
static void shadow_destructor(mobj_t *o);

static mobj_ops_t shadow_mobj_ops = {.get_pframe = shadow_get_pframe,
                                     .fill_pframe = shadow_fill_pframe,
                                     .flush_pframe = shadow_flush_pframe,
                                     .destructor = shadow_destructor};

/*
 * Initialize shadow_allocator using the slab allocator.
 */
void shadow_init()
{
    // NOT_YET_IMPLEMENTED("VM: shadow_init");
    shadow_allocator = slab_allocator_create("shadow", sizeof(mobj_shadow_t));
    KASSERT(shadow_allocator);
}

/*
 * Create a shadow object that shadows the given mobj.
 *
 * Return a new, LOCKED shadow object on success, or NULL upon failure.
 *
 * Hints:
 *  1) Create and initialize a mobj_shadow_t based on the given mobj.
 *  2) Set up the bottom object of the shadow chain, which could have two cases:
 *     a) Either shadowed is a shadow object, and you can use its bottom_mobj
 *     b) Or shadowed is not a shadow object, in which case it is the bottom 
 *        object of this chain.shadow_create
 * 
 *  Make sure to manage the refcounts correctly.
 */
mobj_t *shadow_create(mobj_t *shadowed)
{
    // NOT_YET_IMPLEMENTED("VM: shadow_create");

    if (!shadowed)
    {
        return NULL;
    }

    // create a new shadow object
    mobj_shadow_t *so = (mobj_shadow_t *)slab_obj_alloc(shadow_allocator);
    if (!so)
    {
        return NULL;
    }

    // initialize the mobj_shadow_t
    so->shadowed = shadowed;
    
    // set the bottom_mobj based on the two cases
    if (shadowed->mo_type == MOBJ_SHADOW)
    {
        so->bottom_mobj = MOBJ_TO_SO(so->shadowed)->bottom_mobj;
    }
    else
    {
        so->bottom_mobj = shadowed;
    }
    // init the other fields
    mobj_ref(so->shadowed);
    mobj_lock(so->bottom_mobj);
    mobj_ref(so->bottom_mobj);
    mobj_unlock(so->bottom_mobj);

    // init and lock the shadow object
    mobj_init(&so->mobj, MOBJ_SHADOW, &shadow_mobj_ops);
    mobj_lock(&so->mobj);

    // return the shadow object
    return &so->mobj;
}

/*
 * Given a shadow object o, collapse its shadow chain as far as you can.
 *
 * Hints:
 *  1) You can only collapse if the shadowed object is a shadow object.
 *  2) When collapsing, you must manually migrate pframes from o's shadowed
 *     object to o, checking to see if a copy doesn't already exist in o.
 *  3) Be careful with refcounting! In particular, when you put away o's
 *     shadowed object, its refcount should drop to 0, initiating its
 *     destruction (shadow_destructor).
 *  4) As a reminder, any refcounting done in shadow_collapse() must play nice
 *     with any refcounting done in shadow_destructor().
 *  5) Pay attention to mobj and pframe locking.
 */
void shadow_collapse(mobj_t *o)
{
    // NOT_YET_IMPLEMENTED("VM: shadow_collapse");

    return;

    mobj_shadow_t *so = MOBJ_TO_SO(o);
    mobj_t *next = so->shadowed;
    while(next && next->mo_type==MOBJ_SHADOW)
    {

        // check if the shadowed object is collapsable
        if (next->mo_refcount != 1) 
        {
            // continue on the shadow chain
            so = MOBJ_TO_SO(next);
            next = MOBJ_TO_SO(next)->shadowed;
            continue;
        }

        // migrate the pframes from o's shadowed object to o
        mobj_lock(next); // lock before iterover
        list_iterate(&next->mo_pframes, pf, pframe_t, pf_link)
        {
            // see if shadow pframe is in the current shadow object
            pframe_t *pf_shadow = NULL;
            mobj_lock(&so->mobj);
            mobj_find_pframe(&so->mobj, pf->pf_pagenum, &pf_shadow);
            mobj_unlock(&so->mobj);
            // if so, migrate it
            if (!pf_shadow)
            {
                list_remove(&pf->pf_link);
                // by migrate, just add to the list of pframes
                list_insert_tail(&so->mobj.mo_pframes, &pf->pf_link);
            }
            else
            {
                // if already exists, just release the version we will remove
                pframe_release(&pf_shadow);
            }
        }

        // update the pointer to the shadowed object
        so->shadowed = MOBJ_TO_SO(next)->shadowed;

        // update the refcounts
        mobj_lock(so->shadowed);
        mobj_ref(so->shadowed);
        mobj_put_locked(&next); // "remove" the next

        // update next for next iter, now that it's removed
        next = so->shadowed;
        mobj_unlock(next);
    }
}


/*
 * Obtain the desired pframe from the given mobj, traversing its shadow chain if
 * necessary. This is where copy-on-write logic happens!
 *
 * Arguments: 
 *  o        - The object from which to obtain a pframe
 *  pagenum  - Number of the desired page relative to the object
 *  forwrite - Set if the caller wants to write to the pframe's data, clear if
 *             only reading
 *  pfp      - Upon success, pfp should point to the desired pframe.
 *
 * Return 0 on success, or:
 *  - Propagate errors from mobj_default_get_pframe() and mobj_get_pframe()
 *
 * Hints:
 *  1) If forwrite is set, use mobj_default_get_pframe().
 *  2) If forwrite is clear, check if o already contains the desired frame.
 *     a) If not, iterate through the shadow chain to find the nearest shadow
 *        mobj that has the frame. Do not recurse! If the shadow chain is long,
 *        you will cause a kernel buffer overflow (e.g. from forkbomb).
 *     b) If no shadow objects have the page, call mobj_get_pframe() to get the
 *        page from the bottom object and return what it returns.
 * 
 *  Pay attention to pframe locking.
 */
static long shadow_get_pframe(mobj_t *o, size_t pagenum, long forwrite,
                              pframe_t **pfp)
{
    // NOT_YET_IMPLEMENTED("VM: shadow_get_pframe");
    // return 0;
    if (forwrite)
    {
        return mobj_default_get_pframe(o, pagenum, forwrite, pfp);
    }
    // check if o already contains the desired frame.
    pframe_t *pf = NULL;
    mobj_find_pframe(o, pagenum, &pf);
    if (pf)
    {
        *pfp = pf;
        return 0;
    }

    mobj_shadow_t *shadow = MOBJ_TO_SO(o);
    mobj_t *current = shadow->shadowed; 

    while (current && current->mo_type == MOBJ_SHADOW)
    {
        mobj_lock(current);
        mobj_find_pframe(current, pagenum, &pf);
        mobj_unlock(current);
        
        if (pf)
        {
            *pfp = pf;
            return 0;
        }

        shadow = MOBJ_TO_SO(current);
        current = shadow->shadowed;
    }
    mobj_lock(shadow->bottom_mobj);
    long ret = mobj_get_pframe(shadow->bottom_mobj, pagenum, forwrite, pfp);
    mobj_unlock(shadow->bottom_mobj);
    return ret;
}

/*
 * Use the given mobj's shadow chain to fill the given pframe.
 *
 * Return 0 on success, or:
 *  - Propagate errors from mobj_get_pframe()
 *
 * Hints:
 *  1) Explore mobj_default_get_pframe(), which calls mobj_create_pframe(), to
 *     understand what state pf is in when this function is called, and how you
 *     can use it.
 *  2) As you can see above, shadow_get_pframe would call
 *     mobj_default_get_pframe (when the forwrite is set), which would 
 *     create and then fill the pframe (shadow_fill_pframe is called).
 *  3) Traverse the shadow chain for a copy of the frame, starting at the given
 *     mobj's shadowed object. You can use mobj_find_pframe to look for the 
 *     page frame. pay attention to locking/unlocking, and be sure not to 
 *     recurse when traversing.
 *  4) If none of the shadow objects have a copy of the frame, use
 *     mobj_get_pframe on the bottom object to get it.
 *  5) After obtaining the desired frame, simply copy its contents into pf.
 */
static long shadow_fill_pframe(mobj_t *o, pframe_t *pf)
{
    // NOT_YET_IMPLEMENTED("VM: shadow_fill_pframe");
    // return -1;
    pframe_t *pf_shadow = NULL;
    long ret = 0;


    mobj_shadow_t *so = MOBJ_TO_SO(o);
    mobj_t *shadowed = so->shadowed;
    while (shadowed && shadowed->mo_type == MOBJ_SHADOW)
    {
        mobj_lock(shadowed);
        mobj_find_pframe(shadowed, pf->pf_pagenum, &pf_shadow);
        mobj_unlock(shadowed);

        if (pf_shadow)
        {
            memcpy(pf->pf_addr, pf_shadow->pf_addr, PAGE_SIZE);
            pframe_release(&pf_shadow);
            return 0;
        }

        so = MOBJ_TO_SO(shadowed);
        shadowed = so->shadowed;
    }
    mobj_lock(so->bottom_mobj);
    ret = mobj_get_pframe(so->bottom_mobj, pf->pf_pagenum, 0, &pf_shadow);
    mobj_unlock(so->bottom_mobj);
    if (ret < 0) {
        return ret;
    }
    memcpy(pf->pf_addr, pf_shadow->pf_addr, PAGE_SIZE);
    pframe_release(&pf_shadow);
    return 0;
}

/*
 * Flush a shadow object's pframe to disk.
 *
 * Return 0 on success.
 *
 * Hint:
 *  - Are shadow objects backed to disk? Do you actually need to do anything
 *    here?
 */
static long shadow_flush_pframe(mobj_t *o, pframe_t *pf)
{
    // NOT_YET_IMPLEMENTED("VM: shadow_flush_pframe");
    // return -1;
    return 0;
}

/*
 * Clean up all resources associated with mobj o.
 *
 * Hints:
 *  - Check out mobj_put() to understand how this function gets called.
 *
 *  1) Call mobj_default_destructor() to flush o's pframes.
 *  2) Put the shadow and bottom_mobj members of the shadow object.
 *  3) Free the mobj_shadow_t.
 */
static void shadow_destructor(mobj_t *o)
{
    // NOT_YET_IMPLEMENTED("VM: shadow_destructor");
    mobj_default_destructor(o);

    mobj_shadow_t *so = MOBJ_TO_SO(o);
    // dbg(DBG_PROC, "shadow_destructor: refcount bottom: %d\n", so->bottom_mobj->mo_refcount);
    // dbg(DBG_PROC, "shadow_destructor: refcount: %d\n", so->shadowed->mo_refcount);

    // put the shadow and bottom_mobj members of the shadow object
    mobj_put(&so->shadowed);
    mobj_put(&so->bottom_mobj);

    slab_obj_free(shadow_allocator, so);

    return;
}