package prompts

import (
	"crypto/sha256"
	"embed"
	"encoding/hex"
	"fmt"
	"path"
	"sort"
	"strings"
	"text/template"
)

//go:embed assets/**
var embeddedAssets embed.FS

const (
	SourceBuiltin = "builtin"
	VersionV1     = "v1"
)

const (
	PromptIDModuleGlossaryProposal   = "modules.glossary.proposal"
	PromptIDModuleHomophonesProposal = "modules.homophones.proposal"
	PromptIDModuleSpokenWordProposal = "modules.spoken_word.proposal"
	PromptIDModuleGrammarProposal    = "modules.grammar.proposal"

	PromptIDValidatorSpokenFormPlausibility = "validators.spoken_form_plausibility"
	PromptIDValidatorMeaningReversalReview  = "validators.meaning_reversal_review"
	PromptIDValidatorEditorialReview        = "validators.editorial_review"
	PromptIDValidatorGrammarReview          = "validators.grammar_review"
	PromptIDValidatorSpokenWordReview       = "validators.spoken_word_review"
)

type Metadata struct {
	PromptID      string `json:"prompt_id"`
	PromptVersion string `json:"prompt_version"`
	PromptSource  string `json:"prompt_source"`
	EmbeddedPath  string `json:"embedded_path"`
	SHA256        string `json:"sha256"`
}

type definition struct {
	id          string
	version     string
	embeddedDir string
	systemPath  string
	userPath    string
}

type compiledPrompt struct {
	def        definition
	systemTmpl *template.Template
	userTmpl   *template.Template
	metadata   Metadata
}

var promptRegistry map[string]compiledPrompt
var sharedHardening string

func init() {
	var err error
	sharedHardening, err = readAsset("assets/shared/prompt_hardening.md")
	if err != nil {
		panic(err)
	}

	defs := []definition{
		{id: PromptIDModuleGlossaryProposal, version: VersionV1, embeddedDir: "assets/modules/glossary", systemPath: "assets/modules/glossary/system.md", userPath: "assets/modules/glossary/user.md"},
		{id: PromptIDModuleHomophonesProposal, version: VersionV1, embeddedDir: "assets/modules/homophones", systemPath: "assets/modules/homophones/system.md", userPath: "assets/modules/homophones/user.md"},
		{id: PromptIDModuleSpokenWordProposal, version: VersionV1, embeddedDir: "assets/modules/spoken_word", systemPath: "assets/modules/spoken_word/system.md", userPath: "assets/modules/spoken_word/user.md"},
		{id: PromptIDModuleGrammarProposal, version: VersionV1, embeddedDir: "assets/modules/grammar", systemPath: "assets/modules/grammar/system.md", userPath: "assets/modules/grammar/user.md"},
		{id: PromptIDValidatorSpokenFormPlausibility, version: VersionV1, embeddedDir: "assets/validators/spoken_form_plausibility", systemPath: "assets/validators/spoken_form_plausibility/system.md", userPath: "assets/validators/spoken_form_plausibility/user.md"},
		{id: PromptIDValidatorMeaningReversalReview, version: VersionV1, embeddedDir: "assets/validators/meaning_reversal_review", systemPath: "assets/validators/meaning_reversal_review/system.md", userPath: "assets/validators/meaning_reversal_review/user.md"},
		{id: PromptIDValidatorEditorialReview, version: VersionV1, embeddedDir: "assets/validators/editorial_review", systemPath: "assets/validators/editorial_review/system.md", userPath: "assets/validators/editorial_review/user.md"},
		{id: PromptIDValidatorGrammarReview, version: VersionV1, embeddedDir: "assets/validators/grammar_review", systemPath: "assets/validators/grammar_review/system.md", userPath: "assets/validators/grammar_review/user.md"},
		{id: PromptIDValidatorSpokenWordReview, version: VersionV1, embeddedDir: "assets/validators/spoken_word_review", systemPath: "assets/validators/spoken_word_review/system.md", userPath: "assets/validators/spoken_word_review/user.md"},
	}

	promptRegistry = make(map[string]compiledPrompt, len(defs))
	for _, def := range defs {
		cp, cErr := compilePrompt(def)
		if cErr != nil {
			panic(cErr)
		}
		promptRegistry[def.id] = cp
	}
}

func readAsset(assetPath string) (string, error) {
	b, err := embeddedAssets.ReadFile(assetPath)
	if err != nil {
		return "", fmt.Errorf("read embedded prompt asset %q: %w", assetPath, err)
	}
	return string(b), nil
}

func compilePrompt(def definition) (compiledPrompt, error) {
	systemSource, err := readAsset(def.systemPath)
	if err != nil {
		return compiledPrompt{}, err
	}
	userSource, err := readAsset(def.userPath)
	if err != nil {
		return compiledPrompt{}, err
	}

	funcs := template.FuncMap{
		"hardening": func() string { return sharedHardening },
	}
	systemTmpl, err := template.New(path.Base(def.systemPath)).Option("missingkey=error").Funcs(funcs).Parse(systemSource)
	if err != nil {
		return compiledPrompt{}, fmt.Errorf("parse embedded system prompt %q: %w", def.systemPath, err)
	}
	userTmpl, err := template.New(path.Base(def.userPath)).Option("missingkey=error").Funcs(funcs).Parse(userSource)
	if err != nil {
		return compiledPrompt{}, fmt.Errorf("parse embedded user prompt %q: %w", def.userPath, err)
	}

	hashInput := systemSource + "\n\n" + userSource
	sum := sha256.Sum256([]byte(hashInput))
	meta := Metadata{
		PromptID:      def.id,
		PromptVersion: def.version,
		PromptSource:  SourceBuiltin,
		EmbeddedPath:  def.embeddedDir,
		SHA256:        hex.EncodeToString(sum[:]),
	}

	return compiledPrompt{def: def, systemTmpl: systemTmpl, userTmpl: userTmpl, metadata: meta}, nil
}

func LookupMetadata(promptID string) (Metadata, bool) {
	cp, ok := promptRegistry[strings.TrimSpace(promptID)]
	if !ok {
		return Metadata{}, false
	}
	return cp.metadata, true
}

func MustLookupMetadata(promptID string) Metadata {
	m, ok := LookupMetadata(promptID)
	if !ok {
		panic(fmt.Sprintf("unknown prompt id %q", promptID))
	}
	return m
}

func RegisteredMetadata() []Metadata {
	ids := make([]string, 0, len(promptRegistry))
	for id := range promptRegistry {
		ids = append(ids, id)
	}
	sort.Strings(ids)
	out := make([]Metadata, 0, len(ids))
	for _, id := range ids {
		out = append(out, promptRegistry[id].metadata)
	}
	return out
}

func HardeningText() string {
	return sharedHardening
}