// OPA-backed policy decider for the signing-service.
//
// Loads a Rego policy file (default: policies/sinks/signing_service.rego)
// at startup, compiles it once with the OPA Go SDK, and evaluates it per
// request against the input the HTTP handler assembles from the two
// identity lookups + the requested key_ref.
package policy

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"

	"basicstack.de/proofedmesh/services/signing-service/internal/registry"
	"github.com/open-policy-agent/opa/rego"
)

// Input is the structured input the signing-service HTTP handler builds
// for every /api/v1/sign/dsse call. It is passed verbatim to Rego under
// the top-level `input` document.
type Input struct {
	PocketIDClient string          `json:"pocket_id_client"`
	K8s            InputK8s        `json:"k8s"`
	KeyRef         string          `json:"key_ref"`
	Registry       InputRegistry   `json:"registry"`
}

// InputK8s is the authenticated Kubernetes identity (from TokenReview).
type InputK8s struct {
	Namespace      string `json:"namespace"`
	ServiceAccount string `json:"service_account"`
}

// InputRegistry carries both registry lookups (by pocket id + by k8s).
// Each may be nil if the lookup returned ErrNotFound; the policy treats
// a nil row as a hard identity failure.
type InputRegistry struct {
	ByPocketID *InputRegistryRow `json:"by_pocket_id"`
	ByK8s      *InputRegistryRow `json:"by_k8s"`
}

// InputRegistryRow is a stripped-down `agents` row.
type InputRegistryRow struct {
	AgentID           string   `json:"agent_id"`
	PocketIDClient    string   `json:"pocket_id_client"`
	K8sNamespace      string   `json:"k8s_namespace"`
	K8sServiceAccount string   `json:"k8s_service_account"`
	AllowedActions    []string `json:"allowed_actions"`
}

// FromAgent builds a non-nil InputRegistryRow from a registry.Agent.
func FromAgent(a registry.Agent) *InputRegistryRow {
	return &InputRegistryRow{
		AgentID:           a.AgentID,
		PocketIDClient:    a.PocketIDClient,
		K8sNamespace:      a.K8sNamespace,
		K8sServiceAccount: a.K8sServiceAccount,
		AllowedActions:    a.AllowedActions,
	}
}

// StructuredDecider is a lower-level decider interface that policies that
// need the full Input (not just subject/keyRef) implement. The signing
// handler uses this when a StructuredDecider is wired; otherwise it falls
// back to the subject/keyRef Decider.
type StructuredDecider interface {
	Decide(ctx context.Context, in Input) (Decision, error)
}

// OPAEvaluator is a StructuredDecider backed by an in-process Rego engine.
type OPAEvaluator struct {
	allow     rego.PreparedEvalQuery
	denyReas  rego.PreparedEvalQuery
	modulePkg string
	src       string
}

// NewOPAEvaluator loads and compiles the Rego policy at path.
// The policy package MUST be `proofedmesh.sinks.signing_service` (the
// default policy file ships with that package name).
func NewOPAEvaluator(ctx context.Context, path string) (*OPAEvaluator, error) {
	src, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("policy: read rego: %w", err)
	}
	return newOPAEvaluatorFromSource(ctx, path, string(src))
}

func newOPAEvaluatorFromSource(ctx context.Context, name, src string) (*OPAEvaluator, error) {
	pkg := "proofedmesh.sinks.signing_service"
	allowQ, err := rego.New(
		rego.Query("data."+pkg+".allow"),
		rego.Module(name, src),
	).PrepareForEval(ctx)
	if err != nil {
		return nil, fmt.Errorf("policy: prepare allow: %w", err)
	}
	denyQ, err := rego.New(
		rego.Query("data."+pkg+".deny_reason"),
		rego.Module(name, src),
	).PrepareForEval(ctx)
	if err != nil {
		return nil, fmt.Errorf("policy: prepare deny_reason: %w", err)
	}
	return &OPAEvaluator{allow: allowQ, denyReas: denyQ, modulePkg: pkg, src: src}, nil
}

// Decide evaluates the Rego policy against the given input.
func (o *OPAEvaluator) Decide(ctx context.Context, in Input) (Decision, error) {
	// Rego evaluates `input` as the parsed JSON-like document. Marshal
	// through json tags to get deterministic snake_case keys.
	raw, err := toRegoInput(in)
	if err != nil {
		return Decision{}, err
	}
	rs, err := o.allow.Eval(ctx, rego.EvalInput(raw))
	if err != nil {
		return Decision{}, fmt.Errorf("policy: eval allow: %w", err)
	}
	allowed := false
	if len(rs) > 0 && len(rs[0].Expressions) > 0 {
		if v, ok := rs[0].Expressions[0].Value.(bool); ok {
			allowed = v
		}
	}
	if allowed {
		return Decision{Allow: true}, nil
	}
	// Collect deny_reason set for a usable 403 message.
	drs, err := o.denyReas.Eval(ctx, rego.EvalInput(raw))
	if err != nil {
		return Decision{}, fmt.Errorf("policy: eval deny_reason: %w", err)
	}
	reason := "signing not permitted by policy"
	if len(drs) > 0 && len(drs[0].Expressions) > 0 {
		if arr, ok := drs[0].Expressions[0].Value.([]any); ok && len(arr) > 0 {
			reason = joinReasons(arr)
		}
	}
	return Decision{Allow: false, Reason: reason}, nil
}

// StructuredAdapter wraps a StructuredDecider so it also satisfies the
// legacy subject+keyRef Decider contract (used by the HTTP layer's
// fallback path and existing tests).
type StructuredAdapter struct {
	Decider StructuredDecider
	Build   func(subject, keyRef string) Input
}

func (a StructuredAdapter) Allow(ctx context.Context, subject, keyRef string) (Decision, error) {
	if a.Build == nil {
		return Decision{}, errors.New("policy.StructuredAdapter: Build func is nil")
	}
	return a.Decider.Decide(ctx, a.Build(subject, keyRef))
}

func toRegoInput(in Input) (map[string]any, error) {
	// Round-trip through JSON to collapse to anonymous maps/slices Rego
	// handles directly, honouring our json tags (including nil-ness of
	// *InputRegistryRow).
	b, err := json.Marshal(in)
	if err != nil {
		return nil, fmt.Errorf("policy: marshal input: %w", err)
	}
	var out map[string]any
	if err := json.Unmarshal(b, &out); err != nil {
		return nil, fmt.Errorf("policy: unmarshal input: %w", err)
	}
	return out, nil
}

func joinReasons(arr []any) string {
	// Keep this tiny; policy reasons are small sets of short strings.
	out := ""
	for i, v := range arr {
		s, ok := v.(string)
		if !ok {
			continue
		}
		if i > 0 {
			out += "; "
		}
		out += s
	}
	if out == "" {
		return "signing not permitted by policy"
	}
	return out
}