// Package gen contains the per-file Lua codegen used by protoc-gen-tarantool.
package gen

import (
	"fmt"
	"sort"
	"strconv"
	"strings"

	"google.golang.org/protobuf/compiler/protogen"
	"google.golang.org/protobuf/reflect/protoreflect"
)

const runtimeRequire = "pb"

// Mode controls how generated _encode / _decode wrappers are produced.
//
//   - ModeFull (default): emit per-message inline encode/decode functions
//     that call wire primitives directly, no descriptor dispatch. Faster,
//     more JIT-friendly, larger output.
//   - ModeRuntime: emit thin wrappers that delegate to pb.encode / pb.decode
//     against the (always-emitted) descriptor table. Slower, smaller output,
//     useful for introspection.
//
// The descriptor table is emitted in both modes so users can introspect
// schemas and so future tooling (registry, dynamic types) keeps working.
type Mode int

const (
	ModeFull Mode = iota
	ModeRuntime
)

// ParseMode converts a CLI value (full|runtime) to a Mode. Empty -> default.
func ParseMode(s string) (Mode, error) {
	switch s {
	case "", "full":
		return ModeFull, nil
	case "runtime":
		return ModeRuntime, nil
	}
	return ModeFull, fmt.Errorf("unknown mode %q (want full|runtime)", s)
}

// Config carries per-invocation generator options.
type Config struct {
	Mode Mode
	// Prefix, when non-empty, is prepended to every generated module's Lua
	// require path (and its on-disk subpath). Lets the same .proto be
	// generated under multiple namespaces in one project — e.g. for
	// side-by-side full vs runtime mode comparison in tests.
	Prefix string
}

// GenerateFile emits one `<lua_pkg>.lua` file per input `.proto`.
func GenerateFile(plug *protogen.Plugin, file *protogen.File, cfg Config) error {
	if file.Desc.Syntax() != protoreflect.Proto3 {
		return fmt.Errorf("%s: only proto3 is supported, got %s",
			file.Desc.Path(), file.Desc.Syntax())
	}

	allMsgs := flattenMessagesSkippingMapEntries(file.Messages, nil)
	allEnums := flattenEnums(file.Enums, file.Messages)

	out := plug.NewGeneratedFile(outputFilename(file.Desc, cfg.Prefix), "")
	w := &writer{GeneratedFile: out}

	emitHeader(w, file)
	imports := collectImports(file, allMsgs, cfg.Prefix)
	emitRequires(w, imports)

	w.line("local M = {}")
	w.line("")

	// 1) Enums first (no forward-ref problems).
	for _, e := range allEnums {
		emitEnum(w, file, e)
	}

	// 2) Predeclare all message descriptor tables (so cross-references resolve).
	if len(allMsgs) > 0 {
		w.line("-- Pre-declare message descriptors so cross-references resolve.")
		for _, m := range allMsgs {
			name := luaTypeName(m.Desc.FullName(), file.Desc.Package())
			w.line("M.%s_descriptor = {name = %q}", name, string(m.Desc.FullName()))
		}
		w.line("")
	}

	// 3) Fill in fields[] for each message and finalize.
	for _, m := range allMsgs {
		emitMessageFields(w, file, m, imports, cfg.Prefix)
	}

	// 3a) EmmyLua type annotations (---@class per message, ---@alias per
	// enum). Emitted between the descriptors and the wrappers so the
	// wrapper annotations a few lines down can reference these class names.
	emitEmmyTypesBlock(w, allEnums, allMsgs)

	// 4) Wrappers: _new / _encode / _decode.
	for _, m := range allMsgs {
		switch cfg.Mode {
		case ModeFull:
			emitInlineMessage(w, file, m, imports, cfg.Prefix)
		default:
			emitMessageWrappers(w, file, m)
		}
	}

	// 5) Services (mode-independent — client/server stubs delegate to the
	// per-message _encode/_decode functions emitted in step 4).
	for _, svc := range file.Services {
		emitService(w, file, svc, imports, cfg.Prefix)
	}

	w.line("return M")
	return nil
}

// ----------------------------------------------------------------------------
// writer: thin wrapper for line-oriented emission.
// ----------------------------------------------------------------------------

type writer struct {
	*protogen.GeneratedFile
}

