package dockerfile

import (
	"context"
	"fmt"
	"os"
	"path"
	"path/filepath"
	"slices"
	"strings"

	"github.com/replicate/cog/pkg/config"
	"github.com/replicate/cog/pkg/docker/command"
	"github.com/replicate/cog/pkg/dockercontext"
	"github.com/replicate/cog/pkg/registry"
	"github.com/replicate/cog/pkg/requirements"
	"github.com/replicate/cog/pkg/util/console"
	"github.com/replicate/cog/pkg/util/version"
	"github.com/replicate/cog/pkg/weights"
	"github.com/replicate/cog/pkg/wheels"
)

const DockerignoreHeader = `# generated by replicate/cog
__pycache__
*.pyc
*.pyo
*.pyd
.Python
env
pip-log.txt
pip-delete-this-directory.txt
.tox
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.log
.git
.mypy_cache
.pytest_cache
.hypothesis
`
const LDConfigCacheBuildCommand = "RUN find / -type f -name \"*python*.so\" -printf \"%h\\n\" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig"
const StripDebugSymbolsCommand = "find / -type f -name \"*python*.so\" -not -name \"*cpython*.so\" -exec strip -S {} \\;"
const CFlags = "ENV CFLAGS=\"-O3 -funroll-loops -fno-strict-aliasing -flto -S\""
const UVVersion = "0.9.26"
const uvCacheMount = "--mount=type=cache,target=/root/.cache/uv"
const uvPip = "uv pip"
const PrecompilePythonCommand = "RUN find / -type f -name \"*.py[co]\" -delete && find / -type f -name \"*.py\" -exec touch -t 197001010000 {} \\; && find / -type f -name \"*.py\" -printf \"%h\\n\" | sort -u | /usr/bin/python3 -m compileall --invalidation-mode timestamp -o 2 -j 0"
const STANDARD_GENERATOR_NAME = "STANDARD_GENERATOR"

type StandardGenerator struct {
	Config         *config.Config
	Dir            string
	ConfigFilename string // Base filename like "cog.yaml" or "my-config.yaml"

	// these are here to make this type testable
	GOOS   string
	GOARCH string

	useCudaBaseImage bool
	useCogBaseImage  *bool
	strip            bool
	precompile       bool

	// absolute path to tmpDir, a directory that will be cleaned up
	tmpDir string
	// tmpDir relative to Dir
	relativeTmpDir string

	fileWalker weights.FileWalker

	modelDirs  []string
	modelFiles []string

	pythonRequirementsContents string
	command                    command.Command
	client                     registry.Client
	requiresCog                bool

	// Optional overrides for wheel configs (used by tests for deterministic output).
	// When nil, auto-detection is used (env var → dist/ → PyPI).
	cogWheelConfig    *wheels.WheelConfig
	cogletWheelConfig *wheels.WheelConfig

	// Resolved wheel configs — set once by resolveCogWheelConfigs() and shared
	// between filterManagedPackages() (for warnings) and installCog() (for install).
	resolvedCogConfig    *wheels.WheelConfig
	resolvedCogletConfig *wheels.WheelConfig
}

func NewStandardGenerator(config *config.Config, dir string, configFilename string, command command.Command, client registry.Client, requiresCog bool) (*StandardGenerator, error) {
	tmpDir, err := dockercontext.BuildTempDir(dir)
	if err != nil {
		return nil, err
	}
	// tmpDir, but without dir prefix. This is the path used in the Dockerfile.
	relativeTmpDir, err := filepath.Rel(dir, tmpDir)
	if err != nil {
		return nil, err
	}

	// Default to "cog.yaml" if not specified
	if configFilename == "" {
		configFilename = "cog.yaml"
	}

	return &StandardGenerator{
		Config:         config,
		Dir:            dir,
		ConfigFilename: configFilename,
		// Docker build target is always linux/amd64 (see pkg/docker/buildkit.go).
		// These must match the container platform, not the host.
		GOOS:             "linux",
		GOARCH:           "amd64",
		tmpDir:           tmpDir,
		relativeTmpDir:   relativeTmpDir,
		fileWalker:       filepath.Walk,
		useCudaBaseImage: true,
		useCogBaseImage:  nil,
		strip:            false,
		precompile:       false,
		command:          command,
		client:           client,
		requiresCog:      requiresCog,
	}, nil
}

func (g *StandardGenerator) SetUseCudaBaseImage(argumentValue string) {
	// "false" -> false, "true" -> true, "auto" -> true, "asdf" -> true
	g.useCudaBaseImage = argumentValue != "false"
}

func (g *StandardGenerator) SetUseCogBaseImage(useCogBaseImage bool) {
	g.useCogBaseImage = new(bool)
	*g.useCogBaseImage = useCogBaseImage
}

func (g *StandardGenerator) SetUseCogBaseImagePtr(useCogBaseImage *bool) {
	g.useCogBaseImage = useCogBaseImage
}

