/*
 * Caliptra manifest image lookup and load address helpers.
 *
 * Copyright (C) 2026 ASPEED Technology Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <stdint.h>
#include <stddef.h>
#include <io.h>
#include <manifest.h>
#include <uart_console.h>

/* Define caliptra image identifier */
#define CPTRA_SOC_MANIFEST_HDR_ID (0x0002)
#define CPTRA_FMC_HDR_ID          (0x0003)
#define CPTRA_DDR4_IMEM_HDR_ID    (0x1000)
#define CPTRA_DDR4_DMEM_HDR_ID    (0x1001)
#define CPTRA_DDR4_2D_IMEM_HDR_ID (0x1002)
#define CPTRA_DDR4_2D_DMEM_HDR_ID (0x1003)
#define CPTRA_DDR5_IMEM_HDR_ID    (0x1004)
#define CPTRA_DDR5_DMEM_HDR_ID    (0x1005)
#define CPTRA_DP_FW_HDR_ID        (0x1006)
#define CPTRA_UEFI_HDR_ID         (0x1007)
#define CPTRA_ATF_HDR_ID          (0x1008)
#define CPTRA_OPTEE_HDR_ID        (0x1009)
#define CPTRA_UBOOT_HDR_ID        (0x100A)
#define CPTRA_SSP_HDR_ID          (0x100B)
#define CPTRA_TSP_HDR_ID          (0x100C)

/* Define caliptra image load address */
#define CPTRA_NO_LOAD_ADDR    (0x00000000)
#define CPTRA_FMC_LOAD_ADDR   (0x14b80000)
#define CPTRA_ATF_LOAD_ADDR   (0xb0000000)
#define CPTRA_OPTEE_LOAD_ADDR (0xb0080000)
#define CPTRA_UBOOT_LOAD_ADDR (0x80000000)
#define CPTRA_SSP_LOAD_ADDR   (0xac000000)
#define CPTRA_TSP_LOAD_ADDR   (0xae000000)

/* Define caliptra image loadable property */
#define CPTRA_LOADABLE_MASK GENMASK(31, 30)
#define CPTRA_BOOTMCU_LOADABLE (1)
#define CPTRA_SSP_LOADABLE (2)

struct cptra_load_image {
    char *name;
    uint32_t identifier;
    uint32_t fw_id;
    uintptr_t load_addr;
};

static struct cptra_load_image image_list[] = {
    { "manifest",
      CPTRA_SOC_MANIFEST_HDR_ID, CPTRA_MANIFEST_FW_ID, CPTRA_NO_LOAD_ADDR},
    { "mcu_fmc",
      CPTRA_FMC_HDR_ID, CPTRA_FMC_FW_ID, CPTRA_FMC_LOAD_ADDR},
    { "ddr4_imem",
      CPTRA_DDR4_IMEM_HDR_ID, CPTRA_DDR4_IMEM_FW_ID, CPTRA_NO_LOAD_ADDR},
    { "ddr4_dmem",
      CPTRA_DDR4_DMEM_HDR_ID, CPTRA_DDR4_DMEM_FW_ID, CPTRA_NO_LOAD_ADDR},
    { "ddr4_2d_imem",
      CPTRA_DDR4_2D_IMEM_HDR_ID, CPTRA_DDR4_2D_IMEM_FW_ID, CPTRA_NO_LOAD_ADDR},
    { "ddr4_2d_dmem",
      CPTRA_DDR4_2D_DMEM_HDR_ID, CPTRA_DDR4_2D_DMEM_FW_ID, CPTRA_NO_LOAD_ADDR},
    { "ddr5_imem",
      CPTRA_DDR5_IMEM_HDR_ID, CPTRA_DDR5_IMEM_FW_ID, CPTRA_NO_LOAD_ADDR},
    { "ddr5_dmem",
      CPTRA_DDR5_DMEM_HDR_ID, CPTRA_DDR5_DMEM_FW_ID, CPTRA_NO_LOAD_ADDR},
    { "dp_fw",
      CPTRA_DP_FW_HDR_ID, CPTRA_DP_FW_FW_ID, CPTRA_NO_LOAD_ADDR},
    { "uefi",
      CPTRA_UEFI_HDR_ID, CPTRA_UEFI_FW_ID, CPTRA_NO_LOAD_ADDR},
    { "atf",
      CPTRA_ATF_HDR_ID, CPTRA_ATF_FW_ID, CPTRA_ATF_LOAD_ADDR},
    { "optee",
      CPTRA_OPTEE_HDR_ID, CPTRA_OPTEE_FW_ID, CPTRA_OPTEE_LOAD_ADDR},
    { "uboot",
      CPTRA_UBOOT_HDR_ID, CPTRA_UBOOT_FW_ID, CPTRA_UBOOT_LOAD_ADDR},
    { "ssp",
      CPTRA_SSP_HDR_ID, CPTRA_SSP_FW_ID, CPTRA_SSP_LOAD_ADDR},
    { "tsp",
      CPTRA_TSP_HDR_ID, CPTRA_TSP_FW_ID, CPTRA_TSP_LOAD_ADDR},
};

