package main

import (
	"fmt"
	"slices"
	"strings"
)

const header = "// Code generated by lexgen from lexicons/. DO NOT EDIT.\n\n"

const helpers = `
type Record interface {
	NSID() string
}

func EncodeRecord(r Record) (map[string]any, error) {
	b, err := json.Marshal(r)
	if err != nil {
		return nil, err
	}
	return atdata.UnmarshalJSON(b)
}

func decodeRecord(value map[string]any, nsid string, out any) error {
	if got, _ := value["$type"].(string); got != nsid {
		return fmt.Errorf("%%w: expected $type %%s, got %%q", %[1]s, nsid, got)
	}
	b, err := json.Marshal(b64.NormalizeJSON(value))
	if err != nil {
		return fmt.Errorf("%%w: %%v", %[1]s, err)
	}
	if err := json.Unmarshal(b, out); err != nil {
		return fmt.Errorf("%%w: %%v", %[1]s, err)
	}
	if v, ok := out.(interface{ Validate() error }); ok {
		return v.Validate()
	}
	return nil
}

func marshalOpen(v any, extra, absent map[string]json.RawMessage) ([]byte, error) {
	b, err := json.Marshal(v)
	if err != nil || len(extra)+len(absent) == 0 {
		return b, err
	}
	fields := map[string]json.RawMessage{}
	if err := json.Unmarshal(b, &fields); err != nil {
		return nil, err
	}
	for k, raw := range absent {
		if bytes.Equal(fields[k], raw) {
			delete(fields, k)
		}
	}
	for k, raw := range extra {
		if _, ok := fields[k]; !ok {
			fields[k] = raw
		}
	}
	return json.Marshal(fields)
}

func unmarshalOpen(b []byte, v any) (extra, absent map[string]json.RawMessage, err error) {
	if err := json.Unmarshal(b, v); err != nil {
		return nil, nil, err
	}
	var in map[string]json.RawMessage
	if err := json.Unmarshal(b, &in); err != nil {
		return nil, nil, err
	}
	encoded, err := json.Marshal(v)
	if err != nil {
		return nil, nil, err
	}
	var kept map[string]json.RawMessage
	if err := json.Unmarshal(encoded, &kept); err != nil {
		return nil, nil, err
	}
	return missing(in, kept), missing(kept, in), nil
}

func missing(from, in map[string]json.RawMessage) map[string]json.RawMessage {
	var out map[string]json.RawMessage
	for k, raw := range from {
		if _, ok := in[k]; ok {
			continue
		}
		if out == nil {
			out = map[string]json.RawMessage{}
		}
		out[k] = raw
	}
	return out
}
`

func (g *gen) emit(pkg, errVar string) []byte {
	var b strings.Builder
	b.WriteString(header)
	fmt.Fprintf(&b, "package %s\n\n", pkg)
	b.WriteString("import (\n\t\"bytes\"\n\t\"encoding/json\"\n\t\"fmt\"\n\n\t\"github.com/bluesky-social/indigo/atproto/atdata\"\n\n\t\"github.com/johnandersen777/socialweb-computer-kcp/internal/kit/b64\"\n)\n")
	fmt.Fprintf(&b, helpers, errVar)
	types := slices.Clone(g.order)
	slices.SortFunc(types, func(a, b *goType) int { return strings.Compare(a.name, b.name) })
	for _, t := range types {
		switch t.kind {
		case kindRecord, kindObject:
			g.emitObject(&b, t)
		case kindArray:
			fmt.Fprintf(&b, "\ntype %s []%s\n", t.name, t.elem)
		case kindUnion:
			emitUnion(&b, t)
		}
	}
	return []byte(b.String())
}