func (g *StandardGenerator) IsUsingCogBaseImage() bool {
	useCogBaseImage := g.useCogBaseImage
	if useCogBaseImage != nil {
		return *useCogBaseImage
	}
	return true
}

func (g *StandardGenerator) SetStrip(strip bool) {
	g.strip = strip
}

func (g *StandardGenerator) SetPrecompile(precompile bool) {
	g.precompile = precompile
}

func (g *StandardGenerator) GenerateInitialSteps(ctx context.Context) (string, error) {
	baseImage, err := g.BaseImage(ctx)
	if err != nil {
		return "", err
	}
	installPython, err := g.installPython()
	if err != nil {
		return "", err
	}
	aptInstalls, err := g.aptInstalls()
	if err != nil {
		return "", err
	}
	envs, err := g.envVars()
	if err != nil {
		return "", err
	}
	runCommands, err := g.runCommands()
	if err != nil {
		return "", err
	}
	pipInstalls, err := g.pipInstalls()
	if err != nil {
		return "", err
	}
	installCog, err := g.installCog()
	if err != nil {
		return "", err
	}
	installCACert, err := g.installCACert()
	if err != nil {
		return "", err
	}

	if g.IsUsingCogBaseImage() {
		steps := []string{
			"#syntax=docker/dockerfile:1.4",
			"FROM " + baseImage,
			installCACert, // First! Before any network requests (apt, pip, etc.)
			envs,
			aptInstalls,
			g.installUV(),
		}
		if installCog != "" {
			steps = append(steps, installCog)
		}
		steps = append(steps, pipInstalls)
		if g.precompile {
			steps = append(steps, PrecompilePythonCommand)
		}
		steps = append(steps, runCommands)

		return joinStringsWithoutLineSpace(steps), nil
	}

	// For the CUDA path, uv is installed inside installPython (after the apt step).
	// For all other paths (python:X-slim), install uv after apt.
	uvInstall := ""
	if installPython == "" {
		uvInstall = g.installUV()
	}
	steps := []string{
		"#syntax=docker/dockerfile:1.4",
		"FROM " + baseImage,
		g.preamble(),
		installCACert, // Early! Before tini (uses curl), apt, pip, etc.
		g.installTini(),
		envs,
		aptInstalls,
		uvInstall,
		installPython,
		pipInstalls,
		installCog,
	}
	if g.precompile {
		steps = append(steps, PrecompilePythonCommand)
	}
	steps = append(steps, LDConfigCacheBuildCommand, runCommands)

	return joinStringsWithoutLineSpace(steps), nil
}

func (g *StandardGenerator) GenerateModelBase(ctx context.Context) (string, error) {
	initialSteps, err := g.GenerateInitialSteps(ctx)
	if err != nil {
		return "", err
	}
	steps := []string{
		initialSteps,
		`WORKDIR /src`,
		`EXPOSE 5000`,
	}
	steps = append(steps, g.cogEnvVars()...)
	steps = append(steps, `CMD ["python", "-m", "cog.server.http"]`)
	return strings.Join(steps, "\n"), nil
}

// GenerateDockerfileWithoutSeparateWeights generates a Dockerfile that doesn't write model weights to a separate layer.
func (g *StandardGenerator) GenerateDockerfileWithoutSeparateWeights(ctx context.Context) (string, error) {
	base, err := g.GenerateModelBase(ctx)
	if err != nil {
		return "", err
	}
	bases := []string{
		base,
		`COPY . /src`,
	}
	if m := g.cpCogYaml(); m != "" {
		bases = append(bases, m)
	}
	return joinStringsWithoutLineSpace(bases), nil
}

// GenerateModelBaseWithSeparateWeights creates the Dockerfile and .dockerignore file contents for model weights
// It returns four values:
// - weightsBase: The base image used for Dockerfile generation for model weights.
// - dockerfile: A string that represents the Dockerfile content generated by the function.
// - dockerignoreContents: A string that represents the .dockerignore content.
// - err: An error object if an error occurred during Dockerfile generation; otherwise nil.
func (g *StandardGenerator) GenerateModelBaseWithSeparateWeights(ctx context.Context, imageName string) (weightsBase string, dockerfile string, dockerignoreContents string, err error) {
	weightsBase, g.modelDirs, g.modelFiles, err = g.generateForWeights()
	if err != nil {
		return "", "", "", fmt.Errorf("Failed to generate Dockerfile for model weights files: %w", err)
	}
	initialSteps, err := g.GenerateInitialSteps(ctx)
	if err != nil {
		return "", "", "", err
	}

	// Inject weights base image into initial steps so we can COPY from it
	base := []string{}
	initialStepsLines := strings.Split(initialSteps, "\n")
	for i, line := range initialStepsLines {
		if strings.HasPrefix(line, "FROM ") {
			base = append(base, fmt.Sprintf("FROM %s AS %s", imageName+"-weights", "weights"))
			base = append(base, initialStepsLines[i:]...)
			break
		} else {
			base = append(base, line)
		}
	}

	for _, p := range append(g.modelDirs, g.modelFiles...) {
		base = append(base, "COPY --from=weights --link "+path.Join("/src", p)+" "+path.Join("/src", p))
	}

	base = append(base,
		`WORKDIR /src`,
		`EXPOSE 5000`,
	)
	base = append(base, g.cogEnvVars()...)
	base = append(base,
		`CMD ["python", "-m", "cog.server.http"]`,
		`COPY . /src`,
	)
	if m := g.cpCogYaml(); m != "" {
		base = append(base, m)
	}

	dockerignoreContents = makeDockerignoreForWeights(g.modelDirs, g.modelFiles)
	return weightsBase, joinStringsWithoutLineSpace(base), dockerignoreContents, nil
}

