// SPDX-License-Identifier: (GPL-2.0-only OR MIT)
/*
 * Copyright (C) 2025 Amlogic, Inc. All rights reserved
 */
#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

#include "IONmem.h"
#include "cve.h"
#include "cve_alloc.h"
#include "dmabuf.h"

#define CVE_BUFF_NUM 32

struct cve_buffer *cve_buf = NULL;

static int cve_buffer_deinit(void)
{
    if (cve_buf != NULL) {
        if (cve_buf->dma_dev) {
            dmabuf_dev_close(cve_buf->dma_dev);
        }
        if (cve_buf->ion_dev) {
            ion_mem_exit(cve_buf->ion_dev);
        }
        if (cve_buf->ctx != NULL) {
            free(cve_buf->ctx);
        }
        free(cve_buf);
    }
    return 0;
}

static int cve_check_buffer_init(void)
{
    if (cve_buf == NULL) {
        cve_buf = malloc(sizeof(struct cve_buffer));
        if (cve_buf == NULL) {
            return -1;
        }
        memset(cve_buf, 0, sizeof(struct cve_buffer));
    }
    if (cve_buf->ctx == NULL) {
        cve_buf->ctx = malloc(sizeof(struct cve_ctx_s) * CVE_BUFF_NUM);
        if (cve_buf->ctx == NULL) {
            free(cve_buf);
            return -1;
        }
        memset(cve_buf->ctx, 0, sizeof(struct cve_ctx_s) * CVE_BUFF_NUM);
    }
    if (cve_buf->dma_dev == 0 && cve_buf->ion_dev == 0) {
        cve_buf->dma_dev = dmabuf_dev_open(GENERIC_HEAP);
        if (cve_buf->dma_dev > 0) {
            return 0;
        } else {
            cve_buf->dma_dev = 0;
            cve_buf->ion_dev = ion_mem_init();
            if (cve_buf->ion_dev < 0) {
                cve_buf->ion_dev = 0;
                cve_buffer_deinit();
                return -1;
            }
        }
    }
    return 0;
}

static int find_empty_buffer()
{
    int i;

    for (i = 0; i < CVE_BUFF_NUM; i++) {
        if (cve_buf->ctx[i].used == false) {
            cve_buf->ctx[i].used = true;
            return i;
        }
    }
    return -1;
}

static int find_used_buffer(U64 *u64VirAddr)
{
    int i;
    for (i = 0; i < CVE_BUFF_NUM; i++) {
        if (cve_buf->ctx[i].viraddr == (U64)u64VirAddr) {
            return i;
        }
    }
    return -1;
}

int ion_alloc_buffer(int ion_dev, struct cve_ctx_s *ctx, uint32_t type, size_t len)
{
    int ret = -1;
    int dma_fd = -1;

    switch (type) {
    case INPUT_BUFF_TYPE:
        ctx->cmemParm = malloc(sizeof(IONMEM_AllocParams));
        if (ctx->cmemParm == NULL)
            return -1;
        ret = ion_mem_alloc(ion_dev, len, ctx->cmemParm, false);
        if (ret < 0) {
            printf("%s,%d,Not enough memory\n", __func__, __LINE__);
            return -1;
        }
        dma_fd = ((IONMEM_AllocParams *)ctx->cmemParm)->mImageFd;
        ctx->i_len = len;
        ctx->i_buff = (char *)(mmap(NULL, len, PROT_READ | PROT_WRITE, MAP_SHARED, dma_fd, 0));
        if (!ctx->i_buff) {
            ctx->i_buff = NULL;
            printf("Failed to alloc i_buff:%s\n", strerror(errno));
            return -1;
        }
        ctx->viraddr = (U64)ctx->i_buff;
        ctx->dmafd = dma_fd;
        break;
    default:
        printf("Error no such buff type\n");
        break;
    }
    return dma_fd;
}

int dma_alloc_buffer(int dma_dev, struct cve_ctx_s *ctx, size_t len)
{
    int ret = -1;

    ret = dmabuf_heap_alloc(dma_dev, len, 0, &ctx->dmafd);
    if (ret < 0) {
        printf("%s,%d,Not enough memory\n", __func__, __LINE__);
        return -1;
    }
    ctx->i_len = len;
    ctx->i_buff = (char *)(mmap(NULL, len, PROT_READ | PROT_WRITE, MAP_SHARED, ctx->dmafd, 0));
    if (!ctx->i_buff) {
        ctx->i_buff = NULL;
        printf("Failed to alloc i_buff:%s\n", strerror(errno));
        return -1;
    }
    ctx->viraddr = (U64)ctx->i_buff;

    return ctx->dmafd;
}

int cve_alloc(int *dmafd, void **u64VirAddr, U32 len)
{
    int ret;
    int num;
    struct cve_ctx_s *ctx;

    ret = cve_check_buffer_init();
    if (ret) {
        return ret;
    }

    num = find_empty_buffer();
    if (num == -1) {
        return -1;
    }
    ctx = &cve_buf->ctx[num];

    ctx->size = len;

    if (cve_buf->dma_dev) {
        ret = dma_alloc_buffer(cve_buf->dma_dev, ctx, len);
    } else if (cve_buf->ion_dev) {
        ret = ion_alloc_buffer(cve_buf->ion_dev, ctx, INPUT_BUFF_TYPE, len);
    }
    else {
        ret = -1;
    }
    if (ret < 0) {
        ctx->used = false;
        ctx->size = 0;
        return -1;
    }
    memset((void*)ctx->viraddr, 0, ctx->size);
    *u64VirAddr = (void *)ctx->viraddr;
    *dmafd = ret;
    return 0;
}

int cve_free(void *u64VirAddr)
{
    int num;
    struct cve_ctx_s *ctx;

    num = find_used_buffer((U64 *)u64VirAddr);
    if (num == -1) {
        return -1;
    }
    ctx = &cve_buf->ctx[num];
    if (ctx->i_buff != NULL) {
        munmap(ctx->i_buff, ctx->i_len);
        ctx->i_buff = NULL;
        ctx->i_len = 0;
    }
    if (cve_buf->ion_dev) {
        ion_mem_free(cve_buf->ion_dev, (IONMEM_AllocParams *)ctx->cmemParm);
        if (ctx->cmemParm != NULL) {
            free((IONMEM_AllocParams *)ctx->cmemParm);
            ctx->cmemParm = NULL;
        }
    }
    if (cve_buf->dma_dev) {
        dmabuf_heap_close(ctx->dmafd);
    }

    memset(&cve_buf->ctx[num], 0, sizeof(struct cve_ctx_s));

    return 0;
}