/*
** Job Arranger for ZABBIX
** Copyright (C) 2025 Daiwa Institute of Research Ltd. All Rights Reserved.
**
** This program is free software; you can redistribute it and/or modify
** it under the terms of the GNU General Public License as published by
** the Free Software Foundation; either version 2 of the License, or
** (at your option) any later version.
**
** This program is distributed in the hope that it will be useful,
** but WITHOUT ANY WARRANTY; without even the implied warranty of
** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
** GNU General Public License for more details.
**
** You should have received a copy of the GNU General Public License
** along with this program; if not, write to the Free Software
** Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
**/
package workermanager

import (
	"context"
	"maps"
	"runtime/debug"
	"sync"
	"time"

	"jobarranger2/src/libs/golibs/logger/logger"
)

var (
	heartbeatLockOwner, workerLockOwner string
	ukeyLock                            sync.Mutex
	heartbeatLock                       sync.Mutex
	seenKeys                            = make(map[string]string)
	heartbeatMap                        = make(map[string]time.Time)
	ProcessExitFlag                     bool
	Wm                                  *WorkerManager
)

const (
	daemonWorker       = 0
	oneTimeWorker      = 1
	oneTimeLimitWorker = 2
)

type Worker struct {
	Id         string
	Timeout    int
	WorkerType int
	StopChan   chan bool
}

type WorkerManager struct {
	workers                 map[string]*Worker
	mu                      sync.Mutex
	Wg                      sync.WaitGroup
	Ctx                     context.Context
	Cancel                  context.CancelFunc
	MonitorChan             chan string
	OnetimeLimitWorkerCount int
	OnetimeWorkerLimit      int
}

func InitializeWorkerManager(onetimeWorkerLimit int) {
	ctx, cancel := context.WithCancel(context.Background())
	Wm = &WorkerManager{
		workers:            make(map[string]*Worker),
		Ctx:                ctx,
		Cancel:             cancel,
		MonitorChan:        make(chan string, 20),
		OnetimeWorkerLimit: onetimeWorkerLimit,
	}
}

func StopAllWorkers(doneChan chan bool) {
	Wm.Cancel()
	go func() {
		Wm.Wg.Wait()
		close(Wm.MonitorChan)
		close(doneChan)
	}()
}

func StartWorker(workerFunc func(), id string, timeout int, workerType int, manager string) {
	stopChan := make(chan bool)
	worker := &Worker{
		Id:         id,
		Timeout:    timeout,
		WorkerType: workerType,
		StopChan:   stopChan,
	}

	// Send initial heartbeat
	Wm.MonitorChan <- id

	Wm.mu.Lock()
	defer Wm.mu.Unlock()
	Wm.workers[id] = worker
	Wm.Wg.Add(1)
	if workerType == oneTimeLimitWorker {
		Wm.OnetimeLimitWorkerCount++
	}

	go func() {
		defer func() {
			// Catch panic and log stacktrace
			if r := recover(); r != nil {
				logger.WriteLog("JAFRAMEWORK100002", manager, id, string(debug.Stack()))
			}

			// this will later be used as memory monitoring and gc status

			// var m runtime.MemStats
			// runtime.ReadMemStats(&m)
			// logger.WriteLog("JADEBUG000001", "StartWorker",
			// 	m.HeapAlloc/1048576,             // HeapAlloc in MB
			// 	m.NextGC/1048576,                // Approx next GC threshold in MB
			// 	(m.HeapSys-m.HeapInuse)/1048576, // heap free for reuse in MB
			// 	m.HeapInuse/1048576,             // heap actually in use in MB
			// 	m.HeapSys/1048576,               // total heap obtained from OS in MB
			// 	m.HeapReleased/1048576)          // heap returned to OS in MB

			// JADEBUG000001,0,0,In %s(),
			// HeapAlloc: %d, Approx. next GC threshold: %d bytes,
			// heap free for reuse: %d, heap actually in use: %d,
			// total heap obtained from OS: %d, heap returned to OS: %d,

			Wm.Wg.Done()
			RemoveWorker(id)
			logger.JaLog("JAFRAMEWORK400004", logger.Logging{}, manager, id)
		}()
		workerFunc()
	}()
}

func RemoveWorker(id string) {
	Wm.mu.Lock()
	heartbeatLock.Lock()
	defer Wm.mu.Unlock()
	defer heartbeatLock.Unlock()
	if worker, exists := Wm.workers[id]; exists {

		// Cancel removing from watch list if worker is daemon worker and server is still running
		if worker.WorkerType == daemonWorker && !ProcessExitFlag {
			return
		}

		if worker.WorkerType == oneTimeLimitWorker {
			Wm.OnetimeLimitWorkerCount--
		}
		close(worker.StopChan)

		// Remove from worker map
		delete(Wm.workers, id)

		// Remove from heartbeat map
		delete(heartbeatMap, id)
	}
}

func RegisterUniqueKey(key string, filePath string) {
	ukeyLock.Lock()
	seenKeys[key] = filePath
	ukeyLock.Unlock()
}

func RemoveUniqueKey(key string) {
	ukeyLock.Lock()
	delete(seenKeys, key)
	ukeyLock.Unlock()
}

func IsUniqueKeyExists(key string) bool {
	ukeyLock.Lock()
	_, ok := seenKeys[key]
	ukeyLock.Unlock()
	return ok
}

func GetAllUniqueKeys() map[string]string {
	ukeyLock.Lock()
	defer ukeyLock.Unlock()

	copy := make(map[string]string)
	maps.Copy(copy, seenKeys)
	return copy
}

func ClearAllUniqueKeys() {
	// Clear unique keys
	ukeyLock.Lock()
	seenKeys = make(map[string]string)
	ukeyLock.Unlock()
}

func GetUniqueKeyCount() int {
	ukeyLock.Lock()
	count := len(seenKeys)
	ukeyLock.Unlock()
	return count
}

func CheckWorkerExist(id string) bool {
	Wm.mu.Lock()
	defer Wm.mu.Unlock()

	_, exists := Wm.workers[id]
	return exists
}

func GetWorker(id string) (*Worker, bool) {
	Wm.mu.Lock()
	defer Wm.mu.Unlock()

	worker, ok := Wm.workers[id]
	return worker, ok
}

func RegisterHeartbeat(id string) {
	heartbeatLock.Lock()
	defer heartbeatLock.Unlock()

	heartbeatMap[id] = time.Now()
}

func GetHeartbeatMapCopy() map[string]time.Time {
	heartbeatLock.Lock()
	defer heartbeatLock.Unlock()

	copyMap := make(map[string]time.Time, len(heartbeatMap))
	maps.Copy(copyMap, heartbeatMap)
	return copyMap
}
