package responseschema

import (
	"crypto/sha256"
	"embed"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"sort"
	"strings"
)

// Key identifies one structured response schema used by Audita.
type Key string

const (
	CorrectionSetKey          Key = "correction_set"
	ValidatorDecisionSetKey   Key = "validator_decision_set"
	correctionSetSchemaID         = "audita.correction_set"
	validatorDecisionSchemaID     = "audita.validator_decision_set"
	schemaVersionV1               = "v1"
)

//go:embed assets/*.json
var schemaAssets embed.FS

// Schema describes one registered structured response schema.
type Schema struct {
	ID         string          `json:"id"`
	Version    string          `json:"version"`
	Name       string          `json:"name"`
	JSONSchema json.RawMessage `json:"json_schema"`
	SHA256     string          `json:"sha256"`
}

func (s Schema) DiagnosticsMap() map[string]any {
	return map[string]any{
		"id":      s.ID,
		"version": s.Version,
		"name":    s.Name,
		"sha256":  s.SHA256,
	}
}

var registry = map[Key]Schema{
	CorrectionSetKey: mustBuildSchemaFromAsset(
		correctionSetSchemaID,
		schemaVersionV1,
		"audita_correction_set_v1",
		"assets/correction_set.v1.json",
	),
	ValidatorDecisionSetKey: mustBuildSchemaFromAsset(
		validatorDecisionSchemaID,
		schemaVersionV1,
		"audita_validator_decision_set_v1",
		"assets/validator_decision_set.v1.json",
	),
}

func Registered() []Schema {
	keys := make([]string, 0, len(registry))
	for key := range registry {
		keys = append(keys, string(key))
	}
	sort.Strings(keys)

	out := make([]Schema, 0, len(keys))
	for _, key := range keys {
		out = append(out, cloneSchema(registry[Key(key)]))
	}
	return out
}

// Lookup returns a copy of the registered schema for the provided key.
func Lookup(key Key) (Schema, bool) {
	schema, ok := registry[key]
	if !ok {
		return Schema{}, false
	}
	return cloneSchema(schema), true
}

// MustLookup returns a copy of the registered schema and panics when missing.
func MustLookup(key Key) Schema {
	schema, ok := Lookup(key)
	if !ok {
		panic(fmt.Sprintf("unknown structured response schema key %q", key))
	}
	return schema
}

func cloneSchema(in Schema) Schema {
	out := in
	if in.JSONSchema != nil {
		out.JSONSchema = append(json.RawMessage(nil), in.JSONSchema...)
	}
	return out
}

func mustBuildSchemaFromAsset(id string, version string, name string, assetPath string) Schema {
	rawSchema, err := schemaAssets.ReadFile(assetPath)
	if err != nil {
		panic(fmt.Sprintf("read response schema asset %q: %v", assetPath, err))
	}
	return mustBuildSchema(id, version, name, rawSchema)
}

func mustBuildSchema(id string, version string, name string, rawSchema []byte) Schema {
	id = strings.TrimSpace(id)
	version = strings.TrimSpace(version)
	name = strings.TrimSpace(name)
	if id == "" {
		panic("schema id must not be empty")
	}
	if version == "" {
		panic("schema version must not be empty")
	}
	if name == "" {
		panic("schema name must not be empty")
	}
	if !json.Valid(rawSchema) {
		panic(fmt.Sprintf("schema %s:%s is not valid JSON", id, version))
	}

	hash := sha256.Sum256(rawSchema)
	return Schema{
		ID:         id,
		Version:    version,
		Name:       name,
		JSONSchema: append(json.RawMessage(nil), rawSchema...),
		SHA256:     hex.EncodeToString(hash[:]),
	}
}