package daemon

import (
	"context"
	"errors"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"banger/internal/api"
	"banger/internal/daemon/imagemgr"
	"banger/internal/imagecat"
	"banger/internal/imagepull"
	"banger/internal/model"
	"banger/internal/paths"
	"banger/internal/system"

	"github.com/google/go-containerregistry/pkg/name"
)

// minPullExt4Size keeps the floor consistent with imagepull.MinExt4Size
// when the caller doesn't override --size and the OCI tree is tiny.
const minPullExt4Size int64 = 1 << 30 // 1 GiB

// PullImage downloads an image and registers it as a managed banger
// image. Two paths:
//
//   - Bundle path: `ref` matches an entry in the embedded imagecat
//     catalog. The `.tar.zst` bundle is fetched, `rootfs.ext4` is
//     already flattened + ownership-fixed + agent-injected at build
//     time, so this path is strictly faster than the OCI one.
//   - OCI path: otherwise treat `ref` as an OCI reference, pull its
//     layers, flatten, fix ownership, inject agents.
//
// Kernel info falls back through: `params.KernelRef` → catalog entry's
// `kernel_ref` (bundle path only) → `params.Kernel/Initrd/ModulesDir`.
//
// Concurrency: the slow staging work (network fetch, ext4 build,
// ownership fixup, guest-agent injection) runs WITHOUT imageOpsMu so
// parallel pulls of different images interleave. imageOpsMu is taken
// only for the publish window — recheck name is free, rename the
// staging dir to the final artifact dir, insert the store row. If two
// pulls race to the same name, the loser fails fast at the recheck
// and its staging dir is cleaned up via defer.
func (s *ImageService) PullImage(ctx context.Context, params api.ImagePullParams) (model.Image, error) {
	ref := strings.TrimSpace(params.Ref)
	if ref == "" {
		return model.Image{}, errors.New("reference is required")
	}

	catalog, err := imagecat.LoadEmbedded()
	if err != nil {
		return model.Image{}, fmt.Errorf("load image catalog: %w", err)
	}
	if entry, lookupErr := catalog.Lookup(ref); lookupErr == nil {
		return s.pullFromBundle(ctx, params, entry)
	}
	return s.pullFromOCI(ctx, params)
}

// publishImage is the narrow critical section shared by every image-
// creation path (pull bundle/OCI, register, promote). It re-verifies
// that `image.Name` is still free, atomically renames the staging
// directory to its final home (when applicable), and persists the row.
// The caller owns stagingDir cleanup on failure via its own defer; on
// success, publishImage unsets it so the defer is a no-op.
//
// finalDir == "" means "already published" (the caller built artifacts
// in place, e.g. RegisterImage which only touches the store). When
// non-empty the rename is the publication atom: finalDir must not
// already exist before the rename fires.
func (s *ImageService) publishImage(ctx context.Context, image model.Image, stagingDir, finalDir string) (model.Image, error) {
	s.imageOpsMu.Lock()
	defer s.imageOpsMu.Unlock()

	if existing, err := s.store.GetImageByName(ctx, image.Name); err == nil {
		return model.Image{}, fmt.Errorf("image %q already exists (id=%s); pick a different --name or delete it first", image.Name, existing.ID)
	}
	if finalDir != "" {
		if err := os.Rename(stagingDir, finalDir); err != nil {
			return model.Image{}, fmt.Errorf("publish artifact dir: %w", err)
		}
	}
	if err := s.store.UpsertImage(ctx, image); err != nil {
		if finalDir != "" {
			_ = os.RemoveAll(finalDir)
		}
		return model.Image{}, err
	}
	return image, nil
}

