package migrate

import (
	"archive/tar"
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/arescom/dockmv/internal/dkr"
	"github.com/arescom/dockmv/internal/job"
	"github.com/arescom/dockmv/internal/spec"
)

// PackageFormat selects how the finished package is laid out on disk.
type PackageFormat string

const (
	// FormatDir leaves an unpacked directory, easiest to inspect and to copy
	// onto a USB stick that is already mounted.
	FormatDir PackageFormat = "dir"
	// FormatTar produces a single .tar file, easiest to move around.
	FormatTar PackageFormat = "tar"
)

// Packager writes a self-contained migration package: the container specs, the
// data archives, optionally the images, and a shell installer that replays it
// all on a target that has nothing but docker.
type Packager struct {
	Src        *dkr.Client
	Containers []spec.Container
	Volumes    []spec.Volume
	Networks   []spec.Network
	Plan       spec.Plan

	// OutputDir is the directory packages are created under.
	OutputDir string
	// Format selects a directory or a single tar file.
	Format PackageFormat
	// SourceHost is recorded in the manifest.
	SourceHost string
}

// Result describes the produced package.
type Result struct {
	Path  string `json:"path"`
	Bytes int64  `json:"bytes"`
	Name  string `json:"name"`
}

// Run builds the package, reporting progress into j.
func (p *Packager) Run(ctx context.Context, j *job.Job) (*Result, error) {
	opts := p.Plan.Options
	name := p.Plan.PackageName
	if name == "" {
		name = "docker-migration-" + time.Now().Format("20060102-150405")
	}
	name = sanitize(name)

	root := filepath.Join(p.OutputDir, name)
	if _, err := os.Stat(root); err == nil {
		return nil, fmt.Errorf("package %s already exists in %s", name, p.OutputDir)
	}
	if err := os.MkdirAll(root, 0o755); err != nil {
		return nil, fmt.Errorf("create package directory: %w", err)
	}
	cleanup := true
	defer func() {
		if cleanup {
			os.RemoveAll(root)
		}
	}()

	byID := map[string]*spec.Container{}
	for i := range p.Containers {
		byID[p.Containers[i].ID] = &p.Containers[i]
	}

	var prepared []*Prepared
	for _, sel := range p.Plan.Items {
		if !sel.Include {
			continue
		}
		pr, err := Prepare(byID[sel.ContainerID], sel, p.Volumes, p.Networks)
		if err != nil {
			return nil, fmt.Errorf("container %s: %w", sel.ContainerID, err)
		}
		prepared = append(prepared, pr)
	}
	if len(prepared) == 0 {
		return nil, errors.New("nothing selected to migrate")
	}

	man := spec.Manifest{
		FormatVersion: 1,
		CreatedAt:     time.Now(),
		CreatedBy:     "dockmv",
		SourceHost:    p.SourceHost,
		Options:       opts,
		Items:         p.Plan.Items,
	}
	if v, err := p.Src.Ping(ctx); err == nil {
		man.DockerVersion = v
	}

	savedImages := map[string]string{} // image ref -> payload path

	for _, pr := range prepared {
		item := j.AddItem(pr.Source.ID, pr.Source.Name)
		for _, n := range pr.Source.Warnings {
			j.AddItemWarning(item, "%s: %s", pr.Source.Name, n)
		}
		for _, n := range pr.Notes {
			j.AddItemWarning(item, "%s: %s", pr.Source.Name, n)
		}

		err := p.packOne(ctx, j, item, root, pr, opts, &man, savedImages)
		if err != nil {
			j.SetItemState(item, job.StateFailed, err)
			return nil, fmt.Errorf("%s: %w", pr.Source.Name, err)
		}
		j.SetItemState(item, job.StateSucceeded, nil)

		man.Containers = append(man.Containers, *pr.Target)
		man.Volumes = append(man.Volumes, pr.Volumes...)
		for _, n := range pr.Networks {
			if !hasNetwork(man.Networks, n.Name) {
				man.Networks = append(man.Networks, n)
			}
		}
	}

	if err := writeJSON(filepath.Join(root, "manifest.json"), man); err != nil {
		return nil, err
	}
	installer := renderInstaller(prepared, &man, savedImages, name)
	if err := os.WriteFile(filepath.Join(root, "install.sh"), []byte(installer), 0o755); err != nil {
		return nil, fmt.Errorf("write installer: %w", err)
	}
	if err := os.WriteFile(filepath.Join(root, "README.txt"), []byte(renderReadme(name, prepared, opts)), 0o644); err != nil {
		return nil, fmt.Errorf("write readme: %w", err)
	}
	j.Logf(job.LevelInfo, "", "wrote installer, manifest and readme")

	if p.Format == FormatTar {
		tarPath := root + ".tar"
		j.Logf(job.LevelInfo, "", "packing %s into a single archive", name)
		size, err := tarDirectory(ctx, root, tarPath, name)
		if err != nil {
			os.Remove(tarPath)
			return nil, fmt.Errorf("create package archive: %w", err)
		}
		os.RemoveAll(root)
		cleanup = false
		j.SetArtifact(tarPath, size)
		return &Result{Path: tarPath, Bytes: size, Name: name + ".tar"}, nil
	}

	size, _ := dirSize(root)
	cleanup = false
	j.SetArtifact(root, size)
	return &Result{Path: root, Bytes: size, Name: name}, nil
}