// cogEnvVars returns ENV lines that pass cog.yaml config to the runtime
// so the container doesn't need to parse cog.yaml at startup.
func (g *StandardGenerator) cogEnvVars() []string {
	var envs []string
	if g.Config.Predict != "" {
		envs = append(envs, fmt.Sprintf(`ENV COG_PREDICT_TYPE_STUB="%s"`, g.Config.Predict))
	}
	if g.Config.Train != "" {
		envs = append(envs, fmt.Sprintf(`ENV COG_TRAIN_TYPE_STUB="%s"`, g.Config.Train))
	}
	if g.Config.Concurrency != nil && g.Config.Concurrency.Max > 0 {
		envs = append(envs, fmt.Sprintf(`ENV COG_MAX_CONCURRENCY=%d`, g.Config.Concurrency.Max))
	}
	return envs
}

func (g *StandardGenerator) cpCogYaml() string {
	if g.ConfigFilename == "" || g.ConfigFilename == "cog.yaml" {
		return ""
	}
	// Absolute filename doesn't work anyway, so it's always relative
	return fmt.Sprintf("RUN cp %s /src/cog.yaml", filepath.Join("/src", g.ConfigFilename))
}

func (g *StandardGenerator) generateForWeights() (string, []string, []string, error) {
	modelDirs, modelFiles, err := weights.FindWeights(g.fileWalker)
	if err != nil {
		return "", nil, nil, err
	}
	// generate dockerfile to store these model weights files
	var dockerfileContents strings.Builder
	dockerfileContents.WriteString(`#syntax=docker/dockerfile:1.4
FROM scratch
`)
	for _, p := range append(modelDirs, modelFiles...) {
		fmt.Fprintf(&dockerfileContents, "\nCOPY %s %s", p, path.Join("/src", p))
	}

	return dockerfileContents.String(), modelDirs, modelFiles, nil
}

func makeDockerignoreForWeights(dirs, files []string) string {
	var contents strings.Builder
	for _, p := range dirs {
		fmt.Fprintf(&contents, "%[1]s\n%[1]s/**/*\n", p)
	}
	for _, p := range files {
		fmt.Fprintf(&contents, "%[1]s\n", p)
	}
	return DockerignoreHeader + contents.String()
}

func (g *StandardGenerator) Cleanup() error {
	if err := os.RemoveAll(g.tmpDir); err != nil {
		return fmt.Errorf("Failed to clean up %s: %w", g.tmpDir, err)
	}
	return nil
}

func (g *StandardGenerator) BaseImage(ctx context.Context) (string, error) {
	if g.IsUsingCogBaseImage() {
		baseImage, err := g.determineBaseImageName(ctx)
		if err == nil || g.useCogBaseImage != nil {
			return baseImage, err
		}
		console.Warnf("Could not find a suitable base image, continuing without base image support (%v).", err)
		if g.useCogBaseImage == nil {
			g.useCogBaseImage = new(bool)
			*g.useCogBaseImage = false
		}
	}

	if g.Config.Build.GPU && g.useCudaBaseImage {
		return g.Config.CUDABaseImageTag()
	}
	return "python:" + g.Config.Build.PythonVersion + "-slim", nil
}

func (g *StandardGenerator) Name() string {
	return STANDARD_GENERATOR_NAME
}

func (g *StandardGenerator) BuildDir() (string, error) {
	return dockercontext.StandardBuildDirectory, nil
}

func (g *StandardGenerator) BuildContexts() (map[string]string, error) {
	return map[string]string{}, nil
}

func (g *StandardGenerator) preamble() string {
	return `ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin
ENV NVIDIA_DRIVER_CAPABILITIES=all`
}

