// SPDX-License-Identifier: MIT
/*
 * Copyright 2025 Advanced Micro Devices, Inc.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
 * THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
 * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
 * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 * OTHER DEALINGS IN THE SOFTWARE.
 *
 */

#include <linux/pci.h>
#include "amdgpu.h"
#include "amdgpu_ras.h"
#include "ras_sys.h"
#include "amdgpu_ras_cmd.h"
#include "amdgpu_virt_ras_cmd.h"
#include "amdgpu_ras_mgr.h"

static int amdgpu_virt_ras_remote_ioctl_cmd(struct ras_core_context *ras_core,
			struct ras_cmd_ctx *cmd, void *output_data, uint32_t output_size)
{
	struct amdgpu_device *adev = (struct amdgpu_device *)ras_core->dev;
	uint32_t mem_len = ALIGN(sizeof(*cmd) + output_size, AMDGPU_GPU_PAGE_SIZE);
	struct ras_cmd_ctx *rcmd;
	struct amdgpu_bo *rcmd_bo = NULL;
	uint64_t mc_addr = 0;
	void *cpu_addr = NULL;
	int ret = 0;

	ret = amdgpu_bo_create_kernel(adev, mem_len, PAGE_SIZE,
			AMDGPU_GEM_DOMAIN_VRAM, &rcmd_bo, &mc_addr, (void **)&cpu_addr);
	if (ret)
		return ret;

	rcmd = (struct ras_cmd_ctx *)cpu_addr;
	memset(rcmd, 0, mem_len);
	memcpy(rcmd, cmd, sizeof(*cmd));

	ret = amdgpu_virt_send_remote_ras_cmd(ras_core->dev,
				mc_addr - adev->gmc.vram_start, mem_len);
	if (!ret) {
		if (rcmd->cmd_res) {
			ret = rcmd->cmd_res;
			goto out;
		}

		cmd->cmd_res = rcmd->cmd_res;
		cmd->output_size = rcmd->output_size;
		if (rcmd->output_size && (rcmd->output_size <= output_size) && output_data)
			memcpy(output_data, rcmd->output_buff_raw, rcmd->output_size);
	}

out:
	amdgpu_bo_free_kernel(&rcmd_bo, &mc_addr, &cpu_addr);

	return ret;
}

static int amdgpu_virt_ras_send_remote_cmd(struct ras_core_context *ras_core,
	uint32_t cmd_id, void *input_data, uint32_t input_size,
	void *output_data, uint32_t output_size)
{
	struct ras_cmd_ctx rcmd = {0};
	int ret;

	rcmd.cmd_id = cmd_id;
	rcmd.input_size = input_size;
	memcpy(rcmd.input_buff_raw, input_data, input_size);

	ret = amdgpu_virt_ras_remote_ioctl_cmd(ras_core,
				&rcmd, output_data, output_size);
	if (!ret) {
		if (rcmd.output_size != output_size)
			return RAS_CMD__ERROR_GENERIC;
	}

	return ret;
}

static int amdgpu_virt_ras_get_batch_trace_overview(struct ras_core_context *ras_core,
	struct ras_log_batch_overview *overview)
{
	struct ras_cmd_batch_trace_snapshot_req req = {0};
	struct ras_cmd_batch_trace_snapshot_rsp rsp = {0};
	int ret;

	ret = amdgpu_virt_ras_send_remote_cmd(ras_core, RAS_CMD__GET_BATCH_TRACE_SNAPSHOT,
				&req, sizeof(req), &rsp, sizeof(rsp));
	if (ret)
		return ret;

	overview->first_batch_id = rsp.start_batch_id;
	overview->last_batch_id = rsp.latest_batch_id;
	overview->logged_batch_count = rsp.total_batch_num;

	return RAS_CMD__SUCCESS;
}

