// SPDX-License-Identifier: GPL-2.0
/*
 * Copyright (C) 2025 Ant Group
 * Author: Tiwei Bie <tiwei.btw@antgroup.com>
 *
 * Based on the previous implementation in TT mode
 * Copyright (C) 2000 - 2007 Jeff Dike (jdike@{addtoit,linux.intel}.com)
 */

#include <linux/sched.h>
#include <linux/sched/task.h>
#include <linux/sched/task_stack.h>
#include <linux/module.h>
#include <linux/processor.h>
#include <linux/threads.h>
#include <linux/cpu.h>
#include <linux/hardirq.h>
#include <linux/smp.h>
#include <linux/smp-internal.h>
#include <init.h>
#include <kern.h>
#include <os.h>
#include <smp.h>

enum {
	UML_IPI_RES = 0,
	UML_IPI_CALL_SINGLE,
	UML_IPI_CALL,
	UML_IPI_STOP,
};

void arch_smp_send_reschedule(int cpu)
{
	os_send_ipi(cpu, UML_IPI_RES);
}

void arch_send_call_function_single_ipi(int cpu)
{
	os_send_ipi(cpu, UML_IPI_CALL_SINGLE);
}

void arch_send_call_function_ipi_mask(const struct cpumask *mask)
{
	int cpu;

	for_each_cpu(cpu, mask)
		os_send_ipi(cpu, UML_IPI_CALL);
}

void smp_send_stop(void)
{
	int cpu, me = smp_processor_id();

	for_each_online_cpu(cpu) {
		if (cpu == me)
			continue;
		os_send_ipi(cpu, UML_IPI_STOP);
	}
}

static void ipi_handler(int vector, struct uml_pt_regs *regs)
{
	struct pt_regs *old_regs = set_irq_regs((struct pt_regs *)regs);
	int cpu = raw_smp_processor_id();

	irq_enter();

	if (current->mm)
		os_alarm_process(current->mm->context.id.pid);

	switch (vector) {
	case UML_IPI_RES:
		inc_irq_stat(irq_resched_count);
		scheduler_ipi();
		break;

	case UML_IPI_CALL_SINGLE:
		inc_irq_stat(irq_call_count);
		generic_smp_call_function_single_interrupt();
		break;

	case UML_IPI_CALL:
		inc_irq_stat(irq_call_count);
		generic_smp_call_function_interrupt();
		break;

	case UML_IPI_STOP:
		set_cpu_online(cpu, false);
		while (1)
			pause();
		break;

	default:
		pr_err("CPU#%d received unknown IPI (vector=%d)!\n", cpu, vector);
		break;
	}

	irq_exit();
	set_irq_regs(old_regs);
}

void uml_ipi_handler(int vector)
{
	struct uml_pt_regs r = { .is_user = 0 };

	preempt_disable();
	ipi_handler(vector, &r);
	preempt_enable();
}

/* AP states used only during CPU startup */
enum {
	UML_CPU_PAUSED = 0,
	UML_CPU_RUNNING,
};

static int cpu_states[NR_CPUS];

static int start_secondary(void *unused)
{
	int err, cpu = raw_smp_processor_id();

	notify_cpu_starting(cpu);
	set_cpu_online(cpu, true);

	err = um_setup_timer();
	if (err)
		panic("CPU#%d failed to setup timer, err = %d", cpu, err);

	local_irq_enable();

	cpu_startup_entry(CPUHP_AP_ONLINE_IDLE);

	return 0;
}

void uml_start_secondary(void *opaque)
{
	int cpu = raw_smp_processor_id();
	struct mm_struct *mm = &init_mm;
	struct task_struct *idle;

	stack_protections((unsigned long) &cpu_irqstacks[cpu]);
	set_sigstack(&cpu_irqstacks[cpu], THREAD_SIZE);

	set_cpu_present(cpu, true);
	os_futex_wait(&cpu_states[cpu], UML_CPU_PAUSED);

	smp_rmb(); /* paired with smp_wmb() in __cpu_up() */

	idle = cpu_tasks[cpu];
	idle->thread_info.cpu = cpu;

	mmgrab(mm);
	idle->active_mm = mm;

	idle->thread.request.thread.proc = start_secondary;
	idle->thread.request.thread.arg = NULL;

	new_thread(task_stack_page(idle), &idle->thread.switch_buf,
		   new_thread_handler);
	os_start_secondary(opaque, &idle->thread.switch_buf);
}

void __init smp_prepare_cpus(unsigned int max_cpus)
{
	int err, cpu, me = smp_processor_id();
	unsigned long deadline;

	os_init_smp();

	for_each_possible_cpu(cpu) {
		if (cpu == me)
			continue;

		pr_debug("Booting processor %d...\n", cpu);
		err = os_start_cpu_thread(cpu);
		if (err) {
			pr_crit("CPU#%d failed to start cpu thread, err = %d",
				cpu, err);
			continue;
		}

		deadline = jiffies + msecs_to_jiffies(1000);
		spin_until_cond(cpu_present(cpu) ||
				time_is_before_jiffies(deadline));

		if (!cpu_present(cpu))
			pr_crit("CPU#%d failed to boot\n", cpu);
	}
}

int __cpu_up(unsigned int cpu, struct task_struct *tidle)
{
	cpu_tasks[cpu] = tidle;
	smp_wmb(); /* paired with smp_rmb() in uml_start_secondary() */
	cpu_states[cpu] = UML_CPU_RUNNING;
	os_futex_wake(&cpu_states[cpu]);
	spin_until_cond(cpu_online(cpu));

	return 0;
}

void __init smp_cpus_done(unsigned int max_cpus)
{
}

/* Set in uml_ncpus_setup */
int uml_ncpus = 1;

void __init prefill_possible_map(void)
{
	int cpu;

	for (cpu = 0; cpu < uml_ncpus; cpu++)
		set_cpu_possible(cpu, true);
	for (; cpu < NR_CPUS; cpu++)
		set_cpu_possible(cpu, false);
}

static int __init uml_ncpus_setup(char *line, int *add)
{
	*add = 0;

	if (kstrtoint(line, 10, &uml_ncpus)) {
		os_warn("%s: Couldn't parse '%s'\n", __func__, line);
		return -1;
	}

	uml_ncpus = clamp(uml_ncpus, 1, NR_CPUS);

	return 0;
}

__uml_setup("ncpus=", uml_ncpus_setup,
"ncpus=<# of desired CPUs>\n"
"    This tells UML how many virtual processors to start. The maximum\n"
"    number of supported virtual processors can be obtained by querying\n"
"    the CONFIG_NR_CPUS option using --showconfig.\n\n"
);

EXPORT_SYMBOL(uml_curr_cpu);