// Copyright 2026 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// refgen produces a reference implementation of the SIMD API backed by the spec
// implementation.
package main

import (
	"bytes"
	"cmp"
	"flag"
	"fmt"
	"go/types"
	"log"
	"maps"
	"os"
	"slices"
	"strings"

	"simd/archsimd/_gen/gentools"
	"simd/archsimd/_gen/specgen"
	"simd/archsimd/_gen/specgen/specexpr"
)

func main() {
	gentools.RegisterFlags(nil)

	flag.Usage = func() {
		w := flag.CommandLine.Output()
		fmt.Fprintf(w, "usage: refgen [flags] [spec dir]\n")
		flag.CommandLine.PrintDefaults()
	}

	flag.Parse()
	var specDir string
	switch flag.NArg() {
	case 0:
		specDir = specgen.MustFindSpecDir()
	case 1:
		specDir = flag.Arg(0)
	default:
		flag.Usage()
		os.Exit(1)
	}

	funcs, err := specgen.Load(specDir, nil)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s\n", err.Error())
		os.Exit(1)
	}

	var files gentools.Files
	defer files.FlushOrExit()

	src := &srcWriter{Buffer: files.NewGoFile("simd/internal/simdref/simdref.go")}

	fmt.Fprintf(src, `// Code generated by 'refgen'. DO NOT EDIT.

package simdref

import "simd/internal/spec"

`)

	// Define all vector types
	vecTypeSet := make(map[specexpr.Vector]bool)
	for _, fn := range funcs {
		if fn.Recv.Type != nil {
			vecTypeSet[fn.Recv.Type.(specexpr.Vector)] = true
		}
	}
	vecTypes := slices.SortedFunc(maps.Keys(vecTypeSet), func(a, b specexpr.Vector) int {
		if a.Elem != b.Elem {
			return cmp.Compare(a.Elem.String(), b.Elem.String())
		}
		cmp, ok := a.Width.Compare(b.Width)
		if ok {
			return cmp
		}
		_, aScale := a.Width.(specexpr.ScalableWidth)
		_, bScale := b.Width.(specexpr.ScalableWidth)
		if !aScale && bScale {
			return 1
		}
		return -1
	})
	fmt.Fprintf(src, "type (\n")
	for _, vec := range vecTypes {
		elem := vec.Elem.String()
		if vec.Elem.Base == "Mask" {
			elem = fmt.Sprintf("spec.Mask%d", vec.Elem.Bits)
		}

		fmt.Fprintf(src, "\t%s struct { v []%s }\n", vec.String(), elem)
	}
	fmt.Fprintf(src, ")\n\n")

	// Define functions
	var args []string
	for _, fn := range funcs {
		fmt.Fprintf(src, "%s {\n", fn.Decl())
		src.id = 0

		specName, specSig, typeArgs := fn.SpecFunc()
		specInst, err := types.Instantiate(nil, specSig, typeArgs, false)
		if err != nil {
			panic(fmt.Sprintf("instantiating spec function %s: %s", specName, err))
		}

		specParams := specInst.(*types.Signature).Params()
		args = args[:0]
		if fn.Recv.Type != nil {
			args = append(args, toSpec(fn.Recv.Type, specParams.At(len(args)).Type(), fn.Recv.Name, src))
		}
		for _, in := range fn.In {
			args = append(args, toSpec(in.Type, specParams.At(len(args)).Type(), in.Name, src))
		}

		call := formatCall(specName, typeArgs, args)

		specResults := specInst.(*types.Signature).Results()
		switch len(fn.Out) {
		case 0:
			fmt.Fprintf(src, "\t%s\n", call)
		case 1:
			fmt.Fprintf(src, "\treturn %s\n", fromSpec(fn.Out[0].Type, specResults.At(0).Type(), call, src))
		default:
			var tmps []string
			var res []string
			for i := range fn.Out {
				tmp := fmt.Sprintf("r%d", i+1)
				tmps = append(tmps, tmp)
				res = append(res, fromSpec(fn.Out[i].Type, specResults.At(i).Type(), tmp, src))
			}
			fmt.Fprintf(src, "\t%s := %s\n", strings.Join(tmps, ", "), call)
			fmt.Fprintf(src, "\treturn %s\n", strings.Join(res, ", "))
		}
		fmt.Fprintf(src, "}\n\n")
	}
}

type srcWriter struct {
	*bytes.Buffer
	id int
}

func (w *srcWriter) genIdent() string {
	ident := fmt.Sprintf("tmp%d", w.id)
	w.id++
	return ident
}

