package main

import (
	"fmt"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
)

// member is one wire member of a decoded container (reply or variable struct
// view). Static is the offset from the container's byte 0, or -1 when it
// follows a variable-length list and must be recomputed from earlier members.
type member struct {
	XName, Method, Kind, GoType string
	Elem                        string // annotated element Go type for list views
	Type                        *Type
	Size, Align, Static         int
	Count                       *Node
	ListKind                    string // bytes, scalar, struct, iter
}

// container is the lowered layout shared by replies and variable struct views.
type container struct {
	Name, Recv string
	Members    []member
	Fixed      int // minimum byte length (static prefix)
	Dynamic    int // index of the first variable list, -1 when none
	Reply      bool
	Module     *Module
}

const replyLifetime = "It borrows the reply bytes: valid only until the Handler callback returns, because the read ring is reused. Copy what you keep."

func isByteElem(xml string) bool {
	return xml == "void" || xml == "BYTE" || xml == "CARD8" || xml == "char"
}

func annotatedType(m *Module, n *Node, typ *Type) (string, error) {
	goType := qualified(m, typ.Module, typ.GoType)
	for _, attr := range []string{"enum", "mask"} {
		if name := n.Attr[attr]; name != "" {
			e, owner := findEnum(m, name)
			if e == nil {
				return "", n.fail("unknown %s %s", attr, name)
			}
			goType = qualified(m, owner, e.Name)
		}
	}
	return goType, nil
}

// elemName names the per-package list view for an element type.
func elemName(goType string) string {
	if i := strings.LastIndex(goType, "."); i >= 0 {
		goType = goType[i+1:]
	}
	return strings.ToUpper(goType[:1]) + goType[1:]
}

// layout selects where a container's XML members land on the wire.
type layout int

const (
	layoutStruct     layout = iota // members from byte 0
	layoutReply                    // byte 1, then byte 8 (after seq and length)
	layoutEvent                    // byte 1, then byte 4 (after seq)
	layoutEventNoSeq               // contiguous from byte 1
	layoutXGE                      // from byte 10 (after extension, seq, length, evtype)
)

// headerJump returns the offset following the header data byte, or 0 when the
// layout has no data byte.
func (l layout) headerJump() int {
	switch l {
	case layoutReply:
		return 8
	case layoutEvent:
		return 4
	}
	return 0
}

// buildContainer lowers XML members. Replies and events start at byte 1 (the
// first XML member occupies the header data byte) and continue after the
// sequence (and, for replies, length) header.
func buildContainer(m *Module, name, recv string, nodes []*Node, lay layout, reserved string) (*container, error) {
	reply := lay == layoutReply
	c := &container{Name: name, Recv: recv, Reply: reply, Dynamic: -1, Module: m}
	off := 0
	switch lay {
	case layoutReply, layoutEvent, layoutEventNoSeq:
		off = 1
	case layoutXGE:
		off = 10
	}
	what := map[layout]string{layoutStruct: "struct", layoutReply: "reply", layoutEvent: "event", layoutEventNoSeq: "event", layoutXGE: "event"}[lay]
	first := true
	methods := map[string]bool{}
	for _, n := range nodes {
		if n.Kind == "doc" {
			continue
		}
		mem := member{XName: n.Attr["name"], Kind: n.Kind, Static: off}
		if c.Dynamic >= 0 {
			mem.Static = -1
		}
		switch n.Kind {
		case "pad":
			if a := n.Attr["align"]; a != "" {
				mem.Align, _ = strconv.Atoi(a)
				if c.Dynamic < 0 {
					mem.Size = (-off) & (mem.Align - 1)
				}
			} else {
				mem.Size, _ = strconv.Atoi(n.Attr["bytes"])
			}
		case "field", "list":
			typ, err := lookup(m, n.Attr["type"])
			if err != nil {
				return nil, n.fail("%v", err)
			}
			mem.Type = typ
			mem.Method = goName(mem.XName)
			if contains(reserved, mem.Method) {
				mem.Method += "_"
			}
			if methods[mem.Method] {
				return nil, n.fail("duplicate getter %s", mem.Method)
			}
			methods[mem.Method] = true
			elem, err := annotatedType(m, n, typ)
			if err != nil {
				return nil, err
			}
			mem.Elem = elem
			if n.Kind == "field" {
				if typ.Variable {
					return nil, n.fail("variable-size field outside a list is unsupported")
				}
				mem.GoType = elem
				mem.Size = typ.Size
				break
			}
			if len(n.Children) != 1 {
				return nil, n.fail("reply-side list requires a length expression")
			}
			mem.Count = n.Children[0]
			switch {
			case isByteElem(n.Attr["type"]):
				mem.ListKind, mem.GoType = "bytes", "[]byte"
			case typ.Variable:
				if typ.Size == 0 {
					return nil, n.fail("variable struct %s has an empty fixed prefix", typ.Name)
				}
				mem.ListKind, mem.GoType = "iter", qualified(m, typ.Module, typ.GoType+"Iter")
			case isScalar(typ):
				mem.ListKind, mem.GoType = "scalar", elemName(elem)+"List"
			default:
				mem.ListKind, mem.GoType = "struct", elemName(elem)+"List"
			}
			constant := mem.Count.Kind == "value" && !typ.Variable
			if constant {
				count, _ := strconv.Atoi(mem.Count.Text)
				mem.Size = count * typ.Size
			} else if c.Dynamic < 0 {
				c.Dynamic = len(c.Members)
			}
		default:
			return nil, n.fail("unsupported %s member %s", what, n.Kind)
		}
		if jump := lay.headerJump(); jump > 0 && first {
			if mem.Size != 1 {
				return nil, n.fail("first %s member must occupy the header data byte", what)
			}
			off = jump
			first = false
			c.Members = append(c.Members, mem)
			continue
		}
		first = false
		if c.Dynamic < 0 {
			off += mem.Size
		}
		c.Members = append(c.Members, mem)
	}
	if c.Dynamic >= 0 {
		c.Fixed = c.Members[c.Dynamic].Static
	} else {
		c.Fixed = off
	}
	return c, nil
}

