package queue

import (
	"container/list"
	"fmt"
	"log"
	"math"
	"slices"
	"strconv"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/reimage-ai/broker/job"
)

const priorityCount = 10 // 0-9  0 is highest, 9 is lowest
const priorityLo = priorityCount - 1
const priorityHi = 0
const dropAgeTotalCutoffMin = 30
const dropAgeCheckedCutoffMin = 3

// a queue holds a list of list of workers, one list for each priority
// when delivering jobs higher priority workers are given jobs first
// note: 0 is the highest priority
type queue struct {
	jobs    *list.List
	workers *list.List
	lock    sync.Mutex
	count   uint64
}

// a receiver struct is placed in the waiting queue
// it is used to deliver the job to the worker
// and a context that allows for checking if the worker is still valid
// hints is used as a hint to the scheduler
// in some cases the schedule may prefer a lower priority worker, if they already have a warmed up cache
// for a specific task
type receiver struct {
	channel  chan *job.Job
	hints    []string
	priority int
}

// the lock for all stats
var statsLock sync.Mutex

// statsAvailable1Min stores the average number of available workers by job name over the last minute
var statsAvailable1Min = make(map[string]float64)

// statsQueue1Min stores the average length of job queue by job name over the last minute
var statsQueue1Min = make(map[string]float64)

// jobCheckChannel is used to pass jobs to the checker for sanity checks
var jobCheckerChannel = make(chan *job.Job, 1000)

// CurrentStateToString returns the current state of the queues as a string
func CurrentStateToString(statsMap map[string]map[string]any) string {
	str := "Queue State:\n"
	for k, v := range statsMap {
		str = str + fmt.Sprintf("%12.12s | workers: [%-12s]  queue: %2d  queue_load: %3.2f  avail_load: %3.2f  jobs: %6d \n", k, v["queueDescription"], v["queueLength"], v["queueLoad1Min"], v["availableLoad1Min"], v["count"])
	}
	return str
}

// CurrentState returns the current state of the queues as a string
func CurrentState() map[string]map[string]any {
	var statsMap = make(map[string]map[string]any)
	for name, queue := range getQueueMap() {
		statsMap[name] = make(map[string]any)

		queue.lock.Lock()
		rcvStr := ""
		for e := queue.workers.Front(); e != nil; e = e.Next() {
			rcv, ok := e.Value.(receiver)
			if !ok {
				log.Printf("Invalid type %v", e.Value)
			}
			rcvStr = rcvStr + strconv.FormatInt(int64(rcv.priority), 10)
		}
		statsMap[name]["queueDescription"] = rcvStr
		statsMap[name]["queueLength"] = queue.jobs.Len()
		statsMap[name]["count"] = atomic.LoadUint64(&queue.count)
		queue.lock.Unlock()

		statsLock.Lock()
		statsMap[name]["availableLoad1Min"] = statsAvailable1Min[name]
		statsMap[name]["queueLoad1Min"] = statsQueue1Min[name]
		statsLock.Unlock()
	}
	return statsMap
}

// Wait waits on a given job at the provided priority
func Wait(c *gin.Context, queueName string, priority int, timeoutMs int, worker_name string, hints []string) (*job.Job, error) {
	q, ok := getQueue(queueName)
	if !ok {
		return nil, fmt.Errorf("Unknown queue name %s", queueName)
	}
	if priority < priorityHi || priority > priorityLo {
		return nil, fmt.Errorf("Invalid priority %d", priority)
	}

	q.lock.Lock()
	// if jobs are awaiting, return first job
	for q.jobs.Len() > 0 {
		e := q.jobs.Front()
		q.jobs.Remove(e)

		// if job is super old, just drop it
		now := time.Now().UnixMilli()
		j := e.Value.(*job.Job)
		log.Printf("Queue: %s Size: %d (pop %v) (%.1fs delay)\n", queueName, q.jobs.Len(), j.ID, float64(now-j.GetQueueTime())/1000.0)

		if now-j.StartTime > 1000*60*dropAgeTotalCutoffMin {
			log.Printf("Dropping old job: %v (%d seconds old)\n", j.ID, (now-j.StartTime)/1000)
			j.UpdateStatus("error", map[string]any{"error": "abandoned - expired"})
			continue
		}
		// if the job is somewhat old and the client has not checked on the status recently, just drop it
		if now-j.GetLastCheckTime() > 1000*60*dropAgeCheckedCutoffMin {
			log.Printf("Dropping stale job: %v (%d seconds old)\n", j.ID, (now-j.StartTime)/1000)
			j.UpdateStatus("error", map[string]any{"error": "abandoned - stale"})
			continue
		}

		q.lock.Unlock()
		j.AssignWorker(worker_name)
		j.LogJobStart()
		atomic.AddUint64(&q.count, 1)
		return j, nil
	}

	// otherwise add channel to workers and wait until timeout
	rcv := receiver{
		// should only ever receive one job, but we use a buffered channel because it should never block the sender
		channel:  make(chan *job.Job, 1),
		hints:    hints,
		priority: priority,
	}

	element := q.workers.PushBack(rcv)
	q.lock.Unlock()

	check := func(timeoutMs int) *job.Job {
		select {
		case j := <-rcv.channel:
			j.AssignWorker(worker_name)
			j.LogJobStart()
			atomic.AddUint64(&q.count, 1)
			select {
			case jobCheckerChannel <- j:
			default:
				log.Printf("Warning: Job checker channel blocked")
			}
			return j
		case <-c.Done():
			//log.Printf("Client disconnected")
			return nil
		case <-c.Request.Context().Done():
			//log.Printf("Client disconnected")
			return nil
		case <-time.After(time.Duration(timeoutMs) * time.Millisecond):
			//log.Printf("Timeout")
			return nil
		}
	}

	j := check(timeoutMs)
	if j != nil {
		return j, nil
	}

	q.lock.Lock()
	q.workers.Remove(element)
	q.lock.Unlock()

	// Check one last time after acquiring lock
	return check(0), nil
}