static int amdgpu_virt_ras_get_cper_snapshot(struct ras_core_context *ras_core,
			struct ras_cmd_ctx *cmd, void *data)
{
	struct amdgpu_ras_mgr *ras_mgr = amdgpu_ras_mgr_get_context(ras_core->dev);
	struct amdgpu_virt_ras_cmd *virt_ras =
			(struct amdgpu_virt_ras_cmd *)ras_mgr->virt_ras_cmd;
	int ret;

	if (cmd->input_size != sizeof(struct ras_cmd_cper_snapshot_req))
		return RAS_CMD__ERROR_INVALID_INPUT_SIZE;

	ret = amdgpu_virt_ras_send_remote_cmd(ras_core, cmd->cmd_id,
			cmd->input_buff_raw, cmd->input_size,
			cmd->output_buff_raw, sizeof(struct ras_cmd_cper_snapshot_rsp));
	if (ret)
		return ret;

	memset(&virt_ras->batch_mgr, 0, sizeof(virt_ras->batch_mgr));
	amdgpu_virt_ras_get_batch_trace_overview(ras_core,
					&virt_ras->batch_mgr.batch_overview);

	cmd->output_size = sizeof(struct ras_cmd_cper_snapshot_rsp);
	return RAS_CMD__SUCCESS;
}

static int amdgpu_virt_ras_get_batch_records(struct ras_core_context *ras_core, uint64_t batch_id,
			struct ras_log_info **trace_arr, uint32_t arr_num,
			struct ras_cmd_batch_trace_record_rsp *rsp_cache)
{
	struct ras_cmd_batch_trace_record_req req = {
		.start_batch_id = batch_id,
		.batch_num = RAS_CMD_MAX_BATCH_NUM,
	};
	struct ras_cmd_batch_trace_record_rsp *rsp = rsp_cache;
	struct batch_ras_trace_info *batch;
	int ret = 0;
	uint8_t i;

	if (!rsp->real_batch_num || (batch_id < rsp->start_batch_id) ||
		(batch_id >=  (rsp->start_batch_id + rsp->real_batch_num))) {

		memset(rsp, 0, sizeof(*rsp));
		ret = amdgpu_virt_ras_send_remote_cmd(ras_core, RAS_CMD__GET_BATCH_TRACE_RECORD,
			&req, sizeof(req), rsp, sizeof(*rsp));
		if (ret)
			return -EPIPE;
	}

	batch = &rsp->batchs[batch_id - rsp->start_batch_id];
	if (batch_id != batch->batch_id)
		return -ENODATA;

	for (i = 0; i < batch->trace_num; i++) {
		if (i >= arr_num)
			break;
		trace_arr[i] = &rsp->records[batch->offset + i];
	}

	return i;
}

static int amdgpu_virt_ras_get_cper_records(struct ras_core_context *ras_core,
	struct ras_cmd_ctx *cmd, void *data)
{
	struct amdgpu_ras_mgr *ras_mgr = amdgpu_ras_mgr_get_context(ras_core->dev);
	struct amdgpu_virt_ras_cmd *virt_ras =
			(struct amdgpu_virt_ras_cmd *)ras_mgr->virt_ras_cmd;
	struct ras_cmd_cper_record_req *req =
		(struct ras_cmd_cper_record_req *)cmd->input_buff_raw;
	struct ras_cmd_cper_record_rsp *rsp =
		(struct ras_cmd_cper_record_rsp *)cmd->output_buff_raw;
	struct ras_log_batch_overview *overview = &virt_ras->batch_mgr.batch_overview;
	struct ras_cmd_batch_trace_record_rsp *rsp_cache = &virt_ras->batch_mgr.batch_trace;
	struct ras_log_info **trace;
	uint32_t offset = 0, real_data_len = 0;
	uint64_t batch_id;
	uint8_t *out_buf;
	int ret = 0, i, count;

	if (cmd->input_size != sizeof(struct ras_cmd_cper_record_req))
		return RAS_CMD__ERROR_INVALID_INPUT_SIZE;