// pullFromOCI is the original OCI-registry-pull path. See PullImage for
// the intent.
func (s *ImageService) pullFromOCI(ctx context.Context, params api.ImagePullParams) (image model.Image, err error) {
	ref := strings.TrimSpace(params.Ref)
	parsed, err := name.ParseReference(ref)
	if err != nil {
		return model.Image{}, fmt.Errorf("parse oci ref %q: %w", ref, err)
	}

	imgName := strings.TrimSpace(params.Name)
	if imgName == "" {
		imgName = defaultImageNameFromRef(parsed)
		if imgName == "" {
			return model.Image{}, errors.New("could not derive image name from ref; pass --name")
		}
	}
	if existing, lookupErr := s.store.GetImageByName(ctx, imgName); lookupErr == nil {
		return model.Image{}, fmt.Errorf("image %q already exists (id=%s); pick a different --name or delete it first", imgName, existing.ID)
	}

	kernelPath, initrdPath, modulesDir, err := s.resolveKernelInputs(ctx, params.KernelRef, params.KernelPath, params.InitrdPath, params.ModulesDir)
	if err != nil {
		return model.Image{}, err
	}
	if err := imagemgr.ValidateKernelPaths(kernelPath, initrdPath, modulesDir); err != nil {
		return model.Image{}, err
	}

	id, err := model.NewID()
	if err != nil {
		return model.Image{}, err
	}
	finalDir := filepath.Join(s.layout.ImagesDir, id)
	stagingDir := finalDir + ".staging"
	if err := os.MkdirAll(stagingDir, 0o755); err != nil {
		return model.Image{}, err
	}
	cleanupStaging := true
	defer func() {
		if cleanupStaging {
			_ = os.RemoveAll(stagingDir)
		}
	}()

	// Extract OCI layers into a working tree under TempDir so the
	// state filesystem doesn't temporarily double in size.
	rootfsTree, err := os.MkdirTemp("", "banger-pull-")
	if err != nil {
		return model.Image{}, err
	}
	defer os.RemoveAll(rootfsTree)

	meta, err := s.runPullAndFlatten(ctx, ref, s.layout.OCICacheDir, rootfsTree)
	if err != nil {
		return model.Image{}, fmt.Errorf("pull oci image: %w", err)
	}

	sizeBytes := params.SizeBytes
	if sizeBytes <= 0 {
		treeSize, err := dirSizeBytes(rootfsTree)
		if err != nil {
			return model.Image{}, fmt.Errorf("size oci tree: %w", err)
		}
		sizeBytes = treeSize + treeSize/4 // +25% headroom
		if sizeBytes < minPullExt4Size {
			sizeBytes = minPullExt4Size
		}
	}

	rootfsExt4 := filepath.Join(stagingDir, "rootfs.ext4")
	if err := imagepull.BuildExt4(ctx, s.runner, rootfsTree, rootfsExt4, sizeBytes); err != nil {
		return model.Image{}, fmt.Errorf("build rootfs ext4: %w", err)
	}
	if err := s.runFinalizePulledRootfs(ctx, rootfsExt4, meta); err != nil {
		return model.Image{}, err
	}
	workSeedExt4 := s.runBuildWorkSeed(ctx, rootfsExt4, stagingDir)

	stagedKernel, stagedInitrd, stagedModules, err := imagemgr.StageBootArtifacts(ctx, s.runner, stagingDir, kernelPath, initrdPath, modulesDir)
	if err != nil {
		return model.Image{}, fmt.Errorf("stage boot artifacts: %w", err)
	}

	now := model.Now()
	image = model.Image{
		ID:          id,
		Name:        imgName,
		Managed:     true,
		ArtifactDir: finalDir,
		RootfsPath:  filepath.Join(finalDir, filepath.Base(rootfsExt4)),
		KernelPath:  rebaseUnder(stagedKernel, stagingDir, finalDir),
		InitrdPath:  rebaseUnder(stagedInitrd, stagingDir, finalDir),
		ModulesDir:  rebaseUnder(stagedModules, stagingDir, finalDir),
		CreatedAt:   now,
		UpdatedAt:   now,
	}
	if workSeedExt4 != "" {
		image.WorkSeedPath = filepath.Join(finalDir, filepath.Base(workSeedExt4))
	}
	published, err := s.publishImage(ctx, image, stagingDir, finalDir)
	if err != nil {
		return model.Image{}, err
	}
	cleanupStaging = false
	return published, nil
}