// ref renders a field reference as a getter call on the receiver. Replies may
// refer to the implicit header length word.
func (c *container) ref(s string) string {
	for _, mem := range c.Members {
		if mem.Kind == "field" && mem.XName == s {
			return "int(" + c.Recv + "." + mem.Method + "())"
		}
	}
	if c.Reply && s == "length" {
		return "int(binary.LittleEndian.Uint32(" + c.Recv + ".raw[4:]))"
	}
	return "<unresolved " + s + ">"
}

func (c *container) count(mem member) (string, error) {
	s, err := renderExpr(mem.Count, c.ref, c.Module)
	if err != nil {
		return "", err
	}
	if strings.Contains(s, "<unresolved") {
		return "", mem.Count.fail("list length references an unknown field")
	}
	return s, nil
}

// advance emits code moving off past member i. With fail != "" every step is
// bounds-checked against raw.
func (c *container) advance(b *strings.Builder, i int, fail string) error {
	mem := c.Members[i]
	raw := c.Recv + ".raw"
	switch mem.Kind {
	case "pad":
		if mem.Align > 0 {
			fmt.Fprintf(b, "off=(off+%d)&^%d;", mem.Align-1, mem.Align-1)
		} else {
			fmt.Fprintf(b, "off+=%d;", mem.Size)
		}
	case "field":
		fmt.Fprintf(b, "off+=%d;", mem.Size)
	case "list":
		n, err := c.count(mem)
		if err != nil {
			return err
		}
		if mem.ListKind == "iter" {
			fmt.Fprintf(b, "{ n:=%s;", n)
			if fail != "" {
				fmt.Fprintf(b, "if off>len(%s) || n<0 || n>len(%s) { %s };", raw, raw, fail)
			}
			fmt.Fprintf(b, "it:=%s(%s[off:],n); for it.Next() { off+=it.Value().Size() };", qualified(c.Module, mem.Type.Module, "New"+mem.Type.GoType+"Iter"), raw)
			if fail != "" {
				fmt.Fprintf(b, "if it.Err() { %s };", fail)
			}
			b.WriteString("};")
			return nil
		}
		fmt.Fprintf(b, "{ n:=%s;", n)
		if fail != "" {
			fmt.Fprintf(b, "if n<0 || n>len(%s) { %s };", raw, fail)
		}
		fmt.Fprintf(b, "off+=n*%d };", mem.Type.Size)
		if fail != "" {
			fmt.Fprintf(b, "if off>len(%s) { %s };", raw, fail)
		}
	}
	return nil
}

// offset emits statements leaving member i's offset in off.
func (c *container) offset(b *strings.Builder, i int) error {
	if c.Members[i].Static >= 0 {
		fmt.Fprintf(b, "off:=%d;", c.Members[i].Static)
		return nil
	}
	fmt.Fprintf(b, "off:=%d;", c.Fixed)
	for j := c.Dynamic; j < i; j++ {
		if err := c.advance(b, j, ""); err != nil {
			return err
		}
	}
	return nil
}

