package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"go/format"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strings"
	"unicode"
)

var clientMethods = []string{
	"initialize",
	"thread/start",
	"thread/resume",
	"thread/read",
	"thread/compact/start",
	"turn/start",
	"turn/interrupt",
}

type schema struct {
	any                  bool
	Ref                  string             `json:"$ref"`
	Type                 json.RawMessage    `json:"type"`
	Enum                 []json.RawMessage  `json:"enum"`
	Title                string             `json:"title"`
	Properties           map[string]*schema `json:"properties"`
	Required             []string           `json:"required"`
	Items                *schema            `json:"items"`
	AdditionalProperties *schema            `json:"additionalProperties"`
	OneOf                []*schema          `json:"oneOf"`
	AnyOf                []*schema          `json:"anyOf"`
	AllOf                []*schema          `json:"allOf"`
}

func (s *schema) UnmarshalJSON(b []byte) error {
	b = bytes.TrimSpace(b)
	if bytes.Equal(b, []byte("true")) || bytes.Equal(b, []byte("false")) {
		s.any = true
		return nil
	}
	type plain schema
	return json.Unmarshal(b, (*plain)(s))
}

func (s *schema) types() []string {
	if len(s.Type) == 0 {
		return nil
	}
	if s.Type[0] == '"' {
		t := ""
		err := json.Unmarshal(s.Type, &t)
		if err != nil {
			log.Fatalf("type: %v", err)
		}
		return []string{t}
	}
	ts := []string{}
	err := json.Unmarshal(s.Type, &ts)
	if err != nil {
		log.Fatalf("type: %v", err)
	}
	return ts
}

func (s *schema) isNull() bool {
	ts := s.types()
	return len(ts) == 1 && ts[0] == "null" && s.Ref == ""
}

type bundle struct {
	Definitions map[string]json.RawMessage `json:"definitions"`
}

type gen struct {
	defs    map[string]*schema
	order   []string
	emitted map[string]bool
	kinds   map[string]string
	out     bytes.Buffer
	consts  bytes.Buffer
}

func main() {
	dir, err := os.MkdirTemp("", "codex-schema")
	if err != nil {
		log.Fatalf("mkdir temp: %v", err)
	}
	defer os.RemoveAll(dir)

	cmd := exec.Command("codex", "app-server", "generate-json-schema", "--experimental", "--out", dir)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	err = cmd.Run()
	if err != nil {
		log.Fatalf("generate schema: %v", err)
	}

	version, err := exec.Command("codex", "--version").Output()
	if err != nil {
		log.Fatalf("codex version: %v", err)
	}

	raw, err := os.ReadFile(filepath.Join(dir, "codex_app_server_protocol.schemas.json"))
	if err != nil {
		log.Fatalf("read schema: %v", err)
	}

	b := bundle{}
	err = json.Unmarshal(raw, &b)
	if err != nil {
		log.Fatalf("parse schema: %v", err)
	}

	g := &gen{defs: map[string]*schema{}, emitted: map[string]bool{}, kinds: map[string]string{}}
	g.load(b)

	src := g.generate(strings.TrimSpace(string(version)))
	formatted, err := format.Source(src)
	if err != nil {
		os.WriteFile("protocol.gen.go", src, 0o644)
		log.Fatalf("format: %v", err)
	}
	err = os.WriteFile("protocol.gen.go", formatted, 0o644)
	if err != nil {
		log.Fatalf("write: %v", err)
	}
}

func (g *gen) load(b bundle) {
	top := map[string]json.RawMessage{}
	for name, raw := range b.Definitions {
		if name == "v2" {
			continue
		}
		top[name] = raw
	}

	v2 := map[string]json.RawMessage{}
	err := json.Unmarshal(b.Definitions["v2"], &v2)
	if err != nil {
		log.Fatalf("parse v2 definitions: %v", err)
	}

	for _, set := range []map[string]json.RawMessage{top, v2} {
		for name, raw := range set {
			s := &schema{}
			err := json.Unmarshal(raw, s)
			if err != nil {
				log.Fatalf("parse %s: %v", name, err)
			}
			prev, ok := g.defs[name]
			if ok {
				p, q := *prev, *s
				p.Title, q.Title = "", ""
				a, _ := json.Marshal(p)
				c, _ := json.Marshal(q)
				if !bytes.Equal(a, c) {
					log.Fatalf("definition %s differs between the top level and v2", name)
				}
				continue
			}
			g.defs[name] = s
		}
	}
}

func refName(ref string) string {
	i := strings.LastIndex(ref, "/")
	return ref[i+1:]
}

type method struct {
	name     string
	params   string
	response string
}