func (g *StandardGenerator) installTini() string {
	// Install tini as the image entrypoint to provide signal handling and process
	// reaping appropriate for PID 1.
	//
	// N.B. If you remove/change this, consider removing/changing the `has_init`
	// image label applied in image/build.go.
	lines := []string{
		`RUN --mount=type=cache,target=/var/cache/apt,sharing=locked set -eux; \
apt-get update -qq && \
apt-get install -qqy --no-install-recommends curl; \
rm -rf /var/lib/apt/lists/*; \
TINI_VERSION=v0.19.0; \
TINI_ARCH="$(dpkg --print-architecture)"; \
curl -sSL -o /sbin/tini "https://github.com/krallin/tini/releases/download/${TINI_VERSION}/tini-${TINI_ARCH}"; \
chmod +x /sbin/tini`,
		`ENTRYPOINT ["/sbin/tini", "--"]`,
	}
	return strings.Join(lines, "\n")
}

func (g *StandardGenerator) aptInstalls() (string, error) {
	packages := g.Config.Build.SystemPackages
	if len(packages) == 0 {
		return "", nil
	}

	if g.IsUsingCogBaseImage() {
		// Filter out packages that are already in the base image
		packages = slices.DeleteFunc(slices.Clone(packages), func(pkg string) bool {
			return slices.Contains(baseImageSystemPackages, pkg)
		})
	}

	return "RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy " +
		strings.Join(packages, " ") +
		" && rm -rf /var/lib/apt/lists/*", nil
}

func (g *StandardGenerator) installPython() (string, error) {
	if g.Config.Build.GPU && g.useCudaBaseImage && !g.IsUsingCogBaseImage() {
		return g.installPythonCUDA()
	}
	return "", nil
}

func (g *StandardGenerator) installUV() string {
	return `COPY --from=ghcr.io/astral-sh/uv:` + UVVersion + ` /uv /uvx /usr/local/bin/
ENV UV_SYSTEM_PYTHON=true`
}

func (g *StandardGenerator) installPythonCUDA() (string, error) {
	// TODO: check that python version is valid

	py := g.Config.Build.PythonVersion
	return `RUN --mount=type=cache,target=/var/cache/apt,sharing=locked apt-get update -qq && apt-get install -qqy --no-install-recommends \
	wget \
	curl \
	xz-utils \
	git \
	ca-certificates \
	&& rm -rf /var/lib/apt/lists/*
` + g.installUV() + "\n" + fmt.Sprintf(`RUN uv python install %s && \
	ln -sf $(uv python find %s) /usr/bin/python3
ENV UV_PYTHON=%s
ENV PATH="/usr/local/bin:$PATH"`, py, py, py), nil
}

// resolveCogWheelConfigs resolves and caches the cog and coglet wheel configs.
// It is idempotent — subsequent calls are no-ops. Must be called before
// filterManagedPackages() and installCog().
//
// Precedence for cog SDK:
//  1. Test override (cogWheelConfig field)
//  2. COG_SDK_WHEEL env var
//  3. build.sdk_version in cog.yaml
//  4. Auto-detect dist/ (dev builds only)
//  5. Latest PyPI
//
// Precedence for coglet:
//  1. Test override (cogletWheelConfig field)
//  2. COGLET_WHEEL env var
//  3. Auto-detect dist/ (dev builds only)
//  4. Latest PyPI
func (g *StandardGenerator) resolveCogWheelConfigs() error {
	if g.resolvedCogConfig != nil {
		return nil // already resolved
	}

	var err error

	// Resolve cog SDK
	if g.cogWheelConfig != nil {
		g.resolvedCogConfig = g.cogWheelConfig
	} else if envVal := os.Getenv(wheels.CogSDKWheelEnvVar); envVal != "" {
		g.resolvedCogConfig, err = wheels.GetCogWheelConfig()
		if err != nil {
			return err
		}
	} else if g.Config.Build != nil && g.Config.Build.SDKVersion != "" {
		g.resolvedCogConfig = &wheels.WheelConfig{
			Source:  wheels.WheelSourcePyPI,
			Version: g.Config.Build.SDKVersion,
		}
	} else {
		g.resolvedCogConfig, err = wheels.GetCogWheelConfig()
		if err != nil {
			return err
		}
	}

	// Validate: refuse versions older than the minimum supported SDK
	if err := wheels.ValidateSDKVersion(g.resolvedCogConfig, "cog"); err != nil {
		return err
	}

	// Resolve coglet
	if g.cogletWheelConfig != nil {
		g.resolvedCogletConfig = g.cogletWheelConfig
	} else {
		g.resolvedCogletConfig, err = wheels.GetCogletWheelConfig(g.GOARCH)
		if err != nil {
			return err
		}
	}

	return nil
}

// cogletMinSDKVersion is the minimum SDK version that supports coglet.
// Older SDKs use the built-in Python HTTP server and are incompatible with coglet.
const cogletMinSDKVersion = "0.17.0"