// validate emits the full-length walk used by Parse and Iter.Next.
func (c *container) validate(b *strings.Builder, fail string) error {
	if c.Dynamic < 0 {
		return nil
	}
	fmt.Fprintf(b, "off:=%d;", c.Fixed)
	for j := c.Dynamic; j < len(c.Members); j++ {
		if err := c.advance(b, j, fail); err != nil {
			return err
		}
	}
	fmt.Fprintf(b, "if off>len(%s.raw) { %s };", c.Recv, fail)
	return nil
}

// getters emits one method per field and list.
func (c *container) getters(b *strings.Builder, typ string, lists map[string]string) error {
	r := c.Recv
	for i, mem := range c.Members {
		if mem.Kind == "pad" {
			continue
		}
		fmt.Fprintf(b, "func (%s %s) %s() %s {", r, typ, mem.Method, mem.GoType)
		static := mem.Static >= 0
		at := fmt.Sprintf("%s.raw[%d:]", r, mem.Static)
		if !static {
			if err := c.offset(b, i); err != nil {
				return err
			}
			at = r + ".raw[off:]"
		}
		if mem.Kind == "field" {
			if isScalar(mem.Type) {
				fmt.Fprintf(b, "return %s };\n", getScalar(mem.Type, mem.GoType, at))
			} else {
				fmt.Fprintf(b, "return %s(%s) };\n", qualified(c.Module, mem.Type.Module, "Get"+mem.Type.GoType), at)
			}
			continue
		}
		if static {
			fmt.Fprintf(b, "off:=%d;", mem.Static)
		}
		n, err := c.count(mem)
		if err != nil {
			return err
		}
		switch mem.ListKind {
		case "iter":
			fmt.Fprintf(b, "return %s(%s.raw[off:],%s) };\n", qualified(c.Module, mem.Type.Module, "New"+mem.Type.GoType+"Iter"), r, n)
		case "bytes":
			fmt.Fprintf(b, "end:=off+%s; return %s.raw[off:end:end] };\n", n, r)
		default:
			fmt.Fprintf(b, "end:=off+(%s)*%d; return %s{b:%s.raw[off:end:end]} };\n", n, mem.Type.Size, mem.GoType, r)
			lists[mem.GoType] = listDecl(c.Module, mem.GoType, mem.Elem, mem.Type)
		}
	}
	return nil
}

func listDecl(m *Module, name, goType string, elem *Type) string {
	var b strings.Builder
	fmt.Fprintf(&b, "// %s is a borrowed view over a wire list of %s. %s\n", name, goType, "It shares the lifetime of the bytes it was taken from; copy what you keep.")
	fmt.Fprintf(&b, "type %s struct { b []byte };\n", name)
	fmt.Fprintf(&b, "func (l %s) Len() int { return len(l.b)/%d };\n", name, elem.Size)
	src := fmt.Sprintf("l.b[i*%d:]", elem.Size)
	if isScalar(elem) {
		fmt.Fprintf(&b, "func (l %s) At(i int) %s { return %s };\n", name, goType, getScalar(elem, goType, src))
	} else {
		fmt.Fprintf(&b, "func (l %s) At(i int) %s { return %s(%s) };\n", name, goType, qualified(m, elem.Module, "Get"+elem.GoType), src)
	}
	return b.String()
}

type replyIR struct {
	Req string
	C   *container
}

func resolveReplies(m *Module) ([]replyIR, []*container, error) {
	var replies []replyIR
	for _, n := range m.Root.Children {
		if n.Kind != "request" {
			continue
		}
		for _, r := range n.Children {
			if r.Kind != "reply" {
				continue
			}
			name := goName(n.Attr["name"])
			for _, s := range []string{name + "Reply", "Parse" + name + "Reply"} {
				if err := m.claim(s, r); err != nil {
					return nil, nil, err
				}
			}
			c, err := buildContainer(m, name+"Reply", "r", r.Children, layoutReply, "Raw")
			if err != nil {
				return nil, nil, err
			}
			replies = append(replies, replyIR{Req: name, C: c})
		}
	}
	var views []*container
	for _, t := range m.Types {
		if t.Kind != "struct" || !t.Variable {
			continue
		}
		for _, s := range []string{t.GoType + "View", t.GoType + "Iter", "New" + t.GoType + "Iter"} {
			if err := m.claim(s, t.Node); err != nil {
				return nil, nil, err
			}
		}
		c, err := buildContainer(m, t.GoType+"View", "v", t.Node.Children, layoutStruct, "Raw Size")
		if err != nil {
			return nil, nil, err
		}
		views = append(views, c)
	}
	return replies, views, nil
}

