// Package rosa provides Rosa OS toolchain artifacts and miscellaneous software.
package rosa

import (
	"embed"
	"io/fs"
	"os"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"
	"unsafe"

	"hakurei.app/fhs"
	"hakurei.app/internal/pkg"
)

// Extension is the variant identification string of custom artifact
// implementations registered by package rosa.
const Extension = "rosa"

func init() { pkg.SetExtension(Extension) }

const (
	// kindEtc is the kind of [pkg.Artifact] of cureEtc.
	kindEtc = iota + pkg.KindCustomOffset

	// kindGentooOverlay is the kind of [pkg.Artifact] of gentooOverlay.
	kindGentooOverlay
)

// mustDecode is like [pkg.MustDecode], but replaces the zero value and prints
// a warning.
func mustDecode(s string) pkg.Checksum {
	var fallback = pkg.Checksum{}
	if s == "" {
		println(
			"falling back to",
			pkg.Encode(fallback),
			"for unpopulated checksum",
		)
		return fallback
	}
	return pkg.MustDecode(s)
}

// KV is a key-value pair of strings.
type KV = [2]string

var (
	// AbsUsrSrc is the conventional directory to place source code under.
	AbsUsrSrc = fhs.AbsUsr.Append("src")

	// AbsSystem is the Rosa OS installation prefix.
	AbsSystem = fhs.AbsRoot.Append("system")
)

// linuxArch returns the architecture name used by linux corresponding to arch.
func (s *S) linuxArch() string {
	switch s.arch {
	case "amd64":
		return "x86_64"
	case "arm64":
		return "aarch64"
	case "riscv64":
		return "riscv64"

	default:
		panic("unsupported target " + s.arch)
	}
}

// triple returns the Rosa OS host triple corresponding to arch.
func (s *S) triple() string {
	return s.linuxArch() + "-rosa-linux-musl"
}

const (
	// EnvTriple holds the Rosa OS host triple.
	EnvTriple = "ROSA_TRIPLE"
)

// earlyLDFLAGS returns LDFLAGS corresponding to triplet.
func (s *S) earlyLDFLAGS(static bool) string {
	p := "-fuse-ld=lld " +
		"-L/system/lib -Wl,-rpath=/system/lib " +
		"-L/system/lib/" + s.triple() + " " +
		"-Wl,-rpath=/system/lib/" + s.triple() + " " +
		"-rtlib=compiler-rt " +
		"-unwindlib=libunwind " +
		"-Wl,--as-needed"
	if !static {
		p += " -Wl,--dynamic-linker=/system/bin/linker"
	}
	return p
}

const (
	// stageGentoo denotes the toolchain in a Gentoo stage3 tarball. Special
	// care must be taken to compile correctly against this stage.
	stageGentoo Stage = iota

	// stageIntermediateGentoo is like stageIntermediate, but compiled against
	// stageGentoo.
	stageIntermediateGentoo

	// stageStdGentoo is like Std, but bootstrapped from stageGentoo. This
	// toolchain creates the first stage0 distribution.
	stageStdGentoo

	// stageEarly denotes the stage0 toolchain. Special care must be taken
	// to compile correctly against this toolchain.
	stageEarly

	// stageIntermediate denotes the intermediate toolchain compiled against
	// stageEarly. This toolchain should be functionally identical to [Std]
	// and is used to bootstrap [Std].
	stageIntermediate

	// Std denotes the standard Rosa OS toolchain.
	Std

	// Stage3 denotes the stage3 Rosa OS toolchain built on [Std]. Software
	// built on this toolchain should be identical to [Std]. This is generally
	// for validation of LLVM 3-stage nondeterminism only, due to dynamic
	// linking complications.
	Stage3

	// _stageEnd is the total number of stages available and does not denote a
	// valid toolchain.
	_stageEnd
)

// isStage0 returns whether t is a stage0 toolchain.
func (t Stage) isStage0() bool {
	switch t {
	case stageGentoo, stageEarly:
		return true
	default:
		return false
	}
}

// isIntermediate returns whether t is an intermediate toolchain.
func (t Stage) isIntermediate() bool {
	switch t {
	case stageIntermediateGentoo, stageIntermediate:
		return true
	default:
		return false
	}
}

// isStd returns whether t is considered functionally equivalent to [Std].
func (t Stage) isStd() bool {
	switch t {
	case stageStdGentoo, Std, Stage3:
		return true
	default:
		return false
	}
}

// lastIndexFunc is like [strings.LastIndexFunc] but for [slices].
func lastIndexFunc[S ~[]E, E any](s S, f func(E) bool) (i int) {
	if i = slices.IndexFunc(s, f); i < 0 {
		return
	}
	if i0 := lastIndexFunc[S](s[i+1:], f); i0 >= 0 {
		i = i0
	}
	return
}

