package llm

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

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

// ResponseSchemaKey identifies one structured response schema.
type ResponseSchemaKey string

const (
	DNDSpellsSchemaKey             ResponseSchemaKey = "dnd_spells"
	TestArtifactSchemaKey          ResponseSchemaKey = "test_artifact"
	TestValidatorDecisionSchemaKey ResponseSchemaKey = "test_validator_decision"

	schemaVersionV1 = "v1"
)

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

var responseSchemaRegistry = map[ResponseSchemaKey]ResponseSchema{
	DNDSpellsSchemaKey: mustLoadResponseSchema(
		DNDSpellsSchemaKey,
		"notarius.dnd.spells",
		schemaVersionV1,
		"notarius_dnd_spells_v1",
		"assets/schemas/dnd_spells.v1.json",
	),
	TestArtifactSchemaKey: mustLoadResponseSchema(
		TestArtifactSchemaKey,
		"notarius.test_artifact",
		schemaVersionV1,
		"notarius_test_artifact_v1",
		"assets/schemas/test_artifact.v1.json",
	),
	TestValidatorDecisionSchemaKey: mustLoadResponseSchema(
		TestValidatorDecisionSchemaKey,
		"notarius.test_validator_decision",
		schemaVersionV1,
		"notarius_test_validator_decision_v1",
		"assets/schemas/test_validator_decision.v1.json",
	),
}

// RegisteredResponseSchemas returns all registered response schemas sorted by key.
func RegisteredResponseSchemas() []ResponseSchema {
	keys := make([]string, 0, len(responseSchemaRegistry))
	for key := range responseSchemaRegistry {
		keys = append(keys, string(key))
	}
	sort.Strings(keys)

	out := make([]ResponseSchema, 0, len(keys))
	for _, key := range keys {
		out = append(out, cloneResponseSchema(responseSchemaRegistry[ResponseSchemaKey(key)]))
	}
	return out
}

// LookupResponseSchema returns a copy of the schema for key.
func LookupResponseSchema(key ResponseSchemaKey) (ResponseSchema, bool) {
	schema, ok := responseSchemaRegistry[key]
	if !ok {
		return ResponseSchema{}, false
	}
	return cloneResponseSchema(schema), true
}

// MustLookupResponseSchema returns a copy of the schema for key and panics when missing.
func MustLookupResponseSchema(key ResponseSchemaKey) ResponseSchema {
	schema, ok := LookupResponseSchema(key)
	if !ok {
		panic(fmt.Sprintf("unknown structured response schema key %q", key))
	}
	return schema
}

// DiagnosticsMap returns schema metadata without raw schema content.
func (s ResponseSchema) DiagnosticsMap() map[string]any {
	return map[string]any{
		"key":     s.Key,
		"id":      s.ID,
		"version": s.Version,
		"name":    s.Name,
		"sha256":  s.SHA256,
	}
}

func mustLoadResponseSchema(
	key ResponseSchemaKey,
	id string,
	version string,
	name string,
	path string,
) ResponseSchema {
	key = ResponseSchemaKey(strings.TrimSpace(string(key)))
	id = strings.TrimSpace(id)
	version = strings.TrimSpace(version)
	name = strings.TrimSpace(name)
	path = strings.TrimSpace(path)
	if key == "" {
		panic("response schema key must not be empty")
	}
	if id == "" {
		panic("response schema id must not be empty")
	}
	if version == "" {
		panic("response schema version must not be empty")
	}
	if name == "" {
		panic("response schema name must not be empty")
	}
	if path == "" {
		panic("response schema asset path must not be empty")
	}

	rawSchema, err := schemaAssets.ReadFile(path)
	if err != nil {
		panic(fmt.Sprintf("read response schema %s: %v", path, err))
	}
	if !json.Valid(rawSchema) {
		panic(fmt.Sprintf("response schema %s is not valid JSON", path))
	}

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

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