func (g *gen) variants(def string) []method {
	s := g.defs[def]
	if s == nil {
		log.Fatalf("no definition %s", def)
	}
	ms := []method{}
	for _, v := range s.OneOf {
		m := v.Properties["method"]
		if m == nil || len(m.Enum) != 1 {
			log.Fatalf("%s variant without a single method", def)
		}
		name := ""
		err := json.Unmarshal(m.Enum[0], &name)
		if err != nil {
			log.Fatalf("%s method: %v", def, err)
		}
		params := ""
		if p := v.Properties["params"]; p != nil && p.Ref != "" {
			params = refName(p.Ref)
		}
		response := ""
		if strings.HasSuffix(params, "Params") {
			r := strings.TrimSuffix(params, "Params") + "Response"
			if g.defs[r] != nil {
				response = r
			}
		}
		ms = append(ms, method{name: name, params: params, response: response})
	}
	return ms
}

func (g *gen) generate(version string) []byte {
	client := []method{}
	for _, m := range g.variants("ClientRequest") {
		if slices.Contains(clientMethods, m.name) {
			client = append(client, m)
		}
	}
	if len(client) != len(clientMethods) {
		log.Fatalf("found %d of %d client methods", len(client), len(clientMethods))
	}
	notifications := g.variants("ServerNotification")
	requests := g.variants("ServerRequest")

	roots := []string{}
	for _, set := range [][]method{client, notifications, requests} {
		for _, m := range set {
			if m.params != "" {
				roots = append(roots, m.params)
			}
			if m.response != "" {
				roots = append(roots, m.response)
			}
		}
	}
	for _, r := range roots {
		g.named(r)
	}

	h := bytes.Buffer{}
	fmt.Fprintf(&h, "// Code generated by go run ./gen from the %s app-server schema. DO NOT EDIT.\n\n", version)
	h.WriteString("package codex\n\nimport \"encoding/json\"\n\n")
	fmt.Fprintf(&h, "const SchemaVersion = %q\n\n", version)

	h.WriteString("const (\n")
	seen := map[string]bool{}
	for _, set := range [][]method{client, notifications, requests} {
		for _, m := range set {
			c := "Method" + pascal(m.name)
			if seen[c] {
				continue
			}
			seen[c] = true
			fmt.Fprintf(&h, "%s = %q\n", c, m.name)
		}
	}
	h.WriteString(")\n\n")

	writeTable := func(name string, ms []method, field func(method) string) {
		fmt.Fprintf(&h, "var %s = map[string]func() any{\n", name)
		for _, m := range ms {
			t := field(m)
			if t == "" {
				fmt.Fprintf(&h, "Method%s: nil,\n", pascal(m.name))
				continue
			}
			fmt.Fprintf(&h, "Method%s: func() any { return &%s{} },\n", pascal(m.name), t)
		}
		h.WriteString("}\n\n")
	}
	writeTable("clientResults", client, func(m method) string { return m.response })
	writeTable("notificationParams", notifications, func(m method) string { return m.params })
	writeTable("serverRequestParams", requests, func(m method) string { return m.params })
	writeTable("serverRequestResponses", requests, func(m method) string { return m.response })

	h.Write(g.consts.Bytes())
	h.Write(g.out.Bytes())
	return h.Bytes()
}

func (g *gen) named(name string) string {
	s := g.defs[name]
	if s == nil {
		log.Fatalf("no definition %s", name)
	}
	if g.emitted[name] {
		return name
	}
	g.emitted[name] = true
	g.kinds[name] = g.kindOf(s)
	g.emitType(name, s)
	return name
}

func (g *gen) resolve(s *schema) *schema {
	for s.Ref != "" {
		s = g.defs[refName(s.Ref)]
	}
	return s
}

func nonNull(list []*schema) []*schema {
	out := []*schema{}
	for _, v := range list {
		if !v.isNull() {
			out = append(out, v)
		}
	}
	return out
}

func (g *gen) kindOf(s *schema) string {
	if s.any {
		return "raw"
	}
	if s.Ref != "" {
		return g.kindOf(g.resolve(s))
	}
	if len(s.AllOf) == 1 {
		return g.kindOf(s.AllOf[0])
	}
	alts := nonNull(append(append([]*schema{}, s.OneOf...), s.AnyOf...))
	if len(s.Properties) == 0 && len(alts) == 1 {
		return g.kindOf(alts[0])
	}
	if len(s.Properties) == 0 && len(alts) > 1 {
		kinds := map[string]bool{}
		for _, a := range alts {
			kinds[g.kindOf(a)] = true
		}
		if len(kinds) == 1 && (kinds["string"] || kinds["struct"]) {
			for k := range kinds {
				return k
			}
		}
		return "raw"
	}
	if len(s.Properties) > 0 {
		return "struct"
	}
	ts := []string{}
	for _, t := range s.types() {
		if t != "null" {
			ts = append(ts, t)
		}
	}
	if len(ts) != 1 {
		return "raw"
	}
	switch ts[0] {
	case "string":
		return "string"
	case "integer", "number", "boolean":
		return "scalar"
	case "array":
		return "slice"
	case "object":
		return "map"
	}
	return "raw"
}