func (w *writer) line(format string, args ...any) {
	if len(args) == 0 {
		w.P(format)
	} else {
		w.P(fmt.Sprintf(format, args...))
	}
}

// ----------------------------------------------------------------------------
// Header & requires
// ----------------------------------------------------------------------------

func emitHeader(w *writer, file *protogen.File) {
	w.line("-- Code generated by protoc-gen-tarantool. DO NOT EDIT.")
	w.line("-- source: %s", file.Desc.Path())
	w.line("-- syntax: %s", file.Desc.Syntax())
	if pkg := string(file.Desc.Package()); pkg != "" {
		w.line("-- package: %s", pkg)
	}
	w.line("")
	w.line("local pb = require(%q)", runtimeRequire)
	w.line("local wire = pb.wire")
}

// collectImports returns the deduplicated set of Lua require paths for all
// other .proto files referenced by message fields *and* service inputs/outputs
// in this file.
func collectImports(file *protogen.File, msgs []*protogen.Message, prefix string) map[string]string {
	selfPath := luaPackagePath(file.Desc, prefix)
	out := map[string]string{}
	addType := func(ext protoreflect.FileDescriptor) {
		if ext == nil {
			return
		}
		if isWellKnownTypeFile(ext) {
			return  // pb.wkt is reachable via the existing `pb` require
		}
		lp := luaPackagePath(ext, prefix)
		if lp == selfPath {
			return
		}
		out[lp] = importAlias(lp)
	}
	for _, m := range msgs {
		for _, f := range m.Fields {
			switch {
			case f.Message != nil:
				addType(f.Message.Desc.ParentFile())
			case f.Enum != nil:
				addType(f.Enum.Desc.ParentFile())
			}
		}
	}
	for _, svc := range file.Services {
		for _, meth := range svc.Methods {
			addType(meth.Input.Desc.ParentFile())
			addType(meth.Output.Desc.ParentFile())
		}
	}
	return out
}