	if (!req->buf_size || !req->buf_ptr || !req->cper_num)
		return RAS_CMD__ERROR_INVALID_INPUT_DATA;

	trace = kzalloc_objs(*trace, MAX_RECORD_PER_BATCH);
	if (!trace)
		return RAS_CMD__ERROR_GENERIC;

	out_buf = kzalloc(req->buf_size, GFP_KERNEL);
	if (!out_buf) {
		kfree(trace);
		return RAS_CMD__ERROR_GENERIC;
	}

	memset(out_buf, 0, req->buf_size);

	for (i = 0; i < req->cper_num; i++) {
		batch_id = req->cper_start_id + i;
		if (batch_id >= overview->last_batch_id)
			break;
		count = amdgpu_virt_ras_get_batch_records(ras_core, batch_id,
							  trace, MAX_RECORD_PER_BATCH,
							  rsp_cache);
		if (count > 0) {
			ret = ras_cper_generate_cper(ras_core, trace, count,
					&out_buf[offset], req->buf_size - offset, &real_data_len);
			if (ret)
				break;

			offset += real_data_len;
		}
	}

	if ((ret && (ret != -ENOMEM)) ||
	    copy_to_user(u64_to_user_ptr(req->buf_ptr), out_buf, offset)) {
		kfree(out_buf);
		kfree(trace);
		return RAS_CMD__ERROR_GENERIC;
	}

	rsp->real_data_size = offset;
	rsp->real_cper_num = i;
	rsp->remain_num = (ret == -ENOMEM) ? (req->cper_num - i) : 0;
	rsp->version = 0;

	cmd->output_size = sizeof(struct ras_cmd_cper_record_rsp);

	kfree(out_buf);
	kfree(trace);

	return RAS_CMD__SUCCESS;
}

static int __fill_get_blocks_ecc_cmd(struct amdgpu_device *adev,
			struct vram_blocks_ecc *blks_ecc)
{
	struct ras_cmd_ctx *rcmd;

	if (!blks_ecc || !blks_ecc->bo || !blks_ecc->cpu_addr)
		return -EINVAL;

	rcmd = (struct ras_cmd_ctx *)blks_ecc->cpu_addr;

	rcmd->cmd_id = RAS_CMD__GET_ALL_BLOCK_ECC_STATUS;
	rcmd->input_size = sizeof(struct ras_cmd_blocks_ecc_req);
	rcmd->output_buf_size = blks_ecc->size - sizeof(*rcmd);

	return 0;
}

static int __set_cmd_auto_update(struct amdgpu_device *adev,
			enum ras_cmd_id cmd_id, uint64_t gpa_addr, uint32_t len, bool reg)
{
	struct ras_cmd_auto_update_req req = {0};
	struct ras_cmd_auto_update_rsp rsp = {0};
	int ret;

	req.mode = reg ? 1 : 0;
	req.cmd_id = cmd_id;
	req.addr = gpa_addr;
	req.len = len;
	ret = amdgpu_ras_mgr_handle_ras_cmd(adev, RAS_CMD__SET_CMD_AUTO_UPDATE,
		&req, sizeof(req), &rsp, sizeof(rsp));

	return ret;
}

static int amdgpu_virt_ras_get_block_ecc(struct ras_core_context *ras_core,
				struct ras_cmd_ctx *cmd, void *data)
{
	struct amdgpu_device *adev = ras_core->dev;
	struct amdgpu_ras_mgr *ras_mgr = amdgpu_ras_mgr_get_context(adev);
	struct amdgpu_virt_ras_cmd *virt_ras =
			(struct amdgpu_virt_ras_cmd *)ras_mgr->virt_ras_cmd;
	struct vram_blocks_ecc *blks_ecc = &virt_ras->blocks_ecc;
	struct ras_cmd_ctx *blks_ecc_cmd_ctx;
	struct ras_cmd_blocks_ecc_rsp *blks_ecc_rsp;
	struct ras_cmd_block_ecc_info_req *input_data =
			(struct ras_cmd_block_ecc_info_req *)cmd->input_buff_raw;
	struct ras_cmd_block_ecc_info_rsp *output_data =
			(struct ras_cmd_block_ecc_info_rsp *)cmd->output_buff_raw;
	int ret = 0;