// isLegacySDKVersion returns true if the resolved cog SDK version is explicitly
// pinned below the minimum version that supports coglet. Returns false for
// unpinned, non-PyPI, or unparseable versions (assume modern).
func (g *StandardGenerator) isLegacySDKVersion() bool {
	cfg := g.resolvedCogConfig
	if cfg == nil || cfg.Source != wheels.WheelSourcePyPI || cfg.Version == "" {
		return false
	}
	base := cfg.Version
	if m := wheels.BaseVersionRe.FindString(base); m != "" {
		base = m
	}
	ver, err := version.NewVersion(base)
	if err != nil {
		return false
	}
	return !ver.GreaterOrEqual(version.MustVersion(cogletMinSDKVersion))
}

func (g *StandardGenerator) installCog() (string, error) {
	// Do not install Cog in base images
	if !g.requiresCog {
		return "", nil
	}

	if err := g.resolveCogWheelConfigs(); err != nil {
		return "", err
	}
	wheelConfig := g.resolvedCogConfig

	// Determine if we need --pre flag (pre-release SDK implies pre-release coglet too)
	sdkIsPreRelease := wheelConfig.Source == wheels.WheelSourcePyPI && wheels.IsPreRelease(wheelConfig.Version)

	// Only install coglet explicitly when there's a specific source:
	//   - COGLET_WHEEL env var (explicit override)
	//   - Local wheel from dist/ (dev/CI auto-detect)
	//   - PyPI with pinned version (e.g. COGLET_WHEEL=pypi:0.17.0)
	// Otherwise, let the SDK's own dependency handle it — cog >= 0.17.0 declares
	// coglet as a hard dependency, older versions don't install it.
	//
	// Never install coglet when the SDK is explicitly pinned to < 0.17.0 — those
	// versions use the built-in Python HTTP server and are incompatible with coglet.
	var installLines string
	cogletConfig := g.resolvedCogletConfig
	explicitCoglet := cogletConfig != nil &&
		(cogletConfig.Source == wheels.WheelSourceFile ||
			cogletConfig.Source == wheels.WheelSourceURL ||
			(cogletConfig.Source == wheels.WheelSourcePyPI && cogletConfig.Version != ""))
	if explicitCoglet && g.isLegacySDKVersion() {
		console.Info("Skipping coglet install for legacy SDK")
		explicitCoglet = false
	}
	if explicitCoglet {
		switch cogletConfig.Source {
		case wheels.WheelSourcePyPI:
			console.Infof("Using coglet from PyPI: %s", cogletConfig.PyPIPackageURL("coglet"))
		case wheels.WheelSourceURL:
			console.Infof("Using coglet wheel from URL: %s", cogletConfig.URL)
		case wheels.WheelSourceFile:
			console.Debugf("Using local coglet wheel: %s", cogletConfig.Path)
		}

		cogletIsPreRelease := sdkIsPreRelease ||
			(cogletConfig.Source == wheels.WheelSourcePyPI && wheels.IsPreRelease(cogletConfig.Version))

		cogletInstall, err := g.installCogletWheel(cogletConfig, cogletIsPreRelease)
		if err != nil {
			return "", fmt.Errorf("failed to install coglet wheel: %w", err)
		}
		if cogletInstall != "" {
			installLines = cogletInstall
		}
	}

	// Install cog SDK
	var cogInstall string
	var err error
	switch wheelConfig.Source {
	case wheels.WheelSourcePyPI:
		cogInstall, err = g.installCogFromPyPI(wheelConfig, sdkIsPreRelease)
	case wheels.WheelSourceURL:
		console.Infof("Using cog wheel from URL: %s", wheelConfig.URL)
		cogInstall, err = g.installWheelFromURL(wheelConfig.URL)
	case wheels.WheelSourceFile:
		console.Debugf("Using local cog wheel: %s", wheelConfig.Path)
		cogInstall, err = g.installWheelFromFile(wheelConfig.Path)
	default:
		return "", fmt.Errorf("unknown wheel source: %v", wheelConfig.Source)
	}
	if err != nil {
		return "", err
	}
	if cogInstall != "" {
		if installLines != "" {
			installLines += "\n"
		}
		installLines += cogInstall
	}

	return installLines, nil
}

// installCogFromPyPI installs the cog SDK from PyPI.
// preRelease adds --pre to allow pip to resolve pre-release packages.
func (g *StandardGenerator) installCogFromPyPI(config *wheels.WheelConfig, preRelease bool) (string, error) {
	packageSpec := config.PyPIPackageURL("cog")
	flags := "--no-cache"
	if preRelease {
		flags = "--pre " + flags
	}
	pipInstallLine := "RUN " + uvCacheMount + " " + uvPip + " install " + flags + " " + packageSpec
	if g.strip {
		pipInstallLine += " && " + StripDebugSymbolsCommand
	}
	lines := []string{CFlags, pipInstallLine, "ENV CFLAGS="}
	return strings.Join(lines, "\n"), nil
}