func formatCall(specName string, typeArgs []types.Type, args []string) string {
	var callBuf bytes.Buffer
	fmt.Fprintf(&callBuf, "spec.%s[", specName)
	for i, typeArg := range typeArgs {
		if i > 0 {
			callBuf.WriteString(", ")
		}
		types.WriteType(&callBuf, typeArg, specQualifier)
	}
	fmt.Fprintf(&callBuf, "](%s)", strings.Join(args, ", "))
	return callBuf.String()
}

func specQualifier(pkg *types.Package) string {
	if pkg.Path() == "simd/internal/spec" {
		return "spec"
	}
	return ""
}

// toSpec returns an expression that converts val from the specref Go type for t
// to spec type tt. It may write statements to src.
func toSpec(t specexpr.Type, tt types.Type, val string, src *srcWriter) string {
	arrayOrSlice := func(tElem specexpr.Type, ttElem types.Type) string {
		if eVal := toSpec(tElem, ttElem, val, src); eVal == val {
			// Easy case: the values don't need to change.
			return val
		}
		// Hard case: we need to map each element
		tmp := src.genIdent()
		fmt.Fprintf(src, "var %s %s\n", tmp, types.TypeString(tt, specQualifier))
		fmt.Fprintf(src, "for i := range %s {\n", val)
		eVal := fromSpec(tElem, ttElem, "("+val+")[i]", src)
		fmt.Fprintf(src, "\t%s[i] = %s\n", tmp, eVal)
		fmt.Fprintf(src, "}\n")
		return tmp
	}

	switch t := t.(type) {
	case specexpr.Vector:
		return val + ".v"
	case specexpr.Basic:
		switch tt := tt.(type) {
		case *types.Named:
			if tt.Obj().Name() == "UintN" {
				return "spec.UintN(" + val + ")"
			}
		}
		return val
	case specexpr.Slice:
		tt := tt.Underlying().(*types.Slice)
		return arrayOrSlice(t.Elem, tt.Elem())
	case specexpr.Array:
		switch tt := tt.(type) {
		case *types.Named:
			if tt.Obj().Name() == "Array" {
				return toSpec(t.Elem, tt.TypeArgs().At(0), val, src) + "[:]"
			}
		}
		tt := tt.Underlying().(*types.Array)
		return arrayOrSlice(t.Elem, tt.Elem())
	case specexpr.Pointer:
		tt := tt.(*types.Pointer)
		eVal := toSpec(t.Elem, tt.Elem(), val, src)
		tmp := src.genIdent()
		fmt.Fprintf(src, "var %s %s = %s\n", tmp, types.TypeString(tt.Elem(), specQualifier), eVal)
		return "&" + tmp
	}
	log.Fatalf("unexpected specexpr type %s (%T)", t, t)
	panic("not reachable")
}

// fromSpec returns an expression that converts val from the spec package type
// tt to the specref Go type for t. It may write statements to src.
func fromSpec(t specexpr.Type, tt types.Type, val string, src *srcWriter) string {
	arrayOrSlice := func(tElem specexpr.Type, ttElem types.Type) string {
		if eVal := fromSpec(tElem, ttElem, val, src); eVal == val {
			// Easy case: the values don't need to change.
			return val
		}
		// Hard case: we need to map each element
		tmp := src.genIdent()
		fmt.Fprintf(src, "var %s %s\n", tmp, t)
		fmt.Fprintf(src, "for i := range %s {\n", val)
		eVal := fromSpec(tElem, ttElem, "("+val+")[i]", src)
		fmt.Fprintf(src, "\t%s[i] = %s\n", tmp, eVal)
		fmt.Fprintf(src, "}\n")
		return tmp
	}

	switch t := t.(type) {
	case specexpr.Vector:
		return fmt.Sprintf("%s{%s}", t, val)
	case specexpr.Basic:
		switch tt := tt.(type) {
		case *types.Named:
			if tt.Obj().Name() == "UintN" {
				return fmt.Sprintf("%s(%s)", t, val)
			}
		}
		return val
	case specexpr.Slice:
		tt := tt.Underlying().(*types.Slice)
		return arrayOrSlice(t.Elem, tt.Elem())
	case specexpr.Array:
		switch tt := tt.(type) {
		case *types.Named:
			if tt.Obj().Name() == "Array" {
				return fmt.Sprintf("(%s)(%s)", t, fromSpec(t.Elem, tt.TypeArgs().At(0), val, src))
			}
		}
		tt := tt.Underlying().(*types.Array)
		return arrayOrSlice(t.Elem, tt.Elem())
	case specexpr.Pointer:
		tt := tt.(*types.Pointer)
		eVal := fromSpec(t.Elem, tt.Elem(), val, src)
		tmp := src.genIdent()
		fmt.Fprintf(src, "var %s %s = %s\n", tmp, t, eVal)
		return "&" + tmp
	}
	log.Fatalf("unexpected specexpr type %s (%T)", t, t)
	panic("not reachable")
}