	if (cmd->input_size != sizeof(struct ras_cmd_block_ecc_info_req))
		return RAS_CMD__ERROR_INVALID_INPUT_SIZE;

	if (input_data->block_id >= MAX_RAS_BLOCK_NUM)
		return RAS_CMD__ERROR_INVALID_INPUT_DATA;

	if (__fill_get_blocks_ecc_cmd(adev, blks_ecc))
		return RAS_CMD__ERROR_GENERIC;

	if (!virt_ras->blocks_ecc.auto_update_actived) {
		ret = __set_cmd_auto_update(adev, RAS_CMD__GET_ALL_BLOCK_ECC_STATUS,
				blks_ecc->mc_addr - adev->gmc.vram_start,
				blks_ecc->size, true);
		if (ret)
			return ret;

		blks_ecc->auto_update_actived = true;
	}

	blks_ecc_cmd_ctx = blks_ecc->cpu_addr;
	blks_ecc_rsp = (struct ras_cmd_blocks_ecc_rsp *)blks_ecc_cmd_ctx->output_buff_raw;

	output_data->ce_count = blks_ecc_rsp->blocks[input_data->block_id].ce_count;
	output_data->ue_count = blks_ecc_rsp->blocks[input_data->block_id].ue_count;
	output_data->de_count = blks_ecc_rsp->blocks[input_data->block_id].de_count;

	cmd->output_size = sizeof(struct ras_cmd_block_ecc_info_rsp);
	return RAS_CMD__SUCCESS;
}

static struct ras_cmd_func_map amdgpu_virt_ras_cmd_maps[] = {
	{RAS_CMD__GET_CPER_SNAPSHOT, amdgpu_virt_ras_get_cper_snapshot},
	{RAS_CMD__GET_CPER_RECORD, amdgpu_virt_ras_get_cper_records},
	{RAS_CMD__GET_BLOCK_ECC_STATUS, amdgpu_virt_ras_get_block_ecc},
};

int amdgpu_virt_ras_handle_cmd(struct ras_core_context *ras_core,
		struct ras_cmd_ctx *cmd)
{
	struct ras_cmd_func_map *ras_cmd = NULL;
	int i, res;

	for (i = 0; i < ARRAY_SIZE(amdgpu_virt_ras_cmd_maps); i++) {
		if (cmd->cmd_id == amdgpu_virt_ras_cmd_maps[i].cmd_id) {
			ras_cmd = &amdgpu_virt_ras_cmd_maps[i];
			break;
		}
	}

	if (ras_cmd)
		res = ras_cmd->func(ras_core, cmd, NULL);
	else
		res = amdgpu_virt_ras_remote_ioctl_cmd(ras_core, cmd,
					cmd->output_buff_raw, cmd->output_buf_size);

	cmd->cmd_res = res;

	if (cmd->output_size > cmd->output_buf_size) {
		RAS_DEV_ERR(ras_core->dev,
			"Output data size 0x%x exceeds buffer size 0x%x!\n",
			cmd->output_size, cmd->output_buf_size);
		return RAS_CMD__SUCCESS_EXEED_BUFFER;
	}

	return RAS_CMD__SUCCESS;
}

int amdgpu_virt_ras_sw_init(struct amdgpu_device *adev)
{
	struct amdgpu_ras_mgr *ras_mgr = amdgpu_ras_mgr_get_context(adev);

	ras_mgr->virt_ras_cmd = kzalloc_obj(struct amdgpu_virt_ras_cmd);
	if (!ras_mgr->virt_ras_cmd)
		return -ENOMEM;

	return 0;
}