func replyDecl(r replyIR, lists map[string]string) (string, error) {
	c := r.C
	var b strings.Builder
	fmt.Fprintf(&b, "// %s is a zero-copy view over a %s reply; offsets count from reply byte 0.\n// %s\n", c.Name, r.Req, replyLifetime)
	fmt.Fprintf(&b, "type %s struct { raw []byte };\n", c.Name)
	fmt.Fprintf(&b, "// Parse%s validates raw (the whole reply, 32+4*length bytes) once; getters do not re-check.\n", c.Name)
	fmt.Fprintf(&b, "func Parse%s(raw []byte) (%s, bool) {", c.Name, c.Name)
	fixed := max(c.Fixed, 32)
	fmt.Fprintf(&b, "if len(raw)<%d { return %s{}, false };", fixed, c.Name)
	fmt.Fprintf(&b, "r:=%s{raw:raw};", c.Name)
	if err := c.validate(&b, "return "+c.Name+"{}, false"); err != nil {
		return "", err
	}
	b.WriteString("return r, true };\n")
	fmt.Fprintf(&b, "func (r %s) Raw() []byte { return r.raw };\n", c.Name)
	if err := c.getters(&b, c.Name, lists); err != nil {
		return "", err
	}
	return b.String(), nil
}

func viewDecl(c *container, lists map[string]string) (string, error) {
	base := strings.TrimSuffix(c.Name, "View")
	iter := base + "Iter"
	var b strings.Builder
	fmt.Fprintf(&b, "// %s is a zero-copy view over one wire-encoded %s, yielded by %s.\n// It shares the lifetime of the bytes it was taken from (for replies: until the Handler callback returns). Copy what you keep.\n", c.Name, base, iter)
	fmt.Fprintf(&b, "type %s struct { raw []byte };\n", c.Name)
	fmt.Fprintf(&b, "func (v %s) Raw() []byte { return v.raw };\n", c.Name)
	fmt.Fprintf(&b, "// Size is the element's wire size, including trailing alignment.\nfunc (v %s) Size() int { return len(v.raw) };\n", c.Name)
	// size%s returns the element's wire size at the start of p, -1 when malformed.
	fmt.Fprintf(&b, "func size%s(p []byte) int { if len(p)<%d { return -1 }; v:=%s{raw:p}; _=v;", base, c.Fixed, c.Name)
	if c.Dynamic < 0 {
		fmt.Fprintf(&b, "return %d };\n", c.Fixed)
	} else {
		if err := c.validate(&b, "return -1"); err != nil {
			return "", err
		}
		b.WriteString("return off };\n")
	}
	if err := c.getters(&b, c.Name, lists); err != nil {
		return "", err
	}
	fmt.Fprintf(&b, "// %s walks n consecutive %s elements; Err reports a malformed length.\n// It shares the lifetime of the bytes it walks.\n", iter, base)
	fmt.Fprintf(&b, "type %s struct { b []byte; n int; v %s; err bool };\n", iter, c.Name)
	fmt.Fprintf(&b, "// New%s iterates n elements encoded at the start of b.\nfunc New%s(b []byte, n int) %s { return %s{b:b, n:n} };\n", iter, iter, iter, iter)
	fmt.Fprintf(&b, "func (it *%s) Next() bool { if it.n<=0 || it.err { return false }; s:=size%s(it.b); if s<0 { it.err=true; it.n=0; return false }; it.v=%s{raw:it.b[:s:s]}; it.b=it.b[s:]; it.n--; return true };\n", iter, base, c.Name)
	fmt.Fprintf(&b, "func (it *%s) Value() %s { return it.v };\n", iter, c.Name)
	fmt.Fprintf(&b, "func (it *%s) Err() bool { return it.err };\n", iter)
	return b.String(), nil
}

// touch emits fuzz statements calling every getter and walking every list,
// recursing into variable struct views.
func touch(b *strings.Builder, c *container, expr string, depth int, all map[*Type]*container, failIter string) {
	for _, mem := range c.Members {
		switch {
		case mem.Kind == "field":
			fmt.Fprintf(b, "_ = %s.%s();", expr, mem.Method)
		case mem.Kind == "list" && mem.ListKind == "bytes":
			fmt.Fprintf(b, "_ = len(%s.%s());", expr, mem.Method)
		case mem.Kind == "list" && mem.ListKind == "iter":
			it := fmt.Sprintf("it%d", depth)
			fmt.Fprintf(b, "{ %s := %s.%s(); for %s.Next() {", it, expr, mem.Method, it)
			if child := all[mem.Type]; child != nil {
				touch(b, child, it+".Value()", depth+1, all, failIter)
			}
			fmt.Fprintf(b, "}; if %s.Err() { %s } };", it, failIter)
		case mem.Kind == "list":
			l := fmt.Sprintf("l%d", depth)
			fmt.Fprintf(b, "{ %s := %s.%s(); for i := 0; i < %s.Len(); i++ { _ = %s.At(i) } };", l, expr, mem.Method, l, l)
		}
	}
}

