package rosa

import (
	"runtime"
	"slices"
	"strconv"
	"strings"
	"sync"

	"hakurei.app/container/check"
	"hakurei.app/internal/pkg"
)

// llvmAttr holds the attributes that will be applied to a new [pkg.Artifact]
// containing a LLVM variant.
type llvmAttr struct {
	flags int

	// Concatenated with default environment for CMakeAttr.Env.
	env []string
	// Concatenated with generated entries for CMakeAttr.Cache.
	cmake [][2]string
	// Override CMakeAttr.Append.
	append []string
	// Concatenated with default dependencies for CMakeAttr.Extra.
	extra []pkg.Artifact
	// Concatenated with default fixup for CMakeAttr.Script.
	script string
	// Passed through to CMakeAttr.Prefix.
	prefix *check.Absolute

	// Patch name and body pairs.
	patches [][2]string
}

const (
	llvmProjectClang = 1 << iota
	llvmProjectLld

	llvmProjectAll = 1<<iota - 1

	llvmRuntimeCompilerRT = 1 << iota
	llvmRuntimeLibunwind
	llvmRuntimeLibc
	llvmRuntimeLibcxx
	llvmRuntimeLibcxxABI

	llvmAll        = 1<<iota - 1
	llvmRuntimeAll = llvmAll - (2 * llvmProjectAll) - 1
)

// llvmFlagName resolves a llvmAttr.flags project or runtime flag to its name.
func llvmFlagName(flag int) string {
	switch flag {
	case llvmProjectClang:
		return "clang"
	case llvmProjectLld:
		return "lld"

	case llvmRuntimeCompilerRT:
		return "compiler-rt"
	case llvmRuntimeLibunwind:
		return "libunwind"
	case llvmRuntimeLibc:
		return "libc"
	case llvmRuntimeLibcxx:
		return "libcxx"
	case llvmRuntimeLibcxxABI:
		return "libcxxabi"

	default:
		panic("invalid flag " + strconv.Itoa(flag))
	}
}