// pullFromBundle is the imagecat-backed path: download a ready-to-boot
// bundle (rootfs.ext4 already flattened + ownership-fixed + agent-
// injected at build time), verify its sha256, and register the result
// as a managed image. No flatten / mkfs / debugfs work on the daemon
// host.
func (s *ImageService) pullFromBundle(ctx context.Context, params api.ImagePullParams, entry imagecat.CatEntry) (image model.Image, err error) {
	imgName := strings.TrimSpace(params.Name)
	if imgName == "" {
		imgName = entry.Name
	}
	if existing, lookupErr := s.store.GetImageByName(ctx, imgName); lookupErr == nil {
		return model.Image{}, fmt.Errorf("image %q already exists (id=%s); pick a different --name or delete it first", imgName, existing.ID)
	}

	// Kernel resolution precedence: params > catalog entry's kernel_ref.
	kernelRef := strings.TrimSpace(params.KernelRef)
	if kernelRef == "" && strings.TrimSpace(params.KernelPath) == "" {
		kernelRef = strings.TrimSpace(entry.KernelRef)
	}
	kernelPath, initrdPath, modulesDir, err := s.resolveKernelInputs(ctx, kernelRef, params.KernelPath, params.InitrdPath, params.ModulesDir)
	if err != nil {
		return model.Image{}, err
	}
	if err := imagemgr.ValidateKernelPaths(kernelPath, initrdPath, modulesDir); err != nil {
		return model.Image{}, err
	}

	id, err := model.NewID()
	if err != nil {
		return model.Image{}, err
	}
	finalDir := filepath.Join(s.layout.ImagesDir, id)
	stagingDir := finalDir + ".staging"
	if err := os.MkdirAll(stagingDir, 0o755); err != nil {
		return model.Image{}, err
	}
	cleanupStaging := true
	defer func() {
		if cleanupStaging {
			_ = os.RemoveAll(stagingDir)
		}
	}()

	if _, err := s.runBundleFetch(ctx, stagingDir, entry); err != nil {
		return model.Image{}, fmt.Errorf("fetch bundle: %w", err)
	}
	// manifest.json is metadata we only need at fetch time; strip it
	// so the final artifact dir contains only boot-relevant files.
	_ = os.Remove(filepath.Join(stagingDir, imagecat.ManifestFilename))
	rootfsExt4 := filepath.Join(stagingDir, imagecat.RootfsFilename)
	workSeedExt4 := s.runBuildWorkSeed(ctx, rootfsExt4, stagingDir)

	stagedKernel, stagedInitrd, stagedModules, err := imagemgr.StageBootArtifacts(ctx, s.runner, stagingDir, kernelPath, initrdPath, modulesDir)
	if err != nil {
		return model.Image{}, fmt.Errorf("stage boot artifacts: %w", err)
	}

	now := model.Now()
	image = model.Image{
		ID:          id,
		Name:        imgName,
		Managed:     true,
		ArtifactDir: finalDir,
		RootfsPath:  filepath.Join(finalDir, filepath.Base(rootfsExt4)),
		KernelPath:  rebaseUnder(stagedKernel, stagingDir, finalDir),
		InitrdPath:  rebaseUnder(stagedInitrd, stagingDir, finalDir),
		ModulesDir:  rebaseUnder(stagedModules, stagingDir, finalDir),
		CreatedAt:   now,
		UpdatedAt:   now,
	}
	if workSeedExt4 != "" {
		image.WorkSeedPath = filepath.Join(finalDir, filepath.Base(workSeedExt4))
	}
	published, err := s.publishImage(ctx, image, stagingDir, finalDir)
	if err != nil {
		return model.Image{}, err
	}
	cleanupStaging = false
	return published, nil
}

// runBundleFetch is the seam tests substitute. nil → real implementation.
func (s *ImageService) runBundleFetch(ctx context.Context, destDir string, entry imagecat.CatEntry) (imagecat.Manifest, error) {
	if s.bundleFetch != nil {
		return s.bundleFetch(ctx, destDir, entry)
	}
	return imagecat.Fetch(ctx, nil, destDir, entry)
}

// runPullAndFlatten is the seam tests substitute. nil → real implementation.
func (s *ImageService) runPullAndFlatten(ctx context.Context, ref, cacheDir, destDir string) (imagepull.Metadata, error) {
	if s.pullAndFlatten != nil {
		return s.pullAndFlatten(ctx, ref, cacheDir, destDir)
	}
	pulled, err := imagepull.Pull(ctx, ref, cacheDir)
	if err != nil {
		return imagepull.Metadata{}, err
	}
	return imagepull.Flatten(ctx, pulled, destDir)
}

