// Copyright 2025 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.

package main

import (
	"fmt"
	"strings"
)

const simdssaTmpl = `// Code generated by x/arch/internal/simdgen using 'go run . -xedPath $XED_PATH -o godefs -goroot $GOROOT go.yaml types.yaml categories.yaml'; DO NOT EDIT.

package amd64

import (
	"cmd/compile/internal/ssa"
	"cmd/compile/internal/ssagen"
	"cmd/internal/obj"
	"cmd/internal/obj/x86"
)

func ssaGenSIMDValue(s *ssagen.State, v *ssa.Value) bool {
	p := s.Prog(v.Op.Asm())
	// First arg
	switch v.Op {{"{"}}{{if gt (len .ImmFirst) 0}}
	// Imm
	case {{.ImmFirst}}:
		imm := v.AuxInt
		if imm < 0 || imm > 255 {
			v.Fatalf("Invalid source selection immediate")
		}
		p.From.Offset = imm
		p.From.Type = obj.TYPE_CONST
{{end}}{{if gt (len .VregFirst) 0}}
	// vreg
	case {{.VregFirst}}:
		p.From.Type = obj.TYPE_REG
		p.From.Reg = simdReg(v.Args[0])
{{end}}
	default:
		// At least one arg is required.
		return false
	}

	// Second arg
	switch v.Op {{"{"}}{{if gt (len .VregSecond) 0}}
	// vreg
	case {{.VregSecond}}:
		if p.From.Type == obj.TYPE_CONST {
			p.AddRestSourceReg(simdReg(v.Args[0]))
		} else {
			p.AddRestSourceReg(simdReg(v.Args[1]))
		}{{end}}
	}

	// Third arg
	switch v.Op {{"{"}}{{if gt (len .VregThird) 0}}
	// vreg
	case {{.VregThird}}:
		if p.From.Type == obj.TYPE_CONST {
			p.AddRestSourceReg(simdReg(v.Args[1]))
		} else {
			p.AddRestSourceReg(simdReg(v.Args[2]))
		}
{{end}}{{if gt (len .MaskThird) 0}}
	// k mask
	case {{.MaskThird}}:
		if p.From.Type == obj.TYPE_CONST {
			p.AddRestSourceReg(v.Args[1].Reg())
		} else {
			p.AddRestSourceReg(v.Args[2].Reg())
		}{{end}}
	}

	// Fourth arg
	switch v.Op {{"{"}}{{if gt (len .MaskFourth) 0}}
	case {{.MaskFourth}}:
		if p.From.Type == obj.TYPE_CONST {
			p.AddRestSourceReg(v.Args[2].Reg())
		} else {
			p.AddRestSourceReg(v.Args[3].Reg())
		}{{end}}
	}

	// Output
	switch v.Op {{"{"}}{{if gt (len .VregOut) 0}}
	case {{.VregOut}}:
		p.To.Type = obj.TYPE_REG
		p.To.Reg = simdReg(v)
{{end}}{{if gt (len .MaskOut) 0}}
	case {{.MaskOut}}:
		p.To.Type = obj.TYPE_REG
		p.To.Reg = v.Reg()
{{end}}
	default:
		// One result is required.
		return false
	}
{{if gt (len .ZeroingMask) 0}}
	// Masked operation are always compiled with zeroing.
	switch v.Op {
	case {{.ZeroingMask}}:
		x86.ParseSuffix(p, "Z")
	}
{{end}}
	return true
}
`

// writeSIMDSSA generates the ssa to prog lowering codes and writes it to simdssa.go
// within the specified directory.
func writeSIMDSSA(directory string, ops []Operation) error {
	var ImmFirst []string
	var VregFirst []string
	var VregSecond []string
	var MaskThird []string
	var VregThird []string
	var MaskFourth []string
	var VregOut []string
	var MaskOut []string
	var ZeroingMask []string

	seen := map[string]struct{}{}
	for _, op := range ops {
		asm := op.Asm
		shapeIn, shapeOut, maskType, _, _, gOp, err := op.shape()
		if err != nil {
			return err
		}
		if maskType == 2 {
			asm += "Masked"
		}
		asm = fmt.Sprintf("%s%d", asm, *gOp.Out[0].Bits)
		if _, ok := seen[asm]; ok {
			continue
		}
		seen[asm] = struct{}{}
		caseStr := fmt.Sprintf("ssa.OpAMD64%s", asm)
		if shapeIn == PureVregIn || shapeIn == PureKmaskIn {
			// Masks and vreg are handled together by simdReg()
			VregFirst = append(VregFirst, caseStr)
			if len(gOp.In) > 1 {
				VregSecond = append(VregSecond, caseStr)
			}
		} else if shapeIn == OneKmaskIn {
			VregFirst = append(VregFirst, caseStr)
			VregSecond = append(VregSecond, caseStr)
			MaskThird = append(MaskThird, caseStr)
			if gOp.Zeroing == nil {
				ZeroingMask = append(ZeroingMask, caseStr)
			}
		} else if shapeIn == OneConstImmIn {
			ImmFirst = append(ImmFirst, caseStr)
			VregSecond = append(VregSecond, caseStr)
			VregThird = append(VregThird, caseStr)
		} else {
			// OneKmaskConstImmIn case
			ImmFirst = append(ImmFirst, caseStr)
			VregSecond = append(VregSecond, caseStr)
			VregThird = append(VregThird, caseStr)
			MaskFourth = append(MaskFourth, caseStr)
			if gOp.Zeroing == nil {
				ZeroingMask = append(ZeroingMask, caseStr)
			}
		}
		if shapeOut == OneVregOut || gOp.Out[0].OverwriteClass != nil {
			// If class overwrite is happening, that's not really a mask but a vreg.
			VregOut = append(VregOut, caseStr)
		} else {
			// OneKmaskOut case
			MaskOut = append(MaskOut, caseStr)
		}
	}

	data := struct {
		ImmFirst    string
		VregFirst   string
		VregSecond  string
		MaskThird   string
		VregThird   string
		MaskFourth  string
		VregOut     string
		MaskOut     string
		ZeroingMask string
	}{
		strings.Join(ImmFirst, ", "),
		strings.Join(VregFirst, ", "),
		strings.Join(VregSecond, ", "),
		strings.Join(MaskThird, ", "),
		strings.Join(VregThird, ", "),
		strings.Join(MaskFourth, ", "),
		strings.Join(VregOut, ", "),
		strings.Join(MaskOut, ", "),
		strings.Join(ZeroingMask, ", "),
	}

	file, t, err := openFileAndPrepareTemplate(directory, "src/cmd/compile/internal/amd64/simdssa.go", simdssaTmpl)
	if err != nil {
		return err
	}
	defer file.Close()

	err = t.Execute(file, data)
	if err != nil {
		return fmt.Errorf("failed to execute template: %w", err)
	}

	return nil
}