func (p *Packager) packOne(
	ctx context.Context, j *job.Job, item *job.Item, root string,
	pr *Prepared, opts spec.Options, man *spec.Manifest, savedImages map[string]string,
) error {
	compress := opts.Compress
	level := opts.CompressLevel

	// Image.
	if pr.Selection.MigrateImage && pr.Selection.ImageMode != spec.ImageSkip && pr.Selection.ImageMode != spec.ImagePull {
		ref := pr.Target.Image
		if _, done := savedImages[ref]; !done {
			size := p.Src.ImageSizeBytes(ctx, ref)
			st := j.AddStep(item, "image", "save image "+ref, size)
			j.StartStep(st)
			if opts.DryRun {
				j.SkipStep(st, "dry run: image not written")
			} else {
				rel := filepath.ToSlash(filepath.Join("images", sanitize(ref)+tarExt(compress)))
				payload, err := p.streamToFile(ctx, j, st, filepath.Join(root, filepath.FromSlash(rel)), rel, compress, level,
					func() (io.ReadCloser, error) { return p.Src.SaveImage(ctx, ref) })
				j.FinishStep(st, err)
				if err != nil {
					return fmt.Errorf("save image %s: %w", ref, err)
				}
				payload.Kind, payload.Image = "image", ref
				man.Payloads = append(man.Payloads, *payload)
				savedImages[ref] = rel
			}
		} else {
			st := j.AddStep(item, "image", "image "+ref+" already in package", 0)
			j.StartStep(st)
			j.SkipStep(st, "shared with another container")
		}
	} else if pr.Selection.ImageMode == spec.ImagePull {
		j.Logf(job.LevelInfo, item.ID, "%s: image %s will be pulled by the installer", pr.Source.Name, pr.Target.Image)
	}

	// Data. The source container is stopped for the duration when asked.
	if len(pr.Transfers) == 0 {
		return nil
	}
	restore, err := p.quiesce(ctx, j, item, pr, opts)
	if err != nil {
		return err
	}
	defer func() {
		if restore != nil {
			restore()
		}
	}()

	for i, t := range pr.Transfers {
		st := j.AddStep(item, fmt.Sprintf("data-%d", i), t.Label, t.SizeBytes)
		j.StartStep(st)
		if opts.DryRun {
			j.SkipStep(st, "dry run: data not written")
			continue
		}
		rel := filepath.ToSlash(filepath.Join("data", sanitize(pr.ContainerName()),
			fmt.Sprintf("%02d-%s%s", i, sanitize(strings.Trim(t.Destination, "/")), tarExt(compress))))
		payload, err := p.streamToFile(ctx, j, st, filepath.Join(root, filepath.FromSlash(rel)), rel, compress, level,
			func() (io.ReadCloser, error) { return p.Src.CopyOut(ctx, pr.Source.ID, t.SourcePath) })
		j.FinishStep(st, err)
		if err != nil {
			return fmt.Errorf("archive %s: %w", t.Label, err)
		}
		payload.Kind = "mount"
		payload.Container = pr.ContainerName()
		payload.Destination = t.Destination
		man.Payloads = append(man.Payloads, *payload)
	}
	return nil
}