func (g *gen) emitObject(b *strings.Builder, t *goType) {
	record := t.kind == kindRecord
	typed := record || g.variants[t.ref]
	fmt.Fprintf(b, "\ntype %s struct {\n", t.name)
	switch {
	case record:
		b.WriteString("\tType string `json:\"$type\"`\n")
	case typed:
		b.WriteString("\tType string `json:\"$type,omitempty\"`\n")
	}
	for _, f := range t.fields {
		tag := f.json
		if f.omitempty {
			tag += ",omitempty"
		}
		fmt.Fprintf(b, "\t%s %s `json:%q`\n", f.name, f.goType, tag)
	}
	b.WriteString("\tExtra map[string]json.RawMessage `json:\"-\"`\n\tabsent map[string]json.RawMessage\n}\n")
	if record {
		nsid := typeValue(t.ref)
		fmt.Fprintf(b, "\nfunc (%s) NSID() string {\n\treturn %q\n}\n", t.name, nsid)
		fmt.Fprintf(b, "\nfunc Parse%[1]s(value map[string]any) (*%[1]s, error) {\n\tvar out %[1]s\n\tif err := decodeRecord(value, %[2]q, &out); err != nil {\n\t\treturn nil, err\n\t}\n\treturn &out, nil\n}\n", t.name, nsid)
		fmt.Fprintf(b, "\nfunc (v %[1]s) MarshalJSON() ([]byte, error) {\n\ttype plain %[1]s\n\tif v.Type == \"\" {\n\t\tv.Type = %[2]q\n\t}\n\treturn marshalOpen(plain(v), v.Extra, v.absent)\n}\n", t.name, nsid)
	} else {
		fmt.Fprintf(b, "\nfunc (v %[1]s) MarshalJSON() ([]byte, error) {\n\ttype plain %[1]s\n\treturn marshalOpen(plain(v), v.Extra, v.absent)\n}\n", t.name)
	}
	fmt.Fprintf(b, "\nfunc (v *%[1]s) UnmarshalJSON(b []byte) error {\n\ttype plain %[1]s\n\tvar p plain\n\textra, absent, err := unmarshalOpen(b, &p)\n\tif err != nil {\n\t\treturn err\n\t}\n\t*v = %[1]s(p)\n\tv.Extra = extra\n\tv.absent = absent\n\treturn nil\n}\n", t.name)
}

func emitUnion(b *strings.Builder, t *goType) {
	fmt.Fprintf(b, "\ntype %s struct {\n", t.name)
	for _, v := range t.variants {
		fmt.Fprintf(b, "\t%[1]s *%[1]s\n", v.goType)
	}
	if t.open {
		b.WriteString("\tUnknown json.RawMessage\n")
	}
	b.WriteString("}\n")
	fmt.Fprintf(b, "\nfunc (u %s) MarshalJSON() ([]byte, error) {\n\tswitch {\n", t.name)
	for _, v := range t.variants {
		fmt.Fprintf(b, "\tcase u.%[1]s != nil:\n\t\tv := *u.%[1]s\n\t\tif v.Type == \"\" {\n\t\t\tv.Type = %[2]q\n\t\t}\n\t\treturn json.Marshal(v)\n", v.goType, v.typeValue)
	}
	if t.open {
		b.WriteString("\tcase u.Unknown != nil:\n\t\treturn u.Unknown, nil\n")
	}
	fmt.Fprintf(b, "\t}\n\treturn nil, fmt.Errorf(\"%s: no variant is set\")\n}\n", t.name)
	fmt.Fprintf(b, "\nfunc (u *%[1]s) UnmarshalJSON(b []byte) error {\n\tvar head struct {\n\t\tType string `json:\"$type\"`\n\t}\n\tif err := json.Unmarshal(b, &head); err != nil {\n\t\treturn err\n\t}\n\t*u = %[1]s{}\n\tswitch head.Type {\n", t.name)
	for _, v := range t.variants {
		fmt.Fprintf(b, "\tcase %[2]q:\n\t\tu.%[1]s = new(%[1]s)\n\t\treturn json.Unmarshal(b, u.%[1]s)\n", v.goType, v.typeValue)
	}
	if t.open {
		b.WriteString("\tdefault:\n\t\tu.Unknown = append(json.RawMessage(nil), b...)\n\t\treturn nil\n\t}\n}\n")
	} else {
		fmt.Fprintf(b, "\tdefault:\n\t\treturn fmt.Errorf(\"%s: unknown $type %%q\", head.Type)\n\t}\n}\n", t.name)
	}
}