func (g *gen) emitType(name string, s *schema) {
	switch g.kindOf(s) {
	case "struct":
		fields := g.collectFields(s)
		g.emitStruct(name, fields)
		g.emitDiscriminators(name, s)
	case "string":
		fmt.Fprintf(&g.out, "type %s string\n\n", name)
		values := g.enumValues(s)
		if len(values) > 0 {
			g.consts.WriteString("const (\n")
			for _, v := range values {
				fmt.Fprintf(&g.consts, "%s%s %s = %q\n", name, pascal(v), name, v)
			}
			g.consts.WriteString(")\n\n")
		}
	default:
		fmt.Fprintf(&g.out, "type %s = %s\n\n", name, g.goType(s, name))
	}
}

func (g *gen) enumValues(s *schema) []string {
	if s.Ref != "" {
		return g.enumValues(g.resolve(s))
	}
	out := []string{}
	for _, e := range s.Enum {
		v := ""
		if json.Unmarshal(e, &v) == nil {
			out = append(out, v)
		}
	}
	for _, a := range append(append(append([]*schema{}, s.OneOf...), s.AnyOf...), s.AllOf...) {
		for _, v := range g.enumValues(a) {
			if !slices.Contains(out, v) {
				out = append(out, v)
			}
		}
	}
	return out
}

type field struct {
	json     string
	schemas  []*schema
	required int
	seen     int
}

type objectAlt struct {
	props    map[string]*schema
	required []string
}

func (g *gen) objectAlts(s *schema) []objectAlt {
	if s.Ref != "" {
		return g.objectAlts(g.resolve(s))
	}
	if len(s.AllOf) == 1 && len(s.Properties) == 0 {
		return g.objectAlts(s.AllOf[0])
	}
	alts := nonNull(append(append([]*schema{}, s.OneOf...), s.AnyOf...))
	base := objectAlt{props: s.Properties, required: s.Required}
	if len(alts) == 0 {
		return []objectAlt{base}
	}
	out := []objectAlt{}
	for _, a := range alts {
		for _, sub := range g.objectAlts(a) {
			props := map[string]*schema{}
			for k, v := range base.props {
				props[k] = v
			}
			for k, v := range sub.props {
				props[k] = v
			}
			out = append(out, objectAlt{props: props, required: append(append([]string{}, base.required...), sub.required...)})
		}
	}
	return out
}

func (g *gen) collectFields(s *schema) []*field {
	alts := g.objectAlts(s)
	byName := map[string]*field{}
	names := []string{}
	for _, a := range alts {
		for k, v := range a.props {
			f := byName[k]
			if f == nil {
				f = &field{json: k}
				byName[k] = f
				names = append(names, k)
			}
			f.schemas = append(f.schemas, v)
			f.seen++
			if slices.Contains(a.required, k) {
				f.required++
			}
		}
	}
	slices.Sort(names)
	out := []*field{}
	for _, n := range names {
		f := byName[n]
		if f.seen < len(alts) {
			f.required = 0
		}
		out = append(out, f)
	}
	return out
}

func (g *gen) emitStruct(name string, fields []*field) {
	body := bytes.Buffer{}
	used := map[string]bool{}
	for _, f := range fields {
		goName := fieldName(f.json)
		for used[goName] {
			goName += "_"
		}
		used[goName] = true

		types := map[string]bool{}
		nullable := false
		var first *schema
		for _, fs := range f.schemas {
			t, n := g.fieldType(fs, name+goName)
			types[t] = true
			nullable = nullable || n
			if first == nil {
				first = fs
			}
		}
		t := ""
		if len(types) == 1 {
			for k := range types {
				t = k
			}
		} else {
			t = "json.RawMessage"
		}

		required := f.required > 0 && f.required == f.seen && !nullable
		if !required && g.kindOf(first) == "struct" && len(types) == 1 && !strings.HasPrefix(t, "*") {
			t = "*" + t
		}
		tag := f.json
		if !required {
			tag += ",omitempty"
		}
		fmt.Fprintf(&body, "%s %s `json:%q`\n", goName, t, tag)
	}
	fmt.Fprintf(&g.out, "type %s struct {\n%s}\n\n", name, body.String())
}