func emitReplies(modules []*Module) error {
	type pending struct {
		m       *Module
		replies []replyIR
		views   []*container
	}
	var all []pending
	views := map[*Type]*container{}
	for _, m := range modules {
		replies, vs, err := resolveReplies(m)
		if err != nil {
			return err
		}
		for _, v := range vs {
			for _, t := range m.Types {
				if t.GoType+"View" == v.Name {
					views[t] = v
				}
			}
		}
		all = append(all, pending{m, replies, vs})
	}
	for _, p := range all {
		m := p.m
		lists := map[string]string{}
		var decls []string
		for _, r := range p.replies {
			d, err := replyDecl(r, lists)
			if err != nil {
				return err
			}
			decls = append(decls, d)
		}
		for _, v := range p.views {
			d, err := viewDecl(v, lists)
			if err != nil {
				return err
			}
			decls = append(decls, d)
		}
		var names []string
		for name := range lists {
			if err := m.claim(name, m.Root); err != nil {
				return err
			}
			names = append(names, name)
		}
		sort.Strings(names)
		for _, name := range names {
			decls = append(decls, lists[name])
		}
		if len(decls) == 0 {
			continue
		}
		decls = append([]string{importsFor(m, strings.Join(decls, "\n"))}, decls...)
		data := struct {
			Pkg   string
			Decls []string
		}{Pkg: m.Pkg, Decls: decls}
		if err := writeTemplate(filepath.Join(m.Pkg, "replies_generated.go"), "replies", data); err != nil {
			return err
		}
		if err := emitReplyFuzz(m, p.replies, views); err != nil {
			return err
		}
	}
	return nil
}

// emitReplyFuzz writes FuzzParse<Req>Reply for every reply. Seeds are the
// zero-filled fixed prefix plus any captured goldens <Req>*.reply.hex.
func emitReplyFuzz(m *Module, replies []replyIR, views map[*Type]*container) error {
	var decls []string
	usesGolden := false
	for _, r := range replies {
		var b strings.Builder
		c := r.C
		fmt.Fprintf(&b, "func FuzzParse%s(f *testing.F) {\n", c.Name)
		fmt.Fprintf(&b, "f.Add(make([]byte, %d));\n", (max(c.Fixed, 32)+3)&^3)
		goldens, _ := filepath.Glob(filepath.Join("testdata/golden", m.Pkg, r.Req+"*.reply.hex"))
		sort.Strings(goldens)
		for _, g := range goldens {
			base := filepath.Base(g)
			if strings.TrimSuffix(base, ".reply.hex") != r.Req && !startsUpper(strings.TrimPrefix(strings.TrimSuffix(base, ".reply.hex"), r.Req)) {
				continue
			}
			usesGolden = true
			fmt.Fprintf(&b, "if raw, err := xtest.LoadHex(%q); err != nil { f.Fatal(err) } else { f.Add(raw) };\n", "../"+filepath.ToSlash(g))
		}
		fmt.Fprintf(&b, "f.Fuzz(func(t *testing.T, raw []byte) {\nr, ok := Parse%s(raw);\nif !ok {", c.Name)
		if c.Dynamic < 0 {
			fmt.Fprintf(&b, "if len(raw) >= %d { t.Fatal(\"rejected a complete fixed-size reply\") };", max(c.Fixed, 32))
		}
		b.WriteString("return };\n")
		b.WriteString("if len(r.Raw()) != len(raw) { t.Fatal(\"Raw must return the parsed bytes\") };\n")
		touch(&b, c, "r", 0, views, `t.Fatal("iterator error after successful Parse")`)
		b.WriteString("\n})\n}\n")
		decls = append(decls, b.String())
	}
	imports := "import \"testing\""
	if usesGolden {
		imports += "\nimport \"bnema.dev/xturbo/internal/xtest\""
	}
	decls = append([]string{imports}, decls...)
	data := struct {
		Pkg   string
		Decls []string
	}{Pkg: m.Pkg, Decls: decls}
	return writeTemplate(filepath.Join(m.Pkg, "replies_fuzz_test.go"), "replies", data)
}

func startsUpper(s string) bool { return s != "" && s[0] >= 'A' && s[0] <= 'Z' }