// newLLVMVariant returns a [pkg.Artifact] containing a LLVM variant.
func (t Toolchain) newLLVMVariant(variant string, attr *llvmAttr) pkg.Artifact {
	const (
		version  = "21.1.8"
		checksum = "8SUpqDkcgwOPsqHVtmf9kXfFeVmjVxl4LMn-qSE1AI_Xoeju-9HaoPNGtidyxyka"
	)
	if attr == nil {
		panic("LLVM attr must be non-nil")
	}

	var projects, runtimes []string
	for i := 1; i < llvmProjectAll; i <<= 1 {
		if attr.flags&i != 0 {
			projects = append(projects, llvmFlagName(i))
		}
	}
	for i := (llvmProjectAll + 1) << 1; i < llvmRuntimeAll; i <<= 1 {
		if attr.flags&i != 0 {
			runtimes = append(runtimes, llvmFlagName(i))
		}
	}

	var script, scriptEarly string

	cache := [][2]string{
		{"CMAKE_BUILD_TYPE", "Release"},

		{"LLVM_HOST_TRIPLE", `"${ROSA_TRIPLE}"`},
		{"LLVM_DEFAULT_TARGET_TRIPLE", `"${ROSA_TRIPLE}"`},
	}
	if len(projects) > 0 {
		cache = append(cache,
			[2]string{"LLVM_ENABLE_PROJECTS", `"${ROSA_LLVM_PROJECTS}"`})
	}
	if len(runtimes) > 0 {
		cache = append(cache,
			[2]string{"LLVM_ENABLE_RUNTIMES", `"${ROSA_LLVM_RUNTIMES}"`})
	}

	cmakeAppend := []string{"llvm"}
	if attr.append != nil {
		cmakeAppend = attr.append
	} else {
		cache = append(cache,
			[2]string{"LLVM_ENABLE_LIBCXX", "ON"},
			[2]string{"LLVM_USE_LINKER", "lld"},

			[2]string{"LLVM_INSTALL_BINUTILS_SYMLINKS", "ON"},
			[2]string{"LLVM_INSTALL_CCTOOLS_SYMLINKS", "ON"},
		)
	}

	if attr.flags&llvmProjectClang != 0 {
		cache = append(cache,
			[2]string{"CLANG_DEFAULT_LINKER", "lld"},
			[2]string{"CLANG_DEFAULT_CXX_STDLIB", "libc++"},
			[2]string{"CLANG_DEFAULT_RTLIB", "compiler-rt"},
			[2]string{"CLANG_DEFAULT_UNWINDLIB", "libunwind"},
		)
	}
	if attr.flags&llvmProjectLld != 0 {
		script += `
ln -s ld.lld /work/system/bin/ld
`
	}
	if attr.flags&llvmRuntimeCompilerRT != 0 {
		if attr.append == nil {
			cache = append(cache,
				[2]string{"COMPILER_RT_USE_LLVM_UNWINDER", "ON"})
		}
	}
	if attr.flags&llvmRuntimeLibunwind != 0 {
		cache = append(cache,
			[2]string{"LIBUNWIND_USE_COMPILER_RT", "ON"})
	}
	if attr.flags&llvmRuntimeLibcxx != 0 {
		cache = append(cache,
			[2]string{"LIBCXX_HAS_MUSL_LIBC", "ON"},
			[2]string{"LIBCXX_USE_COMPILER_RT", "ON"},
		)

		if t > toolchainStage3 {
			// libcxxabi fails to compile if c++ headers not prefixed in /usr
			// is found by the compiler, and doing this is easier than
			// overriding CXXFLAGS; not using mv here to avoid chown failures
			scriptEarly += `
cp -r /system/include /usr/include && rm -rf /system/include
`
		}
	}
	if attr.flags&llvmRuntimeLibcxxABI != 0 {
		cache = append(cache,
			[2]string{"LIBCXXABI_USE_COMPILER_RT", "ON"},
			[2]string{"LIBCXXABI_USE_LLVM_UNWINDER", "ON"},
		)
	}

	return t.NewViaCMake("llvm", version, variant, t.NewPatchedSource(
		"llvmorg", version, pkg.NewHTTPGetTar(
			nil, "https://github.com/llvm/llvm-project/archive/refs/tags/"+
				"llvmorg-"+version+".tar.gz",
			mustDecode(checksum),
			pkg.TarGzip,
		), true, attr.patches...,
	), &CMakeAttr{
		Cache:  slices.Concat(cache, attr.cmake),
		Append: cmakeAppend,
		Extra: stage3Concat(t, attr.extra,
			t.Load(Libffi),
			t.Load(Python),
			t.Load(Perl),
			t.Load(Diffutils),
			t.Load(Bash),
			t.Load(Coreutils),

			t.Load(KernelHeaders),
		),
		Prefix: attr.prefix,

		Env: slices.Concat([]string{
			"ROSA_LLVM_PROJECTS=" + strings.Join(projects, ";"),
			"ROSA_LLVM_RUNTIMES=" + strings.Join(runtimes, ";"),
		}, attr.env),
		ScriptEarly: scriptEarly, Script: script + attr.script,

		Exclusive: true,
	})
}