func (g *gen) emitDiscriminators(name string, s *schema) {
	r := g.resolve(s)
	alts := nonNull(append(append([]*schema{}, r.OneOf...), r.AnyOf...))
	if len(alts) < 2 {
		return
	}
	values := []string{}
	for _, a := range alts {
		a = g.resolve(a)
		t := a.Properties["type"]
		if t == nil || len(t.Enum) != 1 {
			return
		}
		v := ""
		if json.Unmarshal(t.Enum[0], &v) != nil {
			return
		}
		values = append(values, v)
	}
	g.consts.WriteString("const (\n")
	for _, v := range values {
		fmt.Fprintf(&g.consts, "%sType%s = %q\n", name, pascal(v), v)
	}
	g.consts.WriteString(")\n\n")
}

func (g *gen) fieldType(s *schema, ctx string) (string, bool) {
	if s.any {
		return "json.RawMessage", false
	}
	if s.Ref != "" {
		return g.named(refName(s.Ref)), false
	}
	if len(s.AllOf) == 1 && len(s.Properties) == 0 {
		return g.fieldType(s.AllOf[0], ctx)
	}
	all := append(append([]*schema{}, s.OneOf...), s.AnyOf...)
	alts := nonNull(all)
	nullable := len(alts) < len(all)
	if len(s.Properties) == 0 && len(alts) == 1 {
		t, n := g.fieldType(alts[0], ctx)
		return t, n || nullable
	}
	ts := s.types()
	if slices.Contains(ts, "null") {
		nullable = true
	}
	return g.goType(s, ctx), nullable
}

func (g *gen) goType(s *schema, ctx string) string {
	switch g.kindOf(s) {
	case "raw":
		return "json.RawMessage"
	case "struct":
		if s.Ref != "" {
			return g.named(refName(s.Ref))
		}
		if !g.emitted[ctx] {
			g.emitted[ctx] = true
			g.emitStruct(ctx, g.collectFields(s))
			g.emitDiscriminators(ctx, s)
		}
		return ctx
	case "string":
		if s.Ref != "" {
			return g.named(refName(s.Ref))
		}
		alts := nonNull(append(append([]*schema{}, s.OneOf...), s.AnyOf...))
		if len(alts) == 1 {
			return g.goType(alts[0], ctx)
		}
		return "string"
	case "scalar":
		if s.Ref != "" {
			return g.named(refName(s.Ref))
		}
		for _, t := range s.types() {
			switch t {
			case "integer":
				return "int64"
			case "number":
				return "float64"
			case "boolean":
				return "bool"
			}
		}
	case "slice":
		if s.Ref != "" {
			return g.named(refName(s.Ref))
		}
		if s.Items == nil {
			return "[]json.RawMessage"
		}
		t, _ := g.fieldType(s.Items, ctx+"Item")
		return "[]" + t
	case "map":
		if s.Ref != "" {
			return g.named(refName(s.Ref))
		}
		if s.AdditionalProperties == nil || s.AdditionalProperties.any {
			return "map[string]any"
		}
		t, _ := g.fieldType(s.AdditionalProperties, ctx+"Value")
		return "map[string]" + t
	}
	if s.Ref != "" {
		return g.named(refName(s.Ref))
	}
	alts := nonNull(append(append([]*schema{}, s.OneOf...), s.AnyOf...))
	if len(alts) == 1 {
		return g.goType(alts[0], ctx)
	}
	return "json.RawMessage"
}

func pascal(s string) string {
	b := strings.Builder{}
	up := true
	for _, r := range s {
		if !unicode.IsLetter(r) && !unicode.IsDigit(r) {
			up = true
			continue
		}
		if up {
			b.WriteRune(unicode.ToUpper(r))
			up = false
			continue
		}
		b.WriteRune(r)
	}
	out := b.String()
	if out == "" || unicode.IsDigit(rune(out[0])) {
		out = "X" + out
	}
	return out
}

func fieldName(s string) string {
	p := []rune(pascal(s))
	out := []rune{}
	for i := 0; i < len(p); i++ {
		end := func(n int) bool { return i+n == len(p) || unicode.IsUpper(p[i+n]) || unicode.IsDigit(p[i+n]) }
		if i+1 < len(p) && p[i] == 'I' && p[i+1] == 'd' && end(2) {
			out = append(out, 'I', 'D')
			i++
			continue
		}
		if i+2 < len(p) && p[i] == 'I' && p[i+1] == 'd' && p[i+2] == 's' && end(3) {
			out = append(out, 'I', 'D', 's')
			i += 2
			continue
		}
		if i+2 < len(p) && p[i] == 'U' && p[i+1] == 'r' && p[i+2] == 'l' && end(3) {
			out = append(out, 'U', 'R', 'L')
			i += 2
			continue
		}
		out = append(out, p[i])
	}
	return string(out)
}