package infrastructure

import (
	"encoding/json"
	"fmt"
	"sync"

	oclib "cloud.o-forge.io/core/oc-lib"
	"cloud.o-forge.io/core/oc-lib/models/common/enum"
	"cloud.o-forge.io/core/oc-lib/models/utils"
	"cloud.o-forge.io/core/oc-lib/models/workflow"
	"cloud.o-forge.io/core/oc-lib/models/workflow_execution"
	"cloud.o-forge.io/core/oc-lib/tools"
	"oc-scheduler/infrastructure/scheduling"
)

type executionConsidersPayload struct {
	ID           string   `json:"id"`
	ExecutionsID string   `json:"executions_id"`
	ExecutionID  string   `json:"execution_id"`
	PeerIDs      []string `json:"peer_ids"`
}

// ---------------------------------------------------------------------------
// Per-execution mutex map (replaces the global stateMu)
// ---------------------------------------------------------------------------

var execLocksMu sync.RWMutex
var execLocks = map[string]*sync.Mutex{} // executionID → per-execution mutex

// RegisterExecLock creates a mutex entry for the execution. Called when a new execution draft is persisted.
func RegisterExecLock(executionID string) {
	execLocksMu.Lock()
	execLocks[executionID] = &sync.Mutex{}
	execLocksMu.Unlock()
}

// UnregisterExecLock removes the mutex entry. Called on unschedule and execution deletion.
func UnregisterExecLock(executionID string) {
	execLocksMu.Lock()
	delete(execLocks, executionID)
	execLocksMu.Unlock()
}

// applyConsidersLocal applies the considers update directly for a confirmed
// booking or purchase (bypasses NATS since updateExecutionState resolves the
// execution from the resource itself).
func applyConsidersLocal(id string, dt tools.DataType) {
	payload, err := json.Marshal(&executionConsidersPayload{ID: id})
	if err != nil {
		return
	}
	updateExecutionState(payload, dt)
}

// EmitConsidersExecution broadcasts a Considers / WORKFLOW_EXECUTION message to all
// storage and compute peers of wf once the execution has transitioned to SCHEDULED.
// Each receiving peer will use it to confirm (IsDraft=false) their local drafts.
func EmitConsidersExecution(exec *workflow_execution.WorkflowExecution, wf *workflow.Workflow) {
	if wf == nil || wf.Graph == nil {
		return
	}
	peerIDs, err := GetWorkflowPeerIDs(wf.GetID(), &tools.APIRequest{Admin: true})
	if err != nil {
		return
	}
	if len(peerIDs) == 0 {
		return
	}
	payload, err := json.Marshal(executionConsidersPayload{
		ID:           exec.GetID(),
		ExecutionID:  exec.GetID(),
		ExecutionsID: exec.ExecutionsID,
		PeerIDs:      peerIDs})
	if err != nil {
		return
	}
	b, err := json.Marshal(tools.PropalgationMessage{
		DataType: int(tools.WORKFLOW_EXECUTION),
		Action:   tools.PB_CONSIDERS,
		Payload:  payload,
	})
	if err != nil {
		return
	}
	tools.NewNATSCaller().SetNATSPub(tools.PROPALGATION_EVENT, tools.NATSResponse{
		FromApp:  "oc-scheduler",
		Datatype: tools.WORKFLOW_EXECUTION,
		Method:   int(tools.PROPALGATION_EVENT),
		Payload:  b,
	})
}

// updateExecutionState sets BookingsState[id]=true (dt==BOOKING) or
// PurchasesState[id]=true (dt==PURCHASE_RESOURCE) on the target execution.
// payload must be JSON-encoded {"id":"...", "execution_id":"..."}.
func updateExecutionState(payload []byte, dt tools.DataType) {
	var data executionConsidersPayload
	if err := json.Unmarshal(payload, &data); err != nil || data.ID == "" {
		return
	}
	schdata := oclib.NewRequestAdmin(oclib.LibDataEnum(dt), nil).LoadOne(data.ID)
	if schdata.Data == nil {
		return
	}
	sch := scheduling.ToSchedulerObject(dt, schdata.Data)
	if sch == nil {
		return
	}
	execID := sch.GetExecutionId()

	execLocksMu.RLock()
	mu := execLocks[execID]
	execLocksMu.RUnlock()
	if mu == nil {
		fmt.Printf("updateExecutionState: no lock for execution %s, skipping\n", execID)
		return
	}
	mu.Lock()
	defer mu.Unlock()

	adminReq := &tools.APIRequest{Admin: true}
	res, _, err := workflow_execution.NewAccessor(adminReq).LoadOne(execID)
	if err != nil || res == nil {
		fmt.Printf("updateExecutionState: could not load execution %s: %v\n", data.ExecutionID, err)
		return
	}

	exec := res.(*workflow_execution.WorkflowExecution)
	fmt.Println("sch.GetExecutionId()", data.ID, exec.BookingsState)

	switch dt {
	case tools.BOOKING:
		if exec.BookingsState == nil {
			exec.BookingsState = map[string]bool{}
		}
		exec.BookingsState[data.ID] = true
		fmt.Println("sch.GetExecutionId()", data.ID)

	case tools.PURCHASE_RESOURCE:
		if exec.PurchasesState == nil {
			exec.PurchasesState = map[string]bool{}
		}
		exec.PurchasesState[data.ID] = true
	}
	allConfirmed := true
	for _, st := range exec.BookingsState {
		if !st {
			allConfirmed = false
			break
		}
	}
	for _, st := range exec.PurchasesState {
		if !st {
			allConfirmed = false
			break
		}
	}
	if allConfirmed {
		exec.State = enum.SCHEDULED
		exec.IsDraft = false
	}
	if _, _, err := utils.GenericRawUpdateOne(exec, exec.GetID(), workflow_execution.NewAccessor(adminReq)); err != nil {
		fmt.Printf("updateExecutionState: could not update execution %s: %v\n", sch.GetExecutionId(), err)
		return
	}
	if allConfirmed {
		// Confirm the order and notify all peers that execution is scheduled.
		go confirmSessionOrder(exec.ExecutionsID, adminReq)
		obj, _, err := workflow.NewAccessor(adminReq).LoadOne(exec.WorkflowID)
		if err == nil && obj != nil {
			go EmitConsidersExecution(exec, obj.(*workflow.Workflow))
		}
	}
}

// confirmExecutionDrafts is called when a Considers/WORKFLOW_EXECUTION message
// is received from oc-discovery, meaning the originating peer has confirmed the
// execution as SCHEDULED. For every booking and purchase ID listed in the
// execution's states, we confirm the local draft (IsDraft=false).
func confirmExecutionDrafts(payload []byte) {
	var data executionConsidersPayload
	if err := json.Unmarshal(payload, &data); err != nil {
		fmt.Printf("confirmExecutionDrafts: could not parse payload: %v\n", err)
		return
	}
	access := oclib.NewRequestAdmin(oclib.LibDataEnum(tools.WORKFLOW_EXECUTION), nil)
	d := access.LoadOne(data.ExecutionID)
	if exec := d.ToWorkflowExecution(); exec != nil {
		for id := range exec.BookingsState {
			go confirmResource(id, tools.BOOKING)
		}
		for id := range exec.PurchasesState {
			go confirmResource(id, tools.PURCHASE_RESOURCE)
		}
	}
}