<commit_before>// Copyright ©2019 The go-hep Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rdict

import (
	"bytes"
	"fmt"
	"go/format"
	"go/types"
	"log"
	"reflect"

	"go-hep.org/x/hep/groot/rmeta"
	"golang.org/x/tools/go/packages"
)

// genStreamer holds the state of the ROOT streamer generation.
type genStreamer struct {
	buf *bytes.Buffer
	pkg *types.Package

	// set of imported packages.
	// usually: "encoding/binary", "math"
	imps map[string]int

	verbose bool // enable verbose mode

	binMa *types.Interface // encoding.BinaryMarshaler
	binUn *types.Interface // encoding.BinaryUnmarshaler
	rvers *types.Interface // rbytes.RVersioner

	gosizes types.Sizes
}

func importPkg(p string) (*types.Package, error) {
	cfg := &packages.Config{Mode: packages.NeedTypes | packages.NeedTypesInfo | packages.NeedTypesSizes | packages.NeedDeps}
	pkgs, err := packages.Load(cfg, p)
	if err != nil {
		return nil, fmt.Errorf("could not load package %q: %w", p, err)
	}

	return pkgs[0].Types, nil
}

// NewGenStreamer returns a new code generator for package p,
// where p is the package's import path.
func NewGenStreamer(p string, verbose bool) (Generator, error) {
	pkg, err := importPkg(p)
	if err != nil {
		return nil, err
	}
	g := &genStreamer{
		buf: new(bytes.Buffer),
		pkg: pkg,
		imps: map[string]int{
			"go-hep.org/x/hep/groot/rbytes": 1,
			"go-hep.org/x/hep/groot/rdict":  1,
			"go-hep.org/x/hep/groot/rmeta":  1,
		},
		verbose: verbose,
	}

	err = g.init()
	if err != nil {
		return nil, err
	}

	return g, nil
}

func (g *genStreamer) init() error {
	pkg, err := importPkg("encoding")
	if err != nil {
		return fmt.Errorf("rdict: could not find package \"encoding\": %w", err)
	}

	o := pkg.Scope().Lookup("BinaryMarshaler")
	if o == nil {
		return fmt.Errorf("rdict: could not find interface encoding.BinaryMarshaler")
	}
	g.binMa = o.(*types.TypeName).Type().Underlying().(*types.Interface)

	o = pkg.Scope().Lookup("BinaryUnmarshaler")
	if o == nil {
		return fmt.Errorf("rdict: could not find interface encoding.BinaryUnmarshaler")
	}
	g.binUn = o.(*types.TypeName).Type().Underlying().(*types.Interface)

	pkg, err = importPkg("go-hep.org/x/hep/groot/rbytes")
	if err != nil {
		return fmt.Errorf("rdict: could not find package %q: %w", "go-hep.org/x/hep/groot/rbytes", err)
	}

	o = pkg.Scope().Lookup("RVersioner")
	if o == nil {
		return fmt.Errorf("rdict: could not find interface rbytes.RVersioner")
	}
	g.rvers = o.(*types.TypeName).Type().Underlying().(*types.Interface)

	sz := int64(reflect.TypeOf(int(0)).Size())
	g.gosizes = &types.StdSizes{WordSize: sz, MaxAlign: sz}
	return nil
}

func (g *genStreamer) printf(format string, args ...interface{}) {
	fmt.Fprintf(g.buf, format, args...)
}

// Generate implements rdict.Generator
func (g *genStreamer) Generate(typeName string) error {
	scope := g.pkg.Scope()
	obj := scope.Lookup(typeName)
	if obj == nil {
		return fmt.Errorf("no such type %q in package %q", typeName, g.pkg.Path()+"/"+g.pkg.Name())
	}

	tn, ok := obj.(*types.TypeName)
	if !ok {
		return fmt.Errorf("%q is not a type (%v)", typeName, obj)
	}

	typ, ok := tn.Type().Underlying().(*types.Struct)
	if !ok {
		return fmt.Errorf("%q is not a named struct (%v)", typeName, tn)
	}
	if g.verbose {
		log.Printf("typ: %q: %+v\n", typeName, typ)
	}

	if !types.Implements(tn.Type(), g.rvers) && !types.Implements(types.NewPointer(tn.Type()), g.rvers) {
		return fmt.Errorf("type %q does not implement %q", tn.Pkg().Path()+"."+tn.Name(), "go-hep.org/x/hep/groot/rbytes.RVersioner")
	}

	g.genStreamer(typ, typeName)
	g.genMarshal(typ, typeName)
	// g.genUnmarshal(typ, typeName)

	return nil
}