// fixupEnviron fixes up PATH, prepends extras and returns the resulting slice.
func fixupEnviron(env, extras []string, paths ...string) []string {
	// some python tools try to be clever and buffers their output, making the
	// build process appear to hang
	env = append(env, "PYTHONUNBUFFERED=1")

	const pathPrefix = "PATH="
	pathVal := strings.Join(paths, ":")

	if i := lastIndexFunc(env, func(s string) bool {
		return strings.HasPrefix(s, pathPrefix)
	}); i < 0 {
		env = append(env, pathPrefix+pathVal)
	} else {
		if len(env[i]) == len(pathPrefix) {
			env[i] = pathPrefix + pathVal
		} else {
			env[i] += ":" + pathVal
		}
	}

	return append(extras, env...)
}

// scriptName is the name of the fixed-up build script.
const scriptName = "all"

// absCureScript is the absolute pathname [Toolchain.New] places the fixed-up
// build script under.
var absCureScript = AbsSystem.Append(scriptName)

var (
	_stage0Dist  = H("stage0-dist")
	_mksh        = H("mksh")
	_toybox      = H("toybox")
	_toyboxEarly = H("toybox-early")
	_llvm        = H("llvm")
	_patch       = H("patch")
)

// HasStageEarly returns whether a stage0 distribution is available.
func (s *S) HasStageEarly() (ok bool) {
	func() {
		defer func() { ok = recover() == nil }()
		s.New(stageEarly).MustLoad(_stage0Dist)
	}()
	return
}

// NewPatchedSource returns [pkg.Artifact] of source with patches applied. If
// passthrough is true, source is returned as is for zero-length patches.
func (t Toolchain) NewPatchedSource(
	name string,
	source pkg.Artifact,
	passthrough bool,
	patches ...KV,
) pkg.Artifact {
	if passthrough && len(patches) == 0 {
		return source
	}

	paths := make([]pkg.ExecPath, len(patches)+1)
	for i, p := range patches {
		paths[i+1] = pkg.Path(
			fhs.AbsRoot.Append("patches", p[0]), false,
			pkg.NewFile(p[0], unsafe.Slice(unsafe.StringData(p[1]), len(p[1]))),
		)
	}
	paths[0] = pkg.Path(fhs.AbsRoot.Append("src"), false, source)

	var buf strings.Builder
	buf.WriteString(`
cp -r /src/. /work/.
chmod -R +w /work && cd /work
`)
	if len(paths) > 1 {
		buf.WriteString(`
function apply {
	echo "applying $1"
	patch -p 1 < "/patches/$1"
}
`)

		for _, p := range patches {
			buf.WriteString("apply '")
			buf.WriteString(p[0])
			buf.WriteString("'\n")
		}
	}

	return t.New(name+"-src", Unversioned, nil, &PackageAttr{
		Paths: paths,
	}, &GenericHelper{
		Build: buf.String(),
	},
		_patch,
	)
}

// helperInPlace is a special directory value for omitting the cd statement.
const helperInPlace = "\x00"

// Helper is a build system helper for [Toolchain.NewPackage].
type Helper interface {
	// extra returns helper-specific dependencies.
	extra() P

	// wantsChmod returns whether the source directory should be made writable.
	wantsChmod() bool
	// wantsWrite returns whether the source directory should be mounted writable.
	wantsWrite() bool
	// scriptEarly returns the helper-specific segment of cure script that goes
	// before the cd statement.
	scriptEarly() string
	// wantsDir returns the directory to enter before script.
	//
	// The zero value implies source directory if [PackageAttr.ScriptEarly] is
	// also empty. The special value helperInPlace omits the cd statement.
	wantsDir() (pathname string, create bool)
	// script returns the helper-specific segment of cure script.
	script(t Toolchain, name string) string
}

// PackageAttr holds build-system-agnostic attributes.
type PackageAttr struct {
	// Measure output if populated. Required by [THostNet].
	KnownChecksum *pkg.Checksum

	// Mount the source tree writable.
	Writable bool
	// Do not pass through [Toolchain.NewPatchedSource].
	Chmod bool
	// Unconditionally enter source directory early.
	EnterSource bool

	// Additional environment variables.
	Env []string
	// Runs before script emitted by [Helper]. Enters source if non-empty.
	ScriptEarly string

	// Passed to [Toolchain.NewPatchedSource].
	Patches []KV
	// Programs to make available in /bin/.
	Bin []string
	// Whether to replace /usr/bin/ with a symlink to /bin/.
	PopulateUsrBin bool

	// Hint for an early variant of toybox to be used when available.
	Early bool
	// Whether to omit output colour environment variables.
	NoColor bool
	// Exclude the LLVM toolchain.
	NoToolchain bool
	// Whether the resulting [pkg.Artifact] is exclusive.
	Exclusive bool
	// Whether to create [pkg.KindExecNet] instead of [pkg.KindExec].
	HostNet bool

	// Unregistered extras.
	Extra []pkg.Artifact
	// Passed through to [Toolchain.New], before source.
	Paths []pkg.ExecPath
}