// Put puts an job in the queue, or delivers directly to highest priority
// receiver if there is one currently waiting
func Put(queueName string, j *job.Job) error {
	q, ok := getQueue(queueName)
	if !ok {
		return fmt.Errorf("Unknown queue name %s", queueName)
	}

	j.AssignQueue(queueName)

	q.lock.Lock()
	defer q.lock.Unlock()

	var bestPriority int = 1000
	var bestReceiver *receiver = nil
	var bestElement *list.Element = nil

	for e := q.workers.Front(); e != nil; e = e.Next() {
		rcv, ok := e.Value.(receiver)
		if !ok {
			return fmt.Errorf("Invalid type %v", e.Value)
		}
		var priority = rcv.priority

		// slightly prioritize receivers with matching hint
		if rcv.hints != nil && slices.Contains(rcv.hints, j.Hint) {
			priority = priority - 4
		}

		if priority < bestPriority {
			bestReceiver = &rcv
			bestPriority = priority
			bestElement = e
		}
	}

	if bestReceiver != nil {
		// if delivering this to a worker with a matching hint, mark it as matched
		if bestReceiver.hints != nil && slices.Contains(bestReceiver.hints, j.Hint) {
			j.HintMatched = true
		}

		// deliver the job to receiver
		q.workers.Remove(bestElement)
		bestReceiver.channel <- j
	} else {
		// if no waiting workers, just add to jobs
		j.LogJobQueued()
		q.jobs.PushBack(j)
		log.Printf("Queue: %s Size: %d (add %v)\n", queueName, q.jobs.Len(), j.ID)
	}

	return nil
}

// QueueLength returns the current length of the queue
func QueueLength(queueName string) (int, error) {
	q, ok := getQueue(queueName)
	if !ok {
		return 0, fmt.Errorf("Unknown queue name %s", queueName)
	}
	q.lock.Lock()
	defer q.lock.Unlock()
	return q.jobs.Len(), nil
}

// JobChecker just checks jobs to assure they are being processed by workers
// it looks at jobs 8 seconds after they have been assigned to workers
// if the worker has not updated the status from "queued" then something is probably wrong
// just print out a warning
func JobChecker() {
	for job := range jobCheckerChannel {
		elapsedMs := time.Now().UnixMilli() - job.GetAssignedTime()
		for elapsedMs < 8000 {
			// sleep until 5 seconds after AssignedTime
			time.Sleep(time.Duration(8000-elapsedMs) * time.Millisecond)
			elapsedMs = time.Now().UnixMilli() - job.GetAssignedTime()
		}

		if job.GetStatus() == "queued" {
			log.Printf("Alert: Job %d picked up by %s %d ms ago but status is %s\n", job.ID, job.GetAssignedWorker(), elapsedMs, job.GetStatus())
			// in theory we could do something here like re-add, but that might cause conflicts if the worker starts?
			// j.UpdateStatus("error", 0.0, "dropped")
		}
	}
}

// Monitor calculates the average number of available workers per queue
// this is used to calculate when new workers are needed
func Monitor() {
	seconds := 5.0
	ticker := time.NewTicker(time.Second * time.Duration(5))

	for true {
		select {
		case <-ticker.C:
			// update stats
			for queueName, queue := range getQueueMap() {
				avail := 0.0
				queue.lock.Lock()
				avail = avail + float64(queue.workers.Len())
				qlen := queue.jobs.Len()
				queue.lock.Unlock()

				statsLock.Lock()
				var c float64 = math.Pow(2.71828, -seconds/60)
				statsAvailable1Min[queueName] = float64(avail)*(1-c) + statsAvailable1Min[queueName]*c
				statsQueue1Min[queueName] = float64(qlen)*(1-c) + statsQueue1Min[queueName]*c
				statsLock.Unlock()
			}
		}
	}

}

// queues stores the queues themselves (mapped by name)
var queues = make(map[string]*queue)

// the lock for queues
var queuesLock sync.Mutex

// getQueue retrieves the named queue, and creates one if necessary
func getQueue(name string) (*queue, bool) {
	queuesLock.Lock()
	defer queuesLock.Unlock()
	q, ok := queues[name]
	if !ok {
		addQueue(name)
		q, ok = queues[name]
		if !ok {
			log.Printf("ERROR: Failed to create queue: %s\n", name)
			return nil, false
		}
	}
	return q, true
}

// getQueueMap returns shallow copy of the map to avoid locking issues
func getQueueMap() map[string]*queue {
	queuesLock.Lock()
	defer queuesLock.Unlock()
	copy := make(map[string]*queue)
	for k, v := range queues {
		copy[k] = v
	}

	return copy
}

// addQueue adds the specified queueName
// must hold the lock before calling this function!
func addQueue(queueName string) {
	log.Printf("Creating queue: %s\n", queueName)
	workers := list.New()
	queues[queueName] = &queue{
		jobs:    list.New(),
		workers: workers,
	}
	statsLock.Lock()
	// start with a high value of available
	// this adjust within 1 minute to the actual value
	// starting with 0 can cause the node-manage to panic and fire up new nodes when we restart broker
	statsAvailable1Min[queueName] = 10.0
	statsQueue1Min[queueName] = 0.0
	statsLock.Unlock()
}
