package handlers

import (
	"context"
	"encoding/json"
	"log"

	"go.starlark.net/starlark"

	"armature/sandbox"
	"armature/store"
)

// ── on_advance Hook ─────────────────────────
//
// Fires synchronously after a stage transition succeeds.
// The hook can enrich/transform stage_data or reject the transition.
//
// Config in stage_config:
//   {"on_advance": {"package_id": "...", "entry_point": "on_advance"}}
//
// Hook receives dict: {stage_data, previous_stage, current_stage, instance_id}
// Hook returns: {stage_data: {...}} (enriched), None (no change),
//               or {error: "msg"} (reject — caller should handle rollback)

// OnAdvanceHookConfig is the on_advance section of stage_config.
type OnAdvanceHookConfig struct {
	PackageID  string `json:"package_id"`
	EntryPoint string `json:"entry_point"`
}

// StageConfigOnAdvance extracts on_advance config from stage_config JSON.
type StageConfigOnAdvance struct {
	OnAdvance *OnAdvanceHookConfig `json:"on_advance,omitempty"`
}

// OnAdvanceResult is the outcome of firing the on_advance hook.
type OnAdvanceResult struct {
	EnrichedData json.RawMessage // updated stage_data (nil = no change)
	Error        string          // non-empty = hook rejected the transition
}

// FireOnAdvanceHook fires the on_advance Starlark hook for a stage transition.
// Returns nil if no hook is configured or the runner is unavailable.
func FireOnAdvanceHook(
	ctx context.Context,
	stores store.Stores,
	runner *sandbox.Runner,
	previousStageConfig json.RawMessage,
	instanceID string,
	previousStage, currentStage int,
	stageData json.RawMessage,
) *OnAdvanceResult {
	if runner == nil || stores.Packages == nil {
		return nil
	}

	var cfg StageConfigOnAdvance
	if len(previousStageConfig) > 0 {
		_ = json.Unmarshal(previousStageConfig, &cfg)
	}
	if cfg.OnAdvance == nil || cfg.OnAdvance.PackageID == "" || cfg.OnAdvance.EntryPoint == "" {
		return nil
	}

	pkg, err := stores.Packages.Get(ctx, cfg.OnAdvance.PackageID)
	if err != nil || pkg == nil {
		log.Printf("[workflow-hooks] on_advance: package %s not found", cfg.OnAdvance.PackageID)
		return nil
	}

	// Build context dict for the hook
	ctxDict := starlark.NewDict(4)
	_ = ctxDict.SetKey(starlark.String("instance_id"), starlark.String(instanceID))
	_ = ctxDict.SetKey(starlark.String("previous_stage"), starlark.MakeInt(previousStage))
	_ = ctxDict.SetKey(starlark.String("current_stage"), starlark.MakeInt(currentStage))

	// Parse stage_data into Starlark dict
	var dataMap map[string]interface{}
	if json.Unmarshal(stageData, &dataMap) == nil {
		_ = ctxDict.SetKey(starlark.String("stage_data"), sandbox.GoToStarlark(dataMap))
	} else {
		_ = ctxDict.SetKey(starlark.String("stage_data"), starlark.NewDict(0))
	}

	val, _, err := runner.CallEntryPoint(ctx, pkg, cfg.OnAdvance.EntryPoint,
		starlark.Tuple{ctxDict}, nil, nil)
	if err != nil {
		log.Printf("[workflow-hooks] on_advance hook error: %v", err)
		return nil
	}

	return parseOnAdvanceResult(val)
}

// parseOnAdvanceResult extracts enriched data or error from the Starlark return value.
func parseOnAdvanceResult(val starlark.Value) *OnAdvanceResult {
	if val == nil || val == starlark.None {
		return nil
	}

	d, ok := val.(*starlark.Dict)
	if !ok {
		return nil
	}

	result := &OnAdvanceResult{}

	// Check for error
	if errVal, found, _ := d.Get(starlark.String("error")); found {
		if s, ok := errVal.(starlark.String); ok {
			result.Error = string(s)
			return result
		}
	}

	// Check for enriched stage_data
	if sdVal, found, _ := d.Get(starlark.String("stage_data")); found {
		if sd, ok := sdVal.(*starlark.Dict); ok {
			goMap := sandbox.DictToMap(sd)
			if data, err := json.Marshal(goMap); err == nil {
				result.EnrichedData = data
				return result
			}
		}
	}

	return nil
}