// pa holds whether an [ArtifactH] is present.
type pa = map[ArtifactH]struct{}

// paPool holds addresses of pa.
var paPool = sync.Pool{New: func() any { return make(pa) }}

// paGet returns the address of a new pa.
func paGet() pa { return paPool.Get().(pa) }

// paPut returns a pa to paPool.
func paPut(pv pa) { clear(pv); paPool.Put(pv) }

// appendHandle recursively appends an [Artifact] named by its handle, and its
// runtime dependencies.
func (t Toolchain) appendHandle(a []pkg.Artifact, pv pa, p ArtifactH) []pkg.Artifact {
	if _, ok := pv[p]; ok {
		return a
	}
	pv[p] = struct{}{}

	meta, u := t.MustLoad(p)
	for _, d := range meta.Dependencies {
		a = t.appendHandle(a, pv, d)
	}
	return append(a, u)
}

// Append recursively appends multiple [Artifact] named by their handles, and
// their runtime dependencies.
func (t Toolchain) Append(a []pkg.Artifact, handles ...ArtifactH) []pkg.Artifact {
	pv := paGet()
	for _, p := range handles {
		a = t.appendHandle(a, pv, p)
	}
	paPut(pv)
	return a
}

// New constructs a [pkg.Artifact] via a build system helper.
func (t Toolchain) New(
	name, version string,
	source pkg.Artifact,
	attr *PackageAttr,
	helper Helper,
	extra ...ArtifactH,
) pkg.Artifact {
	if attr == nil {
		attr = new(PackageAttr)
	}

	if name == "" || version == "" {
		panic("name must be non-empty")
	}
	rn := name
	if version != Unversioned {
		rn = name + "-" + version
	}

	root := make([]pkg.Artifact, 0, 1<<3+len(attr.Extra)+len(extra))
	root = append(root, attr.Extra...)

	const lcMessages = "LC_MESSAGES=C.UTF-8"
	srn := rn
	env := slices.Clone(attr.Env)
	if !attr.NoColor {
		env = append(env,
			"CLICOLOR_FORCE=1",
			"FORCE_COLOR=1",
		)
	}

	var extraBoot []ArtifactH
	switch t.stage {
	case stageGentoo, stageEarly:
		srn += "-boot"
		root = append(root, NewEtc(true))
		if t.stage == stageEarly {
			extraBoot = append(extraBoot, _stage0Dist)
		} else if t.gentooStage3 == nil {
			panic(os.ErrInvalid)
		} else {
			env = append(env,
				"CC=clang",
				"CXX=clang++",
			)
		}
		env = fixupEnviron(env, []string{
			EnvTriple + "=" + t.triple(),
			lcMessages,
			"LDFLAGS=" + t.earlyLDFLAGS(true),
		}, "/system/bin",
			"/usr/bin",
		)

	case stageIntermediateGentoo, stageStdGentoo,
		stageIntermediate, Std, Stage3:
		if t.stage.isIntermediate() {
			srn += "-std"
		}

		if t.stage == Stage3 {
			srn += "-stage4"
		}

		toybox := _toybox
		if attr.Early {
			toybox = _toyboxEarly
		}

		base := _llvm
		if attr.NoToolchain {
			base = _musl
		}

		root = append(root, NewEtc(false))
		extraBoot = append(extraBoot,
			base,
			_mksh,
			toybox,
		)
		env = fixupEnviron(env, []string{
			EnvTriple + "=" + t.triple(),
			lcMessages,
		}, "/system/bin", "/bin")

	default:
		panic("unsupported toolchain " + strconv.Itoa(int(t.stage)))
	}

	wantsChmod, wantsWrite := helper.wantsChmod(), helper.wantsWrite()
	{
		pv := paGet()
		for _, p := range helper.extra() {
			root = t.appendHandle(root, pv, p)
		}
		for _, p := range extra {
			root = t.appendHandle(root, pv, p)
		}
		boot := t.S.New(t.stage - 1)
		for _, p := range extraBoot {
			root = boot.appendHandle(root, pv, p)
		}
		paPut(pv)
	}

	if t.stage == stageGentoo {
		root = append(root,
			t.gentooStage3,
			gentooOverlay{t.gentooStage3},
		)
	}

	var scriptEarly string
	dir, create := helper.wantsDir()
	helperScriptEarly := helper.scriptEarly()
	if attr.EnterSource ||
		dir == "" ||
		attr.ScriptEarly != "" ||
		helperScriptEarly != "" {
		scriptEarly += `
cd '/usr/src/` + name + `/'
`
	}
	scriptEarly += attr.ScriptEarly + helperScriptEarly
	if dir != "" && dir != helperInPlace {
		if create {
			scriptEarly += "\nmkdir -p " + dir
		}
		scriptEarly += "\ncd " + dir + "\n"
	} else if !attr.EnterSource && attr.ScriptEarly == "" && dir != "" {
		panic("cannot remain in root")
	}

	paths := attr.Paths
	if source != nil {
		paths = slices.Concat(attr.Paths, []pkg.ExecPath{
			pkg.Path(AbsUsrSrc.Append(
				name,
			), attr.Writable || wantsWrite, t.NewPatchedSource(
				rn, source, !attr.Chmod && !wantsChmod, attr.Patches...,
			)),
		})
	}

	bin := attr.Bin
	if attr.PopulateUsrBin {
		scriptEarly += `
chmod +w /usr/ /usr/bin/
rm -rf /usr/bin/
ln -s ../bin /usr/`
		bin = append(bin, "env")
	}

	if t.stage != stageGentoo && len(bin) > 0 {
		scriptEarly += "\nchmod +w /bin/" +
			"\n(set -o braceexpand && ln -sf ../system/bin/"
		if len(bin) > 1 {
			scriptEarly += "{'" + strings.Join(bin, "','") + "'}"
		} else {
			scriptEarly += "'" + bin[0] + "'"
		}
		scriptEarly += " /bin/)\n"
	}

	return pkg.NewExec(
		srn, t.arch, attr.KnownChecksum, pkg.ExecTimeoutMax,
		attr.HostNet,
		attr.Exclusive,
		fhs.AbsRoot, env,
		absCureScript,
		nil,

		slices.Concat([]pkg.ExecPath{pkg.Path(
			fhs.AbsRoot, true,
			root...,
		), pkg.Path(
			absCureScript, false,
			pkg.NewFile(scriptName, []byte(
				"#!/system/bin/sh\n"+
					"set -eu -o pipefail\n"+
					scriptEarly+helper.script(t, name),
			)),
		)}, paths)...,
	)
}