// installWheelFromURL installs a wheel from a URL (when COG_SDK_WHEEL=https://...)
func (g *StandardGenerator) installWheelFromURL(url string) (string, error) {
	// Set coglet env vars if this looks like a coglet wheel
	var envLines []string
	if strings.Contains(url, "coglet") {
		if !CheckMajorMinorOnly(g.Config.Build.PythonVersion) {
			return "", fmt.Errorf("Python version must be <major>.<minor> for coglet")
		}
		envLines = []string{
			"ENV R8_COG_VERSION=coglet",
			"ENV R8_PYTHON_VERSION=" + g.Config.Build.PythonVersion,
		}
	}

	// For coglet URLs, uninstall cog first to avoid conflicts with coglet's cog shim package.
	// Some base images (e.g. r8.im/cog-base) have cog pre-installed, which conflicts
	// with coglet's cog compatibility shim that provides the same module paths.
	var pipPrefix string
	if strings.Contains(url, "coglet") {
		pipPrefix = uvPip + " uninstall cog 2>/dev/null || true && "
	}
	pipInstallLine := "RUN " + uvCacheMount + " " + pipPrefix + uvPip + " install --no-cache " + url
	if g.strip {
		pipInstallLine += " && " + StripDebugSymbolsCommand
	}

	envLines = append(envLines, CFlags, pipInstallLine, "ENV CFLAGS=")
	return strings.Join(envLines, "\n"), nil
}

// installWheelFromFile installs a wheel from a local file (when COG_SDK_WHEEL=/path/to/file.whl)
func (g *StandardGenerator) installWheelFromFile(path string) (string, error) {
	// Read the local wheel file
	data, err := os.ReadFile(path)
	if err != nil {
		return "", fmt.Errorf("failed to read wheel file %s: %w", path, err)
	}

	filename := filepath.Base(path)
	lines, containerPath, err := g.writeTemp(filename, data)
	if err != nil {
		return "", err
	}

	// Set coglet env vars if this looks like a coglet wheel
	var pipPrefix string
	if strings.Contains(filename, "coglet") {
		if !CheckMajorMinorOnly(g.Config.Build.PythonVersion) {
			return "", fmt.Errorf("Python version must be <major>.<minor> for coglet")
		}
		lines = append(lines,
			"ENV R8_COG_VERSION=coglet",
			"ENV R8_PYTHON_VERSION="+g.Config.Build.PythonVersion,
		)
		// Uninstall cog first to avoid conflicts with coglet's cog shim package.
		// Some base images (e.g. r8.im/cog-base) have cog pre-installed, which conflicts
		// with coglet's cog compatibility shim that provides the same module paths.
		pipPrefix = uvPip + " uninstall cog 2>/dev/null || true && "
	}

	pipInstallLine := "RUN " + uvCacheMount + " " + pipPrefix + uvPip + " install --no-cache " + containerPath
	if g.strip {
		pipInstallLine += " && " + StripDebugSymbolsCommand
	}
	lines = append(lines, CFlags, pipInstallLine, "ENV CFLAGS=")
	return strings.Join(lines, "\n"), nil
}

// installCogletWheel installs the coglet wheel based on the provided config.
// preRelease adds --pre to allow pip to resolve pre-release packages.
func (g *StandardGenerator) installCogletWheel(config *wheels.WheelConfig, preRelease bool) (string, error) {
	switch config.Source {
	case wheels.WheelSourcePyPI:
		return g.installCogletFromPyPI(config, preRelease)
	case wheels.WheelSourceURL:
		return g.installCogletFromURL(config.URL)
	case wheels.WheelSourceFile:
		return g.installCogletFromFile(config.Path)
	default:
		return "", fmt.Errorf("unknown coglet wheel source: %v", config.Source)
	}
}

// installCogletFromPyPI installs coglet from PyPI.
// preRelease adds --pre to allow pip to resolve pre-release packages.
func (g *StandardGenerator) installCogletFromPyPI(config *wheels.WheelConfig, preRelease bool) (string, error) {
	packageSpec := config.PyPIPackageURL("coglet")
	flags := "--no-cache"
	if preRelease {
		flags = "--pre " + flags
	}
	pipInstallLine := "RUN " + uvCacheMount + " " + uvPip + " install " + flags + " " + packageSpec
	if g.strip {
		pipInstallLine += " && " + StripDebugSymbolsCommand
	}
	lines := []string{CFlags, pipInstallLine, "ENV CFLAGS="}
	return strings.Join(lines, "\n"), nil
}

// installCogletFromURL installs coglet from a URL
func (g *StandardGenerator) installCogletFromURL(url string) (string, error) {
	pipInstallLine := "RUN " + uvCacheMount + " " + uvPip + " install --no-cache " + url
	if g.strip {
		pipInstallLine += " && " + StripDebugSymbolsCommand
	}
	lines := []string{CFlags, pipInstallLine, "ENV CFLAGS="}
	return strings.Join(lines, "\n"), nil
}

