package generator

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
)

// Obj is a JSON object that remembers key order, so generated structs list
// fields in the same order as the schema rather than alphabetically.
type Obj struct {
	Keys []string
	M    map[string]any
}

func newObj() *Obj { return &Obj{M: map[string]any{}} }

func (o *Obj) Get(k string) any { return o.M[k] }

func (o *Obj) Set(k string, v any) {
	if _, ok := o.M[k]; !ok {
		o.Keys = append(o.Keys, k)
	}
	o.M[k] = v
}

func (o *Obj) Delete(k string) {
	if _, ok := o.M[k]; !ok {
		return
	}
	delete(o.M, k)
	for i, key := range o.Keys {
		if key == k {
			o.Keys = append(o.Keys[:i], o.Keys[i+1:]...)
			break
		}
	}
}

func (o *Obj) Obj(k string) *Obj {
	v, _ := o.M[k].(*Obj)
	return v
}

func (o *Obj) Str(k string) string {
	v, _ := o.M[k].(string)
	return v
}

// decodeOrdered parses JSON, returning *Obj for objects, []any for arrays,
// and string / json.Number / bool / nil for scalars.
func decodeOrdered(r io.Reader) (any, error) {
	dec := json.NewDecoder(r)
	dec.UseNumber()
	return parseValue(dec)
}

func parseValue(dec *json.Decoder) (any, error) {
	tok, err := dec.Token()
	if err != nil {
		return nil, err
	}
	switch t := tok.(type) {
	case json.Delim:
		switch t {
		case '{':
			o := newObj()
			for dec.More() {
				kt, err := dec.Token()
				if err != nil {
					return nil, err
				}
				k, ok := kt.(string)
				if !ok {
					return nil, fmt.Errorf("expected object key, got %v", kt)
				}
				v, err := parseValue(dec)
				if err != nil {
					return nil, err
				}
				o.Set(k, v)
			}
			_, err := dec.Token() // '}'
			return o, err
		case '[':
			arr := []any{}
			for dec.More() {
				v, err := parseValue(dec)
				if err != nil {
					return nil, err
				}
				arr = append(arr, v)
			}
			_, err := dec.Token() // ']'
			return arr, err
		}
		return nil, fmt.Errorf("unexpected delimiter %v", t)
	default:
		return t, nil
	}
}

// mergePatch applies an RFC 7386 JSON Merge Patch, which is how OCDS
// extensions are applied to the core schema. A null in the patch deletes
// the key; objects merge recursively; anything else replaces.
func mergePatch(target, patch any) any {
	p, ok := patch.(*Obj)
	if !ok {
		return patch
	}
	t, ok := target.(*Obj)
	if !ok {
		t = newObj()
	}
	for _, k := range p.Keys {
		v := p.M[k]
		if v == nil {
			t.Delete(k)
			continue
		}
		t.Set(k, mergePatch(t.M[k], v))
	}
	return t
}

// MarshalJSON writes the object with its keys in their original order.
func (o *Obj) MarshalJSON() ([]byte, error) {
	var b bytes.Buffer
	b.WriteByte('{')
	for i, k := range o.Keys {
		if i > 0 {
			b.WriteByte(',')
		}
		kb, err := json.Marshal(k)
		if err != nil {
			return nil, err
		}
		vb, err := json.Marshal(o.M[k])
		if err != nil {
			return nil, err
		}
		b.Write(kb)
		b.WriteByte(':')
		b.Write(vb)
	}
	b.WriteByte('}')
	return b.Bytes(), nil
}