// streamToFile copies a stream to a file inside the package, optionally
// gzipping it, while counting bytes and computing a checksum.
func (p *Packager) streamToFile(
	ctx context.Context, j *job.Job, st *job.Step,
	absPath, relPath string, compress bool, level int,
	open func() (io.ReadCloser, error),
) (*spec.Payload, error) {
	if err := os.MkdirAll(filepath.Dir(absPath), 0o755); err != nil {
		return nil, err
	}
	src, err := open()
	if err != nil {
		return nil, err
	}
	defer src.Close()

	f, err := os.Create(absPath)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	hash := sha256.New()
	// The checksum covers the bytes as stored, so the installer can verify the
	// file it is about to feed to docker.
	out := io.MultiWriter(f, hash)

	counted := job.NewCountingReader(src, j, st)
	var copyErr error
	if compress {
		gz, gerr := gzip.NewWriterLevel(out, gzipLevel(nil, level))
		if gerr != nil {
			return nil, gerr
		}
		_, copyErr = io.Copy(gz, counted)
		if cerr := gz.Close(); copyErr == nil {
			copyErr = cerr
		}
	} else {
		_, copyErr = io.Copy(out, counted)
	}
	counted.Flush()
	if copyErr != nil {
		return nil, copyErr
	}
	if err := f.Sync(); err != nil {
		return nil, err
	}
	info, err := f.Stat()
	if err != nil {
		return nil, err
	}
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}
	return &spec.Payload{
		Path:       relPath,
		Bytes:      info.Size(),
		SHA256:     hex.EncodeToString(hash.Sum(nil)),
		Compressed: compress,
	}, nil
}

func (p *Packager) quiesce(ctx context.Context, j *job.Job, item *job.Item, pr *Prepared, opts spec.Options) (func(), error) {
	if opts.DryRun {
		return nil, nil
	}
	wasRunning := pr.Source.State == "running"
	if !pr.Selection.StopSourceDuringCopy {
		if wasRunning {
			j.AddItemWarning(item,
				"archiving %s while it is running; data written during the copy may be inconsistent", pr.Source.Name)
		}
		return nil, nil
	}
	if !wasRunning {
		return nil, nil
	}
	st := j.AddStep(item, "quiesce", "stop source "+pr.Source.Name, 0)
	j.StartStep(st)
	err := p.Src.Stop(ctx, pr.Source.ID, 30*time.Second)
	j.FinishStep(st, err)
	if err != nil {
		return nil, fmt.Errorf("stop source container: %w", err)
	}
	return func() {
		// Building a package does not move the workload anywhere, so the source
		// is always put back the way it was found.
		if err := p.Src.Start(context.WithoutCancel(ctx), pr.Source.ID); err != nil {
			j.AddItemWarning(item, "could not restart source container: %v", err)
		} else {
			j.Logf(job.LevelInfo, item.ID, "source container %s restarted", pr.Source.Name)
		}
	}, nil
}

func tarExt(compress bool) string {
	if compress {
		return ".tar.gz"
	}
	return ".tar"
}

func writeJSON(path string, v any) error {
	b, err := json.MarshalIndent(v, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(path, b, 0o644)
}

func hasNetwork(ns []spec.Network, name string) bool {
	for _, n := range ns {
		if n.Name == name {
			return true
		}
	}
	return false
}

// tarDirectory packs a package directory into a single tar file, keeping the
// directory name as the archive's top-level entry.
func tarDirectory(ctx context.Context, dir, dest, prefix string) (int64, error) {
	f, err := os.Create(dest)
	if err != nil {
		return 0, err
	}
	defer f.Close()
	tw := tar.NewWriter(f)

	err = filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if ctx.Err() != nil {
			return ctx.Err()
		}
		rel, err := filepath.Rel(dir, path)
		if err != nil {
			return err
		}
		name := prefix
		if rel != "." {
			name = prefix + "/" + filepath.ToSlash(rel)
		}
		hdr, err := tar.FileInfoHeader(info, "")
		if err != nil {
			return err
		}
		hdr.Name = name
		if info.IsDir() {
			hdr.Name += "/"
		}
		if err := tw.WriteHeader(hdr); err != nil {
			return err
		}
		if info.IsDir() {
			return nil
		}
		src, err := os.Open(path)
		if err != nil {
			return err
		}
		defer src.Close()
		_, err = io.Copy(tw, src)
		return err
	})
	if err != nil {
		tw.Close()
		return 0, err
	}
	if err := tw.Close(); err != nil {
		return 0, err
	}
	info, err := f.Stat()
	if err != nil {
		return 0, err
	}
	return info.Size(), nil
}

func dirSize(dir string) (int64, error) {
	var total int64
	err := filepath.Walk(dir, func(_ string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() {
			total += info.Size()
		}
		return nil
	})
	return total, err
}