#include <linux/kernel.h>
#include <linux/pci.h>
#include "nnt_device_defs.h"
#include "nnt_pci_conf_access_defs.h"
#include "nnt_pci_conf_access.h"
#include "nnt_defs.h"
#include "nnt_ioctl_defs.h"

/* Forward declarations */
int dma_mapping_page(unsigned int total_pinned,
                     unsigned int page_mapped_counter,
                     struct nnt_device* nnt_device,
                     struct page** current_pages,
                     struct nnt_page_info* page_info);
int pin_user_pages_in_kernel_space(unsigned int total_pages,
                                   unsigned int total_pinned,
                                   int* pinned_pages,
                                   struct nnt_device* nnt_device,
                                   struct nnt_page_info* page_info,
                                   struct page*** current_pages);
int pin_user_memory_in_kernel_space(unsigned int total_pages,
                                    struct nnt_device* nnt_device,
                                    struct nnt_page_info* page_info);
int map_dma_pages(struct nnt_page_info* page_info, struct nnt_device* nnt_device);
int release_dma_pages(struct nnt_page_info* page_info, struct nnt_device* nnt_device);

int dma_mapping_page(unsigned int total_pinned,
                     unsigned int page_mapped_counter,
                     struct nnt_device* nnt_device,
                     struct page** current_pages,
                     struct nnt_page_info* page_info)
{
    int current_page = total_pinned + page_mapped_counter;
    int error_code = 0;

    /* Get the dma address. */
    nnt_device->dma_page.dma_address_list[current_page] =
      dma_map_page(&nnt_device->pci_device->dev, current_pages[current_page], 0, PAGE_SIZE, DMA_BIDIRECTIONAL);
    /* Do we get a valid dma address ? */
    if (dma_mapping_error(&nnt_device->pci_device->dev, nnt_device->dma_page.dma_address_list[current_page]))
    {
        printk(KERN_ERR "Failed to get DMA addresses\n");
        error_code = -EINVAL;
        goto ReturnOnFinsihed;
    }

    page_info->page_address_array[current_page].dma_address = nnt_device->dma_page.dma_address_list[current_page];

ReturnOnFinsihed:
    return error_code;
}

int pin_user_pages_in_kernel_space(unsigned int total_pages,
                                   unsigned int total_pinned,
                                   int* pinned_pages,
                                   struct nnt_device* nnt_device,
                                   struct nnt_page_info* page_info,
                                   struct page*** current_pages)
{
    unsigned long page_pointer_start = page_info->page_pointer_start;
    int error_code = 0;

    /* Remaining pages to pin. */
    int remaining_pages = total_pages - total_pinned;

    /* Point to the current offset. */
    uint64_t current_page_pointer = page_pointer_start + (total_pinned * PAGE_SIZE);

    /* Save the current page. */
    *current_pages = nnt_device->dma_page.page_list + total_pinned;

    /* Returns number of pages pinned - this may be fewer than the number requested
         or -errno in case of error. */
    *pinned_pages = get_user_pages_fast(current_page_pointer, remaining_pages, FOLL_WRITE, *current_pages);

    if (*pinned_pages < 1)
    {
        kfree(nnt_device->dma_page.page_list);
        error_code = -EFAULT;
    }

    /* Allocate dma addresses structure. */
    if ((nnt_device->dma_page.dma_address_list = kcalloc(total_pages, sizeof(dma_addr_t), GFP_KERNEL)) == NULL)
    {
        error_code = -ENOMEM;
    }

    return error_code;
}

int pin_user_memory_in_kernel_space(unsigned int total_pages,
                                    struct nnt_device* nnt_device,
                                    struct nnt_page_info* page_info)
{
    unsigned int page_mapped_counter = 0;
    unsigned int total_pinned = 0;
    unsigned int pinned_pages;
    struct page** current_pages = NULL;
    int error_code = 0;

    while (total_pinned < total_pages)
    {
        /* Pinning user pages in kernel space. */
        if ((error_code = pin_user_pages_in_kernel_space(total_pages, total_pinned, &pinned_pages, nnt_device,
                                                         page_info, &current_pages)) != 0)
            goto ReturnOnFinished;

        /* When the parameter 'inter_iommu' is on, we need to set up
             a mapping on a pages in order to access the physical address. */
        while (page_mapped_counter < pinned_pages)
        {
            if ((error_code =
                   dma_mapping_page(total_pinned, page_mapped_counter, nnt_device, current_pages, page_info)) != 0)
                goto ReturnOnFinished;

            page_mapped_counter++;
        }

        /* Advance the memory that already pinned. */
        total_pinned += pinned_pages;
    }

    /* There is a page not pinned in the kernel space ? */
    if (total_pinned != total_pages)
    {
        return -EFAULT;
    }

ReturnOnFinished:
    return error_code;
}

int map_dma_pages(struct nnt_page_info* page_info, struct nnt_device* nnt_device)
{
    unsigned int total_pages = page_info->total_pages;
    int page_counter = 0;
    int error_code = 0;

    /* Check if we allow locking memory. */
    if (!can_do_mlock())
    {
        error_code = -EPERM;
        goto ReturnOnFinished;
    }

    /* Allocate the page list. */
    if ((nnt_device->dma_page.page_list = kcalloc(total_pages, sizeof(struct page*), GFP_KERNEL)) == NULL)
    {
        error_code = -ENOMEM;
        goto ReturnOnFinished;
    }

    /* Go over the user memory buffer and pin user pages in kernel space */
    if ((error_code = pin_user_memory_in_kernel_space(total_pages, nnt_device, page_info)) != 0)
        goto ReturnOnFinished;

    for (page_counter = 0; page_counter < total_pages; page_counter++)
    {
        printk(KERN_INFO "Page address structure number: %d, device: %04x:%02x:%02x.%0x\n", page_counter,
               pci_domain_nr(nnt_device->pci_device->bus), nnt_device->pci_device->bus->number,
               PCI_SLOT(nnt_device->pci_device->devfn), PCI_FUNC(nnt_device->pci_device->devfn));
    }

ReturnOnFinished:
    return error_code;
}

int release_dma_pages(struct nnt_page_info* page_info, struct nnt_device* nnt_device)
{
    int page_counter;

    if (nnt_device->dma_page.page_list == NULL)
    {
        return 0;
    }

    /* Deallocate the pages. */
    for (page_counter = 0; page_counter < page_info->total_pages; page_counter++)
    {
        /* DMA activity is finished. */
        dma_unmap_page(&nnt_device->pci_device->dev, nnt_device->dma_page.dma_address_list[page_counter], PAGE_SIZE,
                       DMA_BIDIRECTIONAL);

        /* Release the page list. */
        set_page_dirty(nnt_device->dma_page.page_list[page_counter]);
        put_page(nnt_device->dma_page.page_list[page_counter]);
        nnt_device->dma_page.page_list[page_counter] = NULL;
        nnt_device->dma_page.dma_address_list[page_counter] = 0;

        printk(KERN_INFO "Page structure number: %d was released. device:%04x:%02x:%02x.%0x\n", page_counter,
               pci_domain_nr(nnt_device->pci_device->bus), nnt_device->pci_device->bus->number,
               PCI_SLOT(nnt_device->pci_device->devfn), PCI_FUNC(nnt_device->pci_device->devfn));
    }

    // All the pages are clean.
    nnt_device->dma_page.page_list = NULL;

    return 0;
}
