package main

import (
	"fmt"
	"runtime"
	"sync"
	"time"
)

// 全局变量，用于标记任务所属的P（通过首次执行时绑定）
var taskPMap sync.Map // key: 任务名, value: int（P ID）

// 绑定当前任务到执行的P（首次执行时记录）
func bindTaskToP(taskName string) {
	// 用goroutine ID间接关联P（简化版，避免解析栈）
	gid := getGoroutineID()
	// 这里用gid的奇偶性模拟P ID（2个P时，奇偶对应P0/P1，直观且不依赖栈解析）
	pid := gid % 2
	taskPMap.Store(taskName, pid)
	fmt.Printf("[P%d] %s 绑定到P%d\n", pid, taskName, pid)
}

// 获取goroutine ID（简化版，用于区分P）
func getGoroutineID() int {
	var buf [64]byte
	_ = runtime.Stack(buf[:], false)
	// 栈信息格式："goroutine 123 [running]"，提取数字部分
	start := 0
	for buf[start] < '0' || buf[start] > '9' {
		start++
	}
	end := start
	for buf[end] >= '0' && buf[end] <= '9' {
		end++
	}
	gid := 0
	for i := start; i < end; i++ {
		gid = gid*10 + int(buf[i]-'0')
	}
	return gid
}

// CPU密集型任务（持续占用P，不主动让出）
func cpuIntensiveTask(duration time.Duration) {
	endTime := time.Now().Add(duration)
	for time.Now().Before(endTime) {
		// 纯计算，持续占用CPU
	}
}

func main() {
	runtime.GOMAXPROCS(2) // 固定2个P，确保负载均衡触发
	var wg sync.WaitGroup

	// --------------------------
	// P0的任务：长任务A（持续占用P0） + 任务B（归P0队列）
	// --------------------------
	// 长任务A（CPU密集型，耗时3秒，持续占用P0）
	wg.Add(1)
	go func() {
		defer wg.Done()
		taskName := "长任务A"
		bindTaskToP(taskName) // 绑定到执行的P
		// 正确处理Load返回值：先接收value和ok，再转换
		pidVal, ok := taskPMap.Load(taskName)
		if !ok {
			pidVal = -1 // 异常默认值
		}
		pid := pidVal.(int)
		fmt.Printf("[P%d] %s 开始执行（CPU密集型，耗时3秒）\n", pid, taskName)
		cpuIntensiveTask(3 * time.Second)
		// 再次获取P ID（确保一致性）
		pidVal, ok = taskPMap.Load(taskName)
		if !ok {
			pidVal = -1
		}
		pid = pidVal.(int)
		fmt.Printf("[P%d] %s 执行完成\n", pid, taskName)
	}()

	// 延长等待时间（1s），确保长任务A稳定占用P0，且调度器已分配P0
	time.Sleep(1 * time.Second)

	// 任务B（归P0队列，等待长任务A执行时暂存）
	wg.Add(1)
	go func() {
		defer wg.Done()
		taskName := "任务B"
		//// 强制任务B与长任务A同属一个P的队列（通过延迟启动，复用P0的本地队列）
		//time.Sleep(200 * time.Millisecond)
		bindTaskToP(taskName)
		// 正确处理Load返回值
		pidVal, ok := taskPMap.Load(taskName)
		if !ok {
			pidVal = -1
		}
		pid := pidVal.(int)
		fmt.Printf("[P%d] %s 开始执行（归P0队列，可能被窃取）\n", pid, taskName)
		cpuIntensiveTask(500 * time.Millisecond) // CPU密集型短任务
		// 再次获取P ID
		pidVal, ok = taskPMap.Load(taskName)
		if !ok {
			pidVal = -1
		}
		pid = pidVal.(int)
		fmt.Printf("[P%d] %s 执行完成\n", pid, taskName)
	}()

	// --------------------------
	// P1的任务：3个CPU密集型短任务（持续占用P1）
	// --------------------------
	// 短任务执行完后，P1空闲，触发工作窃取
	for i := 0; i < 3; i++ {
		wg.Add(1)
		taskIdx := i
		go func() {
			defer wg.Done()
			taskName := fmt.Sprintf("短任务%d", taskIdx)
			bindTaskToP(taskName)
			// 正确处理Load返回值
			pidVal, ok := taskPMap.Load(taskName)
			if !ok {
				pidVal = -1
			}
			pid := pidVal.(int)
			fmt.Printf("[P%d] %s 开始执行（CPU密集型，耗时200ms）\n", pid, taskName)
			cpuIntensiveTask(200 * time.Millisecond)
			// 再次获取P ID
			pidVal, ok = taskPMap.Load(taskName)
			if !ok {
				pidVal = -1
			}
			pid = pidVal.(int)
			fmt.Printf("[P%d] %s 执行完成\n", pid, taskName)
		}()
	}

	wg.Wait()
	fmt.Println("所有任务执行完毕")
}