int amdgpu_virt_ras_sw_fini(struct amdgpu_device *adev)
{
	struct amdgpu_ras_mgr *ras_mgr = amdgpu_ras_mgr_get_context(adev);

	kfree(ras_mgr->virt_ras_cmd);
	ras_mgr->virt_ras_cmd = NULL;

	return 0;
}

int amdgpu_virt_ras_hw_init(struct amdgpu_device *adev)
{
	struct amdgpu_ras_mgr *ras_mgr = amdgpu_ras_mgr_get_context(adev);
	struct amdgpu_virt_ras_cmd *virt_ras =
			(struct amdgpu_virt_ras_cmd *)ras_mgr->virt_ras_cmd;
	struct vram_blocks_ecc *blks_ecc = &virt_ras->blocks_ecc;

	amdgpu_virt_get_ras_capability(adev);

	memset(blks_ecc, 0, sizeof(*blks_ecc));
	blks_ecc->size = PAGE_SIZE;
	if (amdgpu_bo_create_kernel(adev, blks_ecc->size,
			PAGE_SIZE, AMDGPU_GEM_DOMAIN_VRAM,
			&blks_ecc->bo, &blks_ecc->mc_addr,
			(void **)&blks_ecc->cpu_addr))
		return -ENOMEM;

	return 0;
}

int amdgpu_virt_ras_hw_fini(struct amdgpu_device *adev)
{
	struct amdgpu_ras_mgr *ras_mgr = amdgpu_ras_mgr_get_context(adev);
	struct amdgpu_virt_ras_cmd *virt_ras =
			(struct amdgpu_virt_ras_cmd *)ras_mgr->virt_ras_cmd;
	struct vram_blocks_ecc *blks_ecc = &virt_ras->blocks_ecc;

	if (blks_ecc->bo) {
		__set_cmd_auto_update(adev,
			RAS_CMD__GET_ALL_BLOCK_ECC_STATUS,
			blks_ecc->mc_addr - adev->gmc.vram_start,
			blks_ecc->size, false);

		memset(blks_ecc->cpu_addr, 0, blks_ecc->size);
		amdgpu_bo_free_kernel(&blks_ecc->bo,
			&blks_ecc->mc_addr, &blks_ecc->cpu_addr);

		memset(blks_ecc, 0, sizeof(*blks_ecc));
	}

	return 0;
}

int amdgpu_virt_ras_pre_reset(struct amdgpu_device *adev)
{
	struct amdgpu_ras_mgr *ras_mgr = amdgpu_ras_mgr_get_context(adev);
	struct amdgpu_virt_ras_cmd *virt_ras =
		(struct amdgpu_virt_ras_cmd *)ras_mgr->virt_ras_cmd;

	virt_ras->blocks_ecc.auto_update_actived = false;
	return 0;
}

int amdgpu_virt_ras_post_reset(struct amdgpu_device *adev)
{
	return 0;
}

void amdgpu_virt_ras_set_remote_uniras(struct amdgpu_device *adev, bool en)
{
	struct amdgpu_ras_mgr *ras_mgr = amdgpu_ras_mgr_get_context(adev);
	struct amdgpu_virt_ras_cmd *virt_ras;

	if (!ras_mgr || !ras_mgr->virt_ras_cmd)
		return;

	virt_ras = (struct amdgpu_virt_ras_cmd *)ras_mgr->virt_ras_cmd;
	virt_ras->remote_uniras_supported = en;
}

bool amdgpu_virt_ras_remote_uniras_enabled(struct amdgpu_device *adev)
{
	struct amdgpu_ras_mgr *ras_mgr = amdgpu_ras_mgr_get_context(adev);
	struct amdgpu_virt_ras_cmd *virt_ras;

	if (amdgpu_in_reset(adev))
		return false;

	if (!ras_mgr || !ras_mgr->virt_ras_cmd)
		return false;

	virt_ras = (struct amdgpu_virt_ras_cmd *)ras_mgr->virt_ras_cmd;

	return virt_ras->remote_uniras_supported;
}