// newLLVM returns LLVM toolchain across multiple [pkg.Artifact].
func (t Toolchain) newLLVM() (musl, compilerRT, runtimes, clang pkg.Artifact) {
	var target string
	switch runtime.GOARCH {
	case "386", "amd64":
		target = "X86"
	case "arm64":
		target = "AArch64"

	default:
		panic("unsupported target " + runtime.GOARCH)
	}

	minimalDeps := [][2]string{
		{"LLVM_ENABLE_ZLIB", "OFF"},
		{"LLVM_ENABLE_ZSTD", "OFF"},
		{"LLVM_ENABLE_LIBXML2", "OFF"},
	}

	compilerRT = t.newLLVMVariant("compiler-rt", &llvmAttr{
		env: stage3ExclConcat(t, []string{},
			"LDFLAGS="+earlyLDFLAGS(false),
		),
		cmake: [][2]string{
			// libc++ not yet available
			{"CMAKE_CXX_COMPILER_TARGET", ""},

			{"COMPILER_RT_BUILD_BUILTINS", "ON"},
			{"COMPILER_RT_DEFAULT_TARGET_ONLY", "ON"},
			{"LLVM_ENABLE_PER_TARGET_RUNTIME_DIR", "ON"},

			// does not work without libunwind
			{"COMPILER_RT_BUILD_CTX_PROFILE", "OFF"},
			{"COMPILER_RT_BUILD_LIBFUZZER", "OFF"},
			{"COMPILER_RT_BUILD_MEMPROF", "OFF"},
			{"COMPILER_RT_BUILD_PROFILE", "OFF"},
			{"COMPILER_RT_BUILD_SANITIZERS", "OFF"},
			{"COMPILER_RT_BUILD_XRAY", "OFF"},
		},
		append: []string{"compiler-rt"},
		extra: []pkg.Artifact{t.NewMusl(&MuslAttr{
			Headers: true,
			Env: []string{
				"CC=clang",
			},
		})},
		script: `
mkdir -p "${ROSA_INSTALL_PREFIX}/lib/clang/21/lib/"
ln -s \
	"../../../${ROSA_TRIPLE}" \
	"${ROSA_INSTALL_PREFIX}/lib/clang/21/lib/"

ln -s \
	"clang_rt.crtbegin-` + linuxArch() + `.o" \
	"${ROSA_INSTALL_PREFIX}/lib/${ROSA_TRIPLE}/crtbeginS.o"
ln -s \
	"clang_rt.crtend-` + linuxArch() + `.o" \
	"${ROSA_INSTALL_PREFIX}/lib/${ROSA_TRIPLE}/crtendS.o"
`,
	})

	musl = t.NewMusl(&MuslAttr{
		Extra: []pkg.Artifact{compilerRT},
		Env: stage3ExclConcat(t, []string{
			"CC=clang",
			"LIBCC=/system/lib/clang/21/lib/" +
				triplet() + "/libclang_rt.builtins.a",
			"AR=ar",
			"RANLIB=ranlib",
		},
			"LDFLAGS="+earlyLDFLAGS(false),
		),
	})

	runtimes = t.newLLVMVariant("runtimes", &llvmAttr{
		env: stage3ExclConcat(t, []string{},
			"LDFLAGS="+earlyLDFLAGS(false),
		),
		flags: llvmRuntimeLibunwind | llvmRuntimeLibcxx | llvmRuntimeLibcxxABI,
		cmake: slices.Concat([][2]string{
			// libc++ not yet available
			{"CMAKE_CXX_COMPILER_WORKS", "ON"},

			{"LIBCXX_HAS_ATOMIC_LIB", "OFF"},
			{"LIBCXXABI_HAS_CXA_THREAD_ATEXIT_IMPL", "OFF"},
		}, minimalDeps),
		append: []string{"runtimes"},
		extra: []pkg.Artifact{
			compilerRT,
			musl,
		},
	})

	clang = t.newLLVMVariant("clang", &llvmAttr{
		flags: llvmProjectClang | llvmProjectLld,
		env: stage3ExclConcat(t, []string{},
			"CFLAGS="+earlyCFLAGS,
			"CXXFLAGS="+earlyCXXFLAGS(),
			"LDFLAGS="+earlyLDFLAGS(false),
		),
		cmake: slices.Concat([][2]string{
			{"LLVM_TARGETS_TO_BUILD", target},
			{"CMAKE_CROSSCOMPILING", "OFF"},
			{"CXX_SUPPORTS_CUSTOM_LINKER", "ON"},
		}, minimalDeps),
		extra: []pkg.Artifact{
			musl,
			compilerRT,
			runtimes,
			t.Load(Git),
		},
		script: `
ln -s clang /work/system/bin/cc
ln -s clang++ /work/system/bin/c++

ninja check-all
`,

		patches: [][2]string{
			{"xfail-broken-tests", `diff --git a/clang/test/Modules/timestamps.c b/clang/test/Modules/timestamps.c
index 50fdce630255..4b4465a75617 100644
--- a/clang/test/Modules/timestamps.c
+++ b/clang/test/Modules/timestamps.c
@@ -1,3 +1,5 @@
+// XFAIL: target={{.*-rosa-linux-musl}}
+
 /// Verify timestamps that gets embedded in the module
 #include <c-header.h>

`},

			{"path-system-include", `diff --git a/clang/lib/Driver/ToolChains/Linux.cpp b/clang/lib/Driver/ToolChains/Linux.cpp
index cdbf21fb9026..dd052858700d 100644
--- a/clang/lib/Driver/ToolChains/Linux.cpp
+++ b/clang/lib/Driver/ToolChains/Linux.cpp
@@ -773,6 +773,12 @@ void Linux::AddClangSystemIncludeArgs(const ArgList &DriverArgs,
     addExternCSystemInclude(
         DriverArgs, CC1Args,
         concat(SysRoot, "/usr/include", MultiarchIncludeDir));
+  if (!MultiarchIncludeDir.empty() &&
+      D.getVFS().exists(concat(SysRoot, "/system/include", MultiarchIncludeDir)))
+    addExternCSystemInclude(
+        DriverArgs, CC1Args,
+        concat(SysRoot, "/system/include", MultiarchIncludeDir));
+

   if (getTriple().getOS() == llvm::Triple::RTEMS)
     return;
@@ -783,6 +789,7 @@ void Linux::AddClangSystemIncludeArgs(const ArgList &DriverArgs,
   addExternCSystemInclude(DriverArgs, CC1Args, concat(SysRoot, "/include"));

   addExternCSystemInclude(DriverArgs, CC1Args, concat(SysRoot, "/usr/include"));
+  addExternCSystemInclude(DriverArgs, CC1Args, concat(SysRoot, "/system/include"));

   if (!DriverArgs.hasArg(options::OPT_nobuiltininc) && getTriple().isMusl())
     addSystemInclude(DriverArgs, CC1Args, ResourceDirInclude);
`},

			{"path-system-libraries", `diff --git a/clang/lib/Driver/ToolChains/CommonArgs.cpp b/clang/lib/Driver/ToolChains/CommonArgs.cpp
index 8d3775de9be5..1e126e2d6f24 100644
--- a/clang/lib/Driver/ToolChains/CommonArgs.cpp
+++ b/clang/lib/Driver/ToolChains/CommonArgs.cpp
@@ -463,6 +463,15 @@ void tools::AddLinkerInputs(const ToolChain &TC, const InputInfoList &Inputs,
   if (!TC.isCrossCompiling())
     addDirectoryList(Args, CmdArgs, "-L", "LIBRARY_PATH");

+  const std::string RosaSuffix = "-rosa-linux-musl";
+  if (TC.getTripleString().size() > RosaSuffix.size() &&
+      std::equal(RosaSuffix.rbegin(), RosaSuffix.rend(), TC.getTripleString().rbegin())) {
+    CmdArgs.push_back("-rpath");
+    CmdArgs.push_back("/system/lib");
+    CmdArgs.push_back("-rpath");
+    CmdArgs.push_back(("/system/lib/" + TC.getTripleString()).c_str());
+  }
+
   for (const auto &II : Inputs) {
     // If the current tool chain refers to an OpenMP offloading host, we
     // should ignore inputs that refer to OpenMP offloading devices -
diff --git a/clang/lib/Driver/ToolChains/Linux.cpp b/clang/lib/Driver/ToolChains/Linux.cpp
index 8ac8d4eb9181..795995bb53cb 100644
--- a/clang/lib/Driver/ToolChains/Linux.cpp
+++ b/clang/lib/Driver/ToolChains/Linux.cpp
@@ -324,6 +324,7 @@ Linux::Linux(const Driver &D, const llvm::Triple &Triple, const ArgList &Args)
     Generic_GCC::AddMultilibPaths(D, SysRoot, "libo32", MultiarchTriple, Paths);
     addPathIfExists(D, concat(SysRoot, "/libo32"), Paths);
     addPathIfExists(D, concat(SysRoot, "/usr/libo32"), Paths);
+    addPathIfExists(D, concat(SysRoot, "/system/libo32"), Paths);
   }
   Generic_GCC::AddMultilibPaths(D, SysRoot, OSLibDir, MultiarchTriple, Paths);

@@ -343,16 +344,20 @@ Linux::Linux(const Driver &D, const llvm::Triple &Triple, const ArgList &Args)

   addPathIfExists(D, concat(SysRoot, "/usr/lib", MultiarchTriple), Paths);
   addPathIfExists(D, concat(SysRoot, "/usr", OSLibDir), Paths);
+  addPathIfExists(D, concat(SysRoot, "/system/lib", MultiarchTriple), Paths);
+  addPathIfExists(D, concat(SysRoot, "/system", OSLibDir), Paths);
   if (IsRISCV) {
     StringRef ABIName = tools::riscv::getRISCVABI(Args, Triple);
     addPathIfExists(D, concat(SysRoot, "/", OSLibDir, ABIName), Paths);
     addPathIfExists(D, concat(SysRoot, "/usr", OSLibDir, ABIName), Paths);
+    addPathIfExists(D, concat(SysRoot, "/system", OSLibDir, ABIName), Paths);
   }

   Generic_GCC::AddMultiarchPaths(D, SysRoot, OSLibDir, Paths);

   addPathIfExists(D, concat(SysRoot, "/lib"), Paths);
   addPathIfExists(D, concat(SysRoot, "/usr/lib"), Paths);
+  addPathIfExists(D, concat(SysRoot, "/system/lib"), Paths);
 }

 ToolChain::RuntimeLibType Linux::GetDefaultRuntimeLibType() const {
@@ -457,6 +462,11 @@ std::string Linux::getDynamicLinker(const ArgList &Args) const {
     return Triple.isArch64Bit() ? "/system/bin/linker64" : "/system/bin/linker";
   }
   if (Triple.isMusl()) {
+    const std::string RosaSuffix = "-rosa-linux-musl";
+    if (Triple.str().size() > RosaSuffix.size() &&
+        std::equal(RosaSuffix.rbegin(), RosaSuffix.rend(), Triple.str().rbegin()))
+      return "/system/bin/linker";
+
     std::string ArchName;
     bool IsArm = false;

diff --git a/clang/tools/clang-installapi/Options.cpp b/clang/tools/clang-installapi/Options.cpp
index 64324a3f8b01..15ce70b68217 100644
--- a/clang/tools/clang-installapi/Options.cpp
+++ b/clang/tools/clang-installapi/Options.cpp
@@ -515,7 +515,7 @@ bool Options::processFrontendOptions(InputArgList &Args) {
     FEOpts.FwkPaths = std::move(FrameworkPaths);

   // Add default framework/library paths.
-  PathSeq DefaultLibraryPaths = {"/usr/lib", "/usr/local/lib"};
+  PathSeq DefaultLibraryPaths = {"/usr/lib", "/system/lib", "/usr/local/lib"};
   PathSeq DefaultFrameworkPaths = {"/Library/Frameworks",
                                    "/System/Library/Frameworks"};

`},
		},
	})

	return
}

var (
	// llvm stores the result of Toolchain.newLLVM.
	llvm [_toolchainEnd][4]pkg.Artifact
	// llvmOnce is for lazy initialisation of llvm.
	llvmOnce [_toolchainEnd]sync.Once
)

// NewLLVM returns LLVM toolchain across multiple [pkg.Artifact].
func (t Toolchain) NewLLVM() (musl, compilerRT, runtimes, clang pkg.Artifact) {
	llvmOnce[t].Do(func() {
		llvm[t][0], llvm[t][1], llvm[t][2], llvm[t][3] = t.newLLVM()
	})
	return llvm[t][0], llvm[t][1], llvm[t][2], llvm[t][3]
}