// installCogletFromFile installs coglet from a local wheel file
func (g *StandardGenerator) installCogletFromFile(path string) (string, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return "", fmt.Errorf("failed to read coglet wheel %s: %w", path, err)
	}

	filename := filepath.Base(path)
	lines, containerPath, err := g.writeTemp(filename, data)
	if err != nil {
		return "", err
	}

	pipInstallLine := "RUN " + uvCacheMount + " " + uvPip + " install --no-cache " + containerPath
	if g.strip {
		pipInstallLine += " && " + StripDebugSymbolsCommand
	}

	lines = append(lines, CFlags, pipInstallLine, "ENV CFLAGS=")
	return strings.Join(lines, "\n"), nil
}

// filterManagedPackages strips cog and coglet from user requirements content,
// warning loudly for each occurrence. requirements.txt is not the intended
// mechanism for controlling cog/coglet versions — use build.sdk_version in
// cog.yaml or the COG_SDK_WHEEL / COGLET_WHEEL environment variables instead.
func (g *StandardGenerator) filterManagedPackages(reqContents string) string {
	// Build a human-readable description of what the build system will install
	// for each managed package, so the warning is actionable.
	override := func(pkg string) string {
		var cfg *wheels.WheelConfig
		if pkg == "cog" {
			cfg = g.resolvedCogConfig
		} else {
			cfg = g.resolvedCogletConfig
		}
		if cfg == nil {
			return "latest from PyPI"
		}
		switch cfg.Source {
		case wheels.WheelSourcePyPI:
			if cfg.Version != "" {
				return fmt.Sprintf("%s==%s from PyPI", pkg, cfg.Version)
			}
			return "latest " + pkg + " from PyPI"
		case wheels.WheelSourceURL:
			return cfg.URL
		case wheels.WheelSourceFile:
			return cfg.Path
		default:
			return "unknown source"
		}
	}

	managed := map[string]bool{"cog": true, "coglet": true}
	var filtered []string
	for line := range strings.SplitSeq(reqContents, "\n") {
		trimmed := strings.TrimSpace(line)
		if trimmed == "" || strings.HasPrefix(trimmed, "#") || strings.HasPrefix(trimmed, "-") {
			filtered = append(filtered, line)
			continue
		}
		pkgName := requirements.PackageName(trimmed)
		baseName := strings.ToLower(strings.Split(pkgName, "[")[0])
		if managed[baseName] {
			console.Warnf(
				"'%s' found in requirements — overriding with %s. "+
					"Remove it from requirements and use build.sdk_version in cog.yaml or %s to control the version.",
				trimmed,
				override(baseName),
				map[string]string{"cog": "COG_SDK_WHEEL", "coglet": "COGLET_WHEEL"}[baseName],
			)
			continue
		}
		filtered = append(filtered, line)
	}
	return strings.Join(filtered, "\n")
}

func (g *StandardGenerator) pipInstalls() (string, error) {
	// Resolve wheel configs early so filterManagedPackages can emit precise warnings.
	if g.requiresCog {
		if err := g.resolveCogWheelConfigs(); err != nil {
			return "", err
		}
	}

	var err error
	includePackages := []string{}
	if torchVersion, ok := g.Config.TorchVersion(); ok {
		includePackages = []string{"torch==" + torchVersion}
	}
	if torchvisionVersion, ok := g.Config.TorchvisionVersion(); ok {
		includePackages = append(includePackages, "torchvision=="+torchvisionVersion)
	}
	if torchaudioVersion, ok := g.Config.TorchaudioVersion(); ok {
		includePackages = append(includePackages, "torchaudio=="+torchaudioVersion)
	}
	if tensorflowVersion, ok := g.Config.TensorFlowVersion(); ok {
		includePackages = append(includePackages, "tensorflow=="+tensorflowVersion)
	}
	g.pythonRequirementsContents, err = g.Config.PythonRequirementsForArch(g.GOOS, g.GOARCH, includePackages)
	if err != nil {
		return "", err
	}

	// Strip cog/coglet from user requirements — we always install them ourselves
	// via installCog(). Leaving them in would cause pip to overwrite our version.
	g.pythonRequirementsContents = g.filterManagedPackages(g.pythonRequirementsContents)

	if strings.Trim(g.pythonRequirementsContents, "") == "" {
		return "", nil
	}

	console.Debugf("Generated requirements.txt:\n%s", g.pythonRequirementsContents)
	copyLine, containerPath, err := g.writeTemp("requirements.txt", []byte(g.pythonRequirementsContents))
	if err != nil {
		return "", err
	}

	pipInstallLine := "RUN --mount=type=cache,target=/root/.cache/pip uv run pip install --cache-dir /root/.cache/pip -r " + containerPath
	if g.strip {
		pipInstallLine += " && " + StripDebugSymbolsCommand
	}
	return strings.Join([]string{
		copyLine[0],
		CFlags,
		pipInstallLine,
		"ENV CFLAGS=",
	}, "\n"), nil
}