// GenericHelper directly passes build script segments for invocations of
// one-off build scripts too specific to fit into any other [Helper]
// implementation, but can still benefit from [Toolchain.NewPackage].
type GenericHelper struct {
	// Remain in current directory.
	InPlace bool

	// Concatenated for [Toolchain].
	Build, Check, Install string
}

var _ Helper = new(GenericHelper)

// extra is a noop.
func (*GenericHelper) extra() P { return nil }

// wantsChmod returns false.
func (*GenericHelper) wantsChmod() bool { return false }

// wantsWrite returns false.
func (*GenericHelper) wantsWrite() bool { return false }

// scriptEarly returns the zero value.
func (*GenericHelper) scriptEarly() string { return "" }

// wantsDir requests a new directory, or omits the cd statement if InPlace.
func (attr *GenericHelper) wantsDir() (string, bool) {
	if attr == nil || !attr.InPlace {
		return "/cure/", true
	}
	return helperInPlace, false
}

// script concatenates specified segments.
func (attr *GenericHelper) script(t Toolchain, _ string) string {
	if attr == nil {
		attr = new(GenericHelper)
	}

	script := attr.Build
	if t.opts&OptSkipCheck == 0 {
		script += attr.Check
	}
	script += attr.Install
	return script
}

// native contains natively-implemented and built-in azalea-based [Artifact].
// It is generally recommended to clone this instance for custom [Artifact]
// registrations.
var native S

// Native returns the global [S].
func Native() *S { return &native }

// parseTime is the duration of early parsing of built-in azalea expressions.
var parseTime time.Duration

// ParseTime returns the time taken by early parsing of built-in azalea expressions.
func ParseTime() time.Duration { return parseTime }

// nativeB is the backing directory of built-in azalea-based [Artifact]
// implementations.
//
//go:embed package
var nativeB embed.FS

func init() {
	sub, err := fs.Sub(nativeB, "package")
	if err != nil {
		panic(err)
	}
	t := time.Now()
	if err = native.RegisterFS(sub); err != nil {
		println(err.Error())
		os.Exit(1)
	}
	parseTime = time.Since(t)
}