func emitRequires(w *writer, imports map[string]string) {
	if len(imports) == 0 {
		w.line("")
		return
	}
	keys := make([]string, 0, len(imports))
	for k := range imports {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	for _, k := range keys {
		w.line("local %s = require(%q)", imports[k], k)
	}
	w.line("")
}

// ----------------------------------------------------------------------------
// Enum emission
// ----------------------------------------------------------------------------

func emitEnum(w *writer, file *protogen.File, e *protogen.Enum) {
	name := luaTypeName(e.Desc.FullName(), file.Desc.Package())
	w.line("-- Enum: %s", e.Desc.FullName())
	w.line("M.%s_descriptor = pb.enum(%q, {", name, string(e.Desc.FullName()))
	for _, v := range e.Values {
		w.line("    %s = %d,", string(v.Desc.Name()), v.Desc.Number())
	}
	w.line("})")
	// Convenience aliases the user can reach via `M.MyEnum.RED`, etc.
	w.line("M.%s = M.%s_descriptor.by_name", name, name)
	w.line("")
}

// ----------------------------------------------------------------------------
// Message emission
// ----------------------------------------------------------------------------

// emitMessageFields fills the predeclared M.<Name>_descriptor with its
// fields[] array and finalizes it (which builds field_by_id).
func emitMessageFields(w *writer, file *protogen.File, m *protogen.Message, imports map[string]string, prefix string) {
	name := luaTypeName(m.Desc.FullName(), file.Desc.Package())
	selfPath := luaPackagePath(file.Desc, prefix)

	w.line("-- Message: %s", m.Desc.FullName())
	w.line("M.%s_descriptor.fields = {", name)
	for _, f := range m.Fields {
		w.line("    %s,", renderFieldEntry(file, f, selfPath, imports, prefix))
	}
	w.line("}")
	emitOneofTable(w, name, m)
	w.line("pb.finalize_message(M.%s_descriptor)", name)
	w.line("")
}

// emitOneofTable emits `M.<Name>_descriptor.oneofs = { <name> = {...members...} }`
// when the message has any non-synthetic oneofs.
func emitOneofTable(w *writer, name string, m *protogen.Message) {
	type oneofRow struct {
		name    string
		members []string
	}
	var rows []oneofRow
	for _, oo := range m.Oneofs {
		// Skip synthetic oneofs created for proto3 explicit `optional` — those
		// have a single member that uses HasOptionalKeyword().
		if oo.Fields[0].Desc.HasOptionalKeyword() {
			continue
		}
		row := oneofRow{name: string(oo.Desc.Name())}
		for _, f := range oo.Fields {
			row.members = append(row.members, string(f.Desc.Name()))
		}
		rows = append(rows, row)
	}
	if len(rows) == 0 {
		return
	}
	w.line("M.%s_descriptor.oneofs = {", name)
	for _, row := range rows {
		quoted := make([]string, 0, len(row.members))
		for _, fn := range row.members {
			quoted = append(quoted, fmt.Sprintf("%q", fn))
		}
		w.line("    %s = {%s},", row.name, strings.Join(quoted, ", "))
	}
	w.line("}")
}

// renderFieldEntry produces the Lua table literal for a single field descriptor.
func renderFieldEntry(file *protogen.File, f *protogen.Field, selfPath string, imports map[string]string, prefix string) string {
	parts := []string{
		fmt.Sprintf("name=%q", string(f.Desc.Name())),
		fmt.Sprintf("id=%d", f.Desc.Number()),
	}

	if f.Desc.IsMap() {
		parts = append(parts, "kind='map'")
		parts = append(parts, "key="+renderMapEntry(file, f.Message.Fields[0], selfPath, imports, prefix))
		parts = append(parts, "value="+renderMapEntry(file, f.Message.Fields[1], selfPath, imports, prefix))
		return "{" + strings.Join(parts, ", ") + "}"
	}

	switch {
	case f.Message != nil:
		parts = append(parts, "kind='message'")
		parts = append(parts, "message="+typeRef(file, f.Message.Desc, selfPath, imports, "_descriptor", prefix))
	case f.Enum != nil:
		parts = append(parts, "kind='enum'")
		parts = append(parts, "enum="+typeRef(file, f.Enum.Desc, selfPath, imports, "_descriptor", prefix))
	default:
		s := scalarName(f.Desc.Kind())
		if s == "" {
			panic("unhandled scalar kind: " + f.Desc.Kind().String())
		}
		parts = append(parts, "kind='scalar'")
		parts = append(parts, "proto_type="+strconv.Quote(s))
	}

	if f.Desc.IsList() {
		parts = append(parts, "repeated=true")
		// proto3 packed default for primitives + enums is true; explicit
		// `[packed=false]` flips it. IsPacked() returns the effective value.
		if f.Message == nil && f.Desc.Kind() != protoreflect.StringKind &&
			f.Desc.Kind() != protoreflect.BytesKind {
			if f.Desc.IsPacked() {
				parts = append(parts, "packed=true")
			} else {
				parts = append(parts, "packed=false")
			}
		}
	}

	// Oneof membership. (Skip synthetic oneofs that proto3 explicit `optional`
	// expands into — those are surfaced as `optional=true` instead.)
	if f.Oneof != nil && !f.Desc.HasOptionalKeyword() {
		parts = append(parts, fmt.Sprintf("oneof=%q", string(f.Oneof.Desc.Name())))
	}

	// Proto3 explicit optional (field presence).
	if f.Desc.HasOptionalKeyword() {
		parts = append(parts, "optional=true")
	}

	return "{" + strings.Join(parts, ", ") + "}"
}

// renderMapEntry renders a sub-field descriptor for a map's key or value.
// It mirrors renderFieldEntry but always for a singular non-map value, and
// emits without the `name`/`id` (caller knows: id 1 = key, id 2 = value).
func renderMapEntry(file *protogen.File, f *protogen.Field, selfPath string, imports map[string]string, prefix string) string {
	parts := []string{}
	switch {
	case f.Message != nil:
		parts = append(parts, "kind='message'")
		parts = append(parts, "message="+typeRef(file, f.Message.Desc, selfPath, imports, "_descriptor", prefix))
	case f.Enum != nil:
		parts = append(parts, "kind='enum'")
		parts = append(parts, "enum="+typeRef(file, f.Enum.Desc, selfPath, imports, "_descriptor", prefix))
	default:
		s := scalarName(f.Desc.Kind())
		if s == "" {
			panic("unhandled map sub-field kind: " + f.Desc.Kind().String())
		}
		parts = append(parts, "kind='scalar'")
		parts = append(parts, "proto_type="+strconv.Quote(s))
	}
	return "{" + strings.Join(parts, ", ") + "}"
}

// typeRef returns a Lua expression evaluating to the descriptor of the given
// type (a Message or Enum), resolving cross-file imports as needed.
func typeRef(file *protogen.File, td protoreflect.Descriptor, selfPath string, imports map[string]string, suffix string, prefix string) string {
	parent := td.ParentFile()
	if isWellKnownTypeFile(parent) {
		return "pb.wkt." + wktTypeName(td.FullName()) + suffix
	}
	luaName := luaTypeName(td.FullName(), parent.Package())
	parentPath := luaPackagePath(parent, prefix)
	if parentPath == selfPath {
		return "M." + luaName + suffix
	}
	alias, ok := imports[parentPath]
	if !ok {
		// Should be impossible if collectImports walked all fields.
		alias = importAlias(parentPath)
	}
	return alias + "." + luaName + suffix
}

// emitMessageWrappers emits the small _new / _encode / _decode helpers plus
// has_<field> / clear_<field> for each explicit-optional field.
func emitMessageWrappers(w *writer, file *protogen.File, m *protogen.Message) {
	name := luaTypeName(m.Desc.FullName(), file.Desc.Package())
	full := emmyMessageFullName(m)

	emitEmmyWrapperAnnotations(w, name, full, wrapperNew)
	w.line("function M.%s_new(t) return t or {} end", name)
	emitEmmyWrapperAnnotations(w, name, full, wrapperEncode)
	w.line("function M.%s_encode(t) return pb.encode(M.%s_descriptor, t) end", name, name)
	emitEmmyWrapperAnnotations(w, name, full, wrapperDecode)
	w.line("function M.%s_decode(b) return pb.decode(M.%s_descriptor, b) end", name, name)
	emitEmmyWrapperAnnotations(w, name, full, wrapperDecodeLazy)
	w.line("function M.%s_decode_lazy(b) return pb.decode_lazy(M.%s_descriptor, b) end", name, name)
	emitOptionalAccessors(w, name, m, full)
	w.line("")
}

// emitOptionalAccessors writes M.<Name>_has_<field>(t) and _clear_<field>(t)
// for every field marked with proto3 explicit `optional`.
func emitOptionalAccessors(w *writer, name string, m *protogen.Message, fullName string) {
	for _, f := range m.Fields {
		if !f.Desc.HasOptionalKeyword() {
			continue
		}
		fname := string(f.Desc.Name())
		emitEmmyWrapperAnnotations(w, name, fullName, wrapperHas)
		w.line("function M.%s_has_%s(t) return t.%s ~= nil end", name, fname, fname)
		emitEmmyWrapperAnnotations(w, name, fullName, wrapperClear)
		w.line("function M.%s_clear_%s(t) t.%s = nil end", name, fname, fname)
	}
}

// ----------------------------------------------------------------------------
// Flattening helpers
// ----------------------------------------------------------------------------

// flattenMessages returns top-level + all nested messages in declaration order.
func flattenMessages(top []*protogen.Message, acc []*protogen.Message) []*protogen.Message {
	for _, m := range top {
		acc = append(acc, m)
		acc = flattenMessages(m.Messages, acc)
	}
	return acc
}

// flattenMessagesSkippingMapEntries is like flattenMessages but excludes the
// synthetic <Field>Entry messages protoc generates for `map<K,V>` fields.
// Those don't get their own Lua descriptor — map handling is inline.
func flattenMessagesSkippingMapEntries(top []*protogen.Message, acc []*protogen.Message) []*protogen.Message {
	for _, m := range top {
		if m.Desc.IsMapEntry() {
			continue
		}
		acc = append(acc, m)
		acc = flattenMessagesSkippingMapEntries(m.Messages, acc)
	}
	return acc
}

// flattenEnums returns top-level enums + all enums nested inside messages.
func flattenEnums(topEnums []*protogen.Enum, msgs []*protogen.Message) []*protogen.Enum {
	out := append([]*protogen.Enum{}, topEnums...)
	var walk func(ms []*protogen.Message)
	walk = func(ms []*protogen.Message) {
		for _, m := range ms {
			out = append(out, m.Enums...)
			walk(m.Messages)
		}
	}
	walk(msgs)
	return out
}