func (g *StandardGenerator) runCommands() (string, error) {
	runCommands := g.Config.Build.Run

	// For backwards compatibility
	for _, command := range g.Config.Build.PreInstall {
		runCommands = append(runCommands, config.RunItem{Command: command})
	}

	lines := []string{}
	for _, run := range runCommands {
		command := strings.TrimSpace(run.Command)
		if strings.Contains(command, "\n") {
			return "", fmt.Errorf(`One of the commands in 'run' contains a new line, which won't work. You need to create a new list item in YAML prefixed with '-' for each command.

This is the offending line: %s`, command)
		}

		if len(run.Mounts) > 0 {
			mounts := []string{}
			for _, mount := range run.Mounts {
				if mount.Type == "secret" {
					secretMount := fmt.Sprintf("--mount=type=secret,id=%s,target=%s", mount.ID, mount.Target)
					mounts = append(mounts, secretMount)
				}
			}
			lines = append(lines, fmt.Sprintf("RUN %s %s", strings.Join(mounts, " "), command))
		} else {
			lines = append(lines, "RUN "+command)
		}
	}
	return strings.Join(lines, "\n"), nil
}

func (g *StandardGenerator) envVars() (string, error) {
	return envLineFromConfig(g.Config)
}

// installCACert generates Dockerfile lines to install a custom CA certificate.
// If COG_CA_CERT is not set, returns empty string (no-op).
func (g *StandardGenerator) installCACert() (string, error) {
	certData, err := ReadCACert()
	if err != nil {
		return "", err
	}
	if certData == nil {
		return "", nil
	}
	return GenerateCACertInstall(certData, g.writeTemp)
}

// writeTemp writes a temporary file that can be used as part of the build process
// It returns the lines to add to Dockerfile to make it available and the filename it ends up as inside the container
func (g *StandardGenerator) writeTemp(filename string, contents []byte) ([]string, string, error) {
	path := filepath.Join(g.tmpDir, filename)
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return []string{}, "", fmt.Errorf("Failed to write %s: %w", filename, err)
	}
	if err := os.WriteFile(path, contents, 0o644); err != nil {
		return []string{}, "", fmt.Errorf("Failed to write %s: %w", filename, err)
	}
	return []string{fmt.Sprintf("COPY %s /tmp/%s", filepath.Join(g.relativeTmpDir, filename), filename)}, "/tmp/" + filename, nil
}

func joinStringsWithoutLineSpace(chunks []string) string {
	lines := []string{}
	for _, chunk := range chunks {
		chunkLines := strings.Split(chunk, "\n")
		lines = append(lines, chunkLines...)
	}
	return strings.Join(filterEmpty(lines), "\n")
}

func filterEmpty(list []string) []string {
	filtered := []string{}
	for _, s := range list {
		if s != "" {
			filtered = append(filtered, s)
		}
	}
	return filtered
}

func (g *StandardGenerator) GenerateWeightsManifest(ctx context.Context) (*weights.Manifest, error) {
	m := weights.NewManifest()

	for _, dir := range g.modelDirs {
		err := g.fileWalker(dir, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if info.IsDir() {
				return nil
			}

			return m.AddFile(path)
		})
		if err != nil {
			return nil, err
		}
	}

	for _, path := range g.modelFiles {
		err := m.AddFile(path)
		if err != nil {
			return nil, err
		}
	}

	return m, nil
}

func (g *StandardGenerator) determineBaseImageName(ctx context.Context) (string, error) {
	var changed bool
	var err error

	cudaVersion := g.Config.Build.CUDA

	pythonVersion := g.Config.Build.PythonVersion
	pythonVersion, changed, err = stripPatchVersion(pythonVersion)
	if err != nil {
		return "", err
	}
	if changed {
		console.Warnf("Stripping patch version from Python version %s to %s", g.Config.Build.PythonVersion, pythonVersion)
	}

	torchVersion, _ := g.Config.TorchVersion()

	// validate that the base image configuration exists
	imageGenerator, err := NewBaseImageGenerator(ctx, g.client, cudaVersion, pythonVersion, torchVersion, g.command, false)
	if err != nil {
		return "", err
	}
	baseImage := BaseImageName(imageGenerator.cudaVersion, imageGenerator.pythonVersion, imageGenerator.torchVersion)
	return baseImage, nil
}

func stripPatchVersion(versionString string) (string, bool, error) {
	if versionString == "" {
		return "", false, nil
	}

	v, err := version.NewVersion(versionString)
	if err != nil {
		return "", false, fmt.Errorf("Invalid version: %s", versionString)
	}

	strippedVersion := fmt.Sprintf("%d.%d", v.Major, v.Minor)
	changed := strippedVersion != versionString

	return strippedVersion, changed, nil
}