func (g *genStreamer) genMarshal(t types.Type, typeName string) {
	g.printf(`// MarshalROOT implements rbytes.Marshaler
func (o *%[1]s) MarshalROOT(w *rbytes.WBuffer) (int, error) {
	if w.Err() != nil {
		return 0, w.Err()
	}

	pos := w.WriteVersion(o.RVersion())

`,
		typeName,
	)

	typ := t.Underlying().(*types.Struct)
	for i := 0; i < typ.NumFields(); i++ {
		ft := typ.Field(i)
		n := ft.Name() // no `groot:"foo"` redirection.
		g.genMarshalType(ft.Type(), n)
	}

	g.printf("\n\treturn w.SetByteCount(pos, o.Class())\n}\n\n")
}

// func (g *genStreamer) genUnmarshal(t types.Type, typeName string) {
// 	g.printf(`// UnmarshalROOT implements rbytes.Unmarshaler
// func (o *%[1]s) UnmarshalROOT(r *rbytes.RBuffer) error {
// 	rs, err := rdict.RStreamer(r, o)
// 	if err != nil {
// 		return err
// 	}
// 	return rs.RStreamROOT(r)
// }
// `,
// 		typeName,
// 	)
// }

func (g *genStreamer) genStreamer(t types.Type, typeName string) {
	g.printf(`func init() {<commit_msg>groot/rdict: generate RVersioner if not implemented<commit_after>// Copyright ©2019 The go-hep Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package rdict

import (
	"bytes"
	"fmt"
	"go/format"
	"go/types"
	"log"
	"reflect"

	"go-hep.org/x/hep/groot/rmeta"
	"golang.org/x/tools/go/packages"
)

// genStreamer holds the state of the ROOT streamer generation.
type genStreamer struct {
	buf *bytes.Buffer
	pkg *types.Package

	// set of imported packages.
	// usually: "encoding/binary", "math"
	imps map[string]int

	verbose bool // enable verbose mode

	binMa *types.Interface // encoding.BinaryMarshaler
	binUn *types.Interface // encoding.BinaryUnmarshaler
	rvers *types.Interface // rbytes.RVersioner

	gosizes types.Sizes
}

func importPkg(p string) (*types.Package, error) {
	cfg := &packages.Config{Mode: packages.NeedTypes | packages.NeedTypesInfo | packages.NeedTypesSizes | packages.NeedDeps}
	pkgs, err := packages.Load(cfg, p)
	if err != nil {
		return nil, fmt.Errorf("could not load package %q: %w", p, err)
	}

	return pkgs[0].Types, nil
}

// NewGenStreamer returns a new code generator for package p,
// where p is the package's import path.
func NewGenStreamer(p string, verbose bool) (Generator, error) {
	pkg, err := importPkg(p)
	if err != nil {
		return nil, err
	}
	g := &genStreamer{
		buf: new(bytes.Buffer),
		pkg: pkg,
		imps: map[string]int{
			"go-hep.org/x/hep/groot/rbytes": 1,
			"go-hep.org/x/hep/groot/rdict":  1,
			"go-hep.org/x/hep/groot/rmeta":  1,
		},
		verbose: verbose,
	}

	err = g.init()
	if err != nil {
		return nil, err
	}

	return g, nil
}

func (g *genStreamer) init() error {
	pkg, err := importPkg("encoding")
	if err != nil {
		return fmt.Errorf("rdict: could not find package \"encoding\": %w", err)
	}

	o := pkg.Scope().Lookup("BinaryMarshaler")
	if o == nil {
		return fmt.Errorf("rdict: could not find interface encoding.BinaryMarshaler")
	}
	g.binMa = o.(*types.TypeName).Type().Underlying().(*types.Interface)

	o = pkg.Scope().Lookup("BinaryUnmarshaler")
	if o == nil {
		return fmt.Errorf("rdict: could not find interface encoding.BinaryUnmarshaler")
	}
	g.binUn = o.(*types.TypeName).Type().Underlying().(*types.Interface)

	pkg, err = importPkg("go-hep.org/x/hep/groot/rbytes")
	if err != nil {
		return fmt.Errorf("rdict: could not find package %q: %w", "go-hep.org/x/hep/groot/rbytes", err)
	}

	o = pkg.Scope().Lookup("RVersioner")
	if o == nil {
		return fmt.Errorf("rdict: could not find interface rbytes.RVersioner")
	}
	g.rvers = o.(*types.TypeName).Type().Underlying().(*types.Interface)

	sz := int64(reflect.TypeOf(int(0)).Size())
	g.gosizes = &types.StdSizes{WordSize: sz, MaxAlign: sz}
	return nil
}

func (g *genStreamer) printf(format string, args ...interface{}) {
	fmt.Fprintf(g.buf, format, args...)
}

// Generate implements rdict.Generator
func (g *genStreamer) Generate(typeName string) error {
	scope := g.pkg.Scope()
	obj := scope.Lookup(typeName)
	if obj == nil {
		return fmt.Errorf("no such type %q in package %q", typeName, g.pkg.Path()+"/"+g.pkg.Name())
	}

	tn, ok := obj.(*types.TypeName)
	if !ok {
		return fmt.Errorf("%q is not a type (%v)", typeName, obj)
	}

	typ, ok := tn.Type().Underlying().(*types.Struct)
	if !ok {
		return fmt.Errorf("%q is not a named struct (%v)", typeName, tn)
	}
	if g.verbose {
		log.Printf("typ: %q: %+v\n", typeName, typ)
	}

	if !types.Implements(tn.Type(), g.rvers) && !types.Implements(types.NewPointer(tn.Type()), g.rvers) {
		g.genRVersioner(typ, typeName)
	}

	g.genStreamer(typ, typeName)
	g.genMarshal(typ, typeName)
	// g.genUnmarshal(typ, typeName)

	return nil
}

func (g *genStreamer) genMarshal(t types.Type, typeName string) {
	g.printf(`// MarshalROOT implements rbytes.Marshaler
func (o *%[1]s) MarshalROOT(w *rbytes.WBuffer) (int, error) {
	if w.Err() != nil {
		return 0, w.Err()
	}

	pos := w.WriteVersion(o.RVersion())

`,
		typeName,
	)

	typ := t.Underlying().(*types.Struct)
	for i := 0; i < typ.NumFields(); i++ {
		ft := typ.Field(i)
		n := ft.Name() // no `groot:"foo"` redirection.
		g.genMarshalType(ft.Type(), n)
	}

	g.printf("\n\treturn w.SetByteCount(pos, o.Class())\n}\n\n")
}

// func (g *genStreamer) genUnmarshal(t types.Type, typeName string) {
// 	g.printf(`// UnmarshalROOT implements rbytes.Unmarshaler
// func (o *%[1]s) UnmarshalROOT(r *rbytes.RBuffer) error {
// 	rs, err := rdict.RStreamer(r, o)
// 	if err != nil {
// 		return err
// 	}
// 	return rs.RStreamROOT(r)
// }
// `,
// 		typeName,
// 	)
// }

func (g *genStreamer) genRVersioner(t types.Type, typeName string) {
	g.printf("func (*%s) RVersion() int16 { return 1 }\n\n", typeName)
}

func (g *genStreamer) genStreamer(t types.Type, typeName string) {
	g.printf(`func init() {
{'fraction_non_alphanumeric': 0.17475013150973173, 'fraction_numerical': 0.0034718569174118884, 'mean_word_length': 3.70827142149579, 'pattern_counts': {'":': 8, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1752653', 'n_tokens_mistral': 3745, 'n_tokens_neox': 3562, 'n_words': 885}