// runFinalizePulledRootfs applies ownership fixup and injects banger's
// guest agents. Tests substitute via s.finalizePulledRootfs; nil →
// real implementation using debugfs + the companion vsock-agent
// binary resolved via paths.CompanionBinaryPath.
func (s *ImageService) runFinalizePulledRootfs(ctx context.Context, ext4File string, meta imagepull.Metadata) error {
	if s.finalizePulledRootfs != nil {
		return s.finalizePulledRootfs(ctx, ext4File, meta)
	}
	if err := imagepull.ApplyOwnership(ctx, s.runner, ext4File, meta); err != nil {
		return fmt.Errorf("apply ownership: %w", err)
	}
	vsockBin, err := paths.CompanionBinaryPath("banger-vsock-agent")
	if err != nil {
		return fmt.Errorf("locate vsock agent binary: %w", err)
	}
	if err := imagepull.InjectGuestAgents(ctx, s.runner, ext4File, imagepull.GuestAgentAssets{
		VsockAgentBin: vsockBin,
	}); err != nil {
		return fmt.Errorf("inject guest agents: %w", err)
	}
	return nil
}

// runBuildWorkSeed extracts /root from the pulled rootfs into a
// sibling work-seed ext4 image. Any failure is treated as non-fatal:
// the image is still publishable without a seed, and VM create falls
// back to the empty-work-disk path (losing distro dotfiles but keeping
// every other guarantee). Returns the work-seed path on success, "" on
// failure (with a warn logged). Tests substitute via s.workSeedBuilder.
func (s *ImageService) runBuildWorkSeed(ctx context.Context, rootfsExt4, stagingDir string) string {
	outPath := filepath.Join(stagingDir, "work-seed.ext4")
	var err error
	if s.workSeedBuilder != nil {
		err = s.workSeedBuilder(ctx, rootfsExt4, outPath)
	} else {
		err = system.BuildWorkSeedImage(ctx, s.runner, rootfsExt4, outPath)
	}
	if err != nil {
		if s.logger != nil {
			s.logger.Warn("work-seed build failed; VMs using this image will start with an empty /root", "rootfs", rootfsExt4, "error", err.Error())
		}
		_ = os.Remove(outPath)
		return ""
	}
	return outPath
}

// nameSanitize keeps lowercase alphanumerics + hyphens, collapses runs.
var nameSanitizeRE = regexp.MustCompile(`[^a-z0-9]+`)

// defaultImageNameFromRef derives a friendly name like "debian-bookworm"
// from "docker.io/library/debian:bookworm". Returns "" if it can't.
func defaultImageNameFromRef(ref name.Reference) string {
	repo := ref.Context().RepositoryStr() // e.g. library/debian
	parts := strings.Split(repo, "/")
	base := parts[len(parts)-1]

	suffix := ""
	switch r := ref.(type) {
	case name.Tag:
		if t := r.TagStr(); t != "" && t != "latest" {
			suffix = "-" + t
		}
	case name.Digest:
		// take the first 12 hex chars after sha256:
		d := r.DigestStr()
		if i := strings.Index(d, ":"); i >= 0 && len(d) >= i+13 {
			suffix = "-" + d[i+1:i+13]
		}
	}

	out := nameSanitizeRE.ReplaceAllString(strings.ToLower(base+suffix), "-")
	out = strings.Trim(out, "-")
	return out
}

// rebaseUnder rewrites a path that points inside oldRoot to point inside
// newRoot. Empty input returns empty (kept by StageBootArtifacts when an
// optional artifact is absent).
func rebaseUnder(path, oldRoot, newRoot string) string {
	if path == "" {
		return ""
	}
	if rel, err := filepath.Rel(oldRoot, path); err == nil && !strings.HasPrefix(rel, "..") {
		return filepath.Join(newRoot, rel)
	}
	return path
}

// dirSizeBytes returns the sum of regular-file sizes under root, following
// no symlinks (lstat). Suitable for sizing an ext4 image.
func dirSizeBytes(root string) (int64, error) {
	var total int64
	err := filepath.WalkDir(root, func(_ string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if !d.Type().IsRegular() {
			return nil
		}
		info, err := d.Info()
		if err != nil {
			return err
		}
		total += info.Size()
		return nil
	})
	return total, err
}