package kk_pprof

import (
	"context"
	"gitee.com/cruvie/kk_go_kit/kk_func"
	"gitee.com/cruvie/kk_go_kit/kk_stage"
	"gitee.com/cruvie/kk_go_kit/kk_time"
	"log/slog"
	"os"
	"runtime/pprof"
	"sync"
	"time"
)

// StartProfileCollect 收集 CPU 使用情况
// dur 收集时间, period 收集间隔（可选）
func StartProfileCollect(dur time.Duration, period time.Duration) {
	stage := kk_stage.NewStage(context.Background(), kk_func.GetCurrentFunctionName(), true)
	var lock sync.Mutex
	if period == 0 {
		//only once
		go startProfileCollect(stage, dur, &lock)
	} else {
		for {
			go startProfileCollect(stage, dur, &lock)
			time.Sleep(period)
		}
	}
}
func startProfileCollect(stage *kk_stage.Stage, dur time.Duration, lock *sync.Mutex) {
	lock.Lock()
	timeStr := time.Now().Format(kk_time.DateTime)
	slog.Info("StartProfileCollect", kk_stage.NewLog(stage).
		Any("start time", timeStr).Args()...)
	fileName := "profile" + timeStr + ".prof"
	// 创建输出文件
	file, err := os.Create(fileName)
	if err != nil {

		slog.Error("Failed to create CPU profile", kk_stage.NewLog(stage).Error(err).Args()...)
		return
	}
	// 收集 CPU 使用情况
	errProfile := pprof.StartCPUProfile(file)
	defer func(file *os.File) {
		pprof.StopCPUProfile()
		lock.Unlock()
		err := file.Close()
		if err != nil {
			slog.Error("Failed to close CPU profile", kk_stage.NewLog(stage).Any("webUI", "go tool pprof -http=:9995 "+fileName).Error(err).Args()...)
		}
		slog.Info("StartProfileCollect finished", kk_stage.NewLog(stage).Any("webUI", "go tool pprof -http=:9995 "+fileName).Args()...)
	}(file)
	if errProfile != nil {
		slog.Error("Failed to start CPU profile", kk_stage.NewLog(stage).Error(errProfile).Args()...)
		return
	}
	// 等待一段时间以收集足够的数据
	time.Sleep(dur)
}

// StartHeapCollect 收集内存分配情况,只收集在被调用时正在使用的内存分配情况
func StartHeapCollect(period time.Duration) {
	stage := kk_stage.NewStage(context.Background(), kk_func.GetCurrentFunctionName(), true)
	if period == 0 {
		//only once
		startHeapCollect(stage)
	} else {
		for {
			go startHeapCollect(stage)
			time.Sleep(period)
		}
	}
}

func startHeapCollect(stage *kk_stage.Stage) {
	timeStr := time.Now().Format(kk_time.DateTime)

	slog.Info("StartHeapCollect", kk_stage.NewLog(stage).
		Any("start time", timeStr).Args()...)
	fileName := "heap" + timeStr + ".prof"
	// 创建输出文件
	file, err := os.Create(fileName)
	if err != nil {
		slog.Error("Failed to create memory profile", kk_stage.NewLog(stage).Error(err).Args()...)
		return
	}
	defer func(file *os.File) {
		err := file.Close()
		if err != nil {
			slog.Error("Failed to close memory profile", kk_stage.NewLog(stage).Any("webUI", "go tool pprof -http=:9996 "+fileName).Args()...)
		}
		slog.Info("StartHeapCollect finished", kk_stage.NewLog(stage).Any("webUI", "go tool pprof -http=:9996 "+fileName).Args()...)
	}(file)
	// 收集内存分配情况
	errHeap := pprof.WriteHeapProfile(file)
	if errHeap != nil {
		slog.Error("Failed to write memory profile", kk_stage.NewLog(stage).Error(errHeap).Args()...)
		return
	}
}
func StartGoroutineCollect(period time.Duration) {
	stage := kk_stage.NewStage(context.Background(), kk_func.GetCurrentFunctionName(), true)

	if period == 0 {
		startGoroutineCollect(stage)
	} else {
		for {
			go startGoroutineCollect(stage)
			time.Sleep(period)
		}
	}
}

func startGoroutineCollect(stage *kk_stage.Stage) {
	timeStr := time.Now().Format(kk_time.DateTime)

	slog.Info("StartGoroutineCollect", kk_stage.NewLog(stage).
		Any("start time", timeStr).Args()...)
	fileName := "goroutine" + timeStr + ".prof"
	// 创建输出文件
	file, err := os.Create(fileName)
	if err != nil {
		slog.Error("Failed to create goroutine profile", kk_stage.NewLog(stage).Error(err).Args()...)
		return
	}
	defer func(file *os.File) {
		err := file.Close()
		if err != nil {
			slog.Error("Failed to close goroutine profile", kk_stage.NewLog(stage).Any("webUI", "go tool pprof -http=:9997 "+fileName).Error(err).Args()...)
		}
		slog.Info("StartGoroutineCollect finished", kk_stage.NewLog(stage).
			Any("start time", timeStr).Args()...)
	}(file)
	// 收集 goroutine 信息
	errW := pprof.Lookup("goroutine").WriteTo(file, 0)
	if errW != nil {
		slog.Error("Failed to write goroutine profile", kk_stage.NewLog(stage).Error(errW).Args()...)
		return
	}
}