static struct cptra_image_info *cptra_find_image_info(
            struct cptra_image_context *ctx, uint32_t identifier)
{
    bool match = false;
    uint32_t img_num = ctx->hdr->img_count;
    struct cptra_image_info *img_info = ctx->img_info;

    if (img_num > CPTRA_IMC_ENTRY_COUNT + 2) {
        return NULL;
    }

    for (img_info = ctx->img_info;
         img_info < ctx->img_info + img_num;
         img_info++) {
        if (img_info->identifier == identifier) {
            match = true;
            break;
        }
    }

    return match ? img_info : NULL;
}

static int cptra_img_info_get_offset(struct cptra_image_context *ctx,
                                     uint32_t identifier)
{
    struct cptra_image_info *img_info = NULL;

    img_info = cptra_find_image_info(ctx, identifier);

    return img_info ? (int)img_info->offset : CPTRA_ERR_IMAGE_OFFSET_INVALID;
}

static int cptra_img_info_get_size(struct cptra_image_context *ctx,
                                   uint32_t identifier)
{
    struct cptra_image_info *img_info = NULL;

    img_info = cptra_find_image_info(ctx, identifier);

    return img_info ? (int)img_info->size : CPTRA_ERR_IMAGE_SIZE_INVALID;
}

static struct cptra_load_image *cptra_find_load_image(uint32_t fw_id)
{
    int i = 0;
    int match = -1;

    for (i = 0; i < (int)ARRAY_SIZE(image_list) && match == -1; i++) {
        match = image_list[i].fw_id == fw_id ? i : match;
    }

    return match != -1 ? &image_list[match] : NULL;
}

bool cptra_ime_loadable_image(struct cptra_manifest_ime *ime)
{
    if (!ime) {
        return false;
    }

    if (FIELD_GET(CPTRA_LOADABLE_MASK, ime->flags) == CPTRA_BOOTMCU_LOADABLE) {
        return true;
    }

    return false;
}

int cptra_ime_image_offset(struct cptra_image_context *ctx, uint32_t fw_id)
{
    struct cptra_load_image *img = NULL;

    if (!ctx) {
        return CPTRA_ERR_IMAGE_OFFSET_INVALID;
    }

    img = cptra_find_load_image(fw_id);
    if (!img) {
        uprintf("Cannot find image with fw_id 0x%x.\n", fw_id);
        return CPTRA_ERR_IMAGE_OFFSET_INVALID;
    }

    return cptra_img_info_get_offset(ctx, img->identifier);
}

char *cptra_ime_get_image_name(uint32_t fw_id)
{
    struct cptra_load_image *img = NULL;

    img = cptra_find_load_image(fw_id);

    return img ? img->name : "Unknown";
}

int cptra_ime_image_size(struct cptra_image_context *ctx, uint32_t fw_id)
{
    struct cptra_load_image *img = NULL;

    if (!ctx) {
        return CPTRA_ERR_IMAGE_SIZE_INVALID;
    }

    img = cptra_find_load_image(fw_id);
    if (!img) {
        uprintf("Cannot find image with fw_id 0x%x.\n", fw_id);
        return CPTRA_ERR_IMAGE_SIZE_INVALID;
    }

    return cptra_img_info_get_size(ctx, img->identifier);
}

uintptr_t cptra_ime_get_load_addr(uint32_t fw_id)
{
    struct cptra_load_image *img = NULL;

    img = cptra_find_load_image(fw_id);
    if (!img) {
        uprintf("Cannot find image with fw_id 0x%x.\n", fw_id);
        return (uintptr_t)NULL;
    }

    return img->load_addr;
}

