package kvm2

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/digitalocean/go-libvirt"
	"github.com/go-logr/logr"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/trace"

	"gitea.home.hrajfrisbee.cz/kacerr/egress-proxies-operator/internal/provider"
	"gitea.home.hrajfrisbee.cz/kacerr/egress-proxies-operator/internal/tracing"
)

const (
	// markerCapacityBytes is the marker volume's size: one lowercase hex
	// sha256 digest, nothing else.
	markerCapacityBytes = sha256.Size * 2
	// markerReadCap bounds what is read back from a marker, so a wrong file
	// under that name cannot grow memory.
	markerReadCap = 1 << 10
	// seedReserveBytes is what the quota precheck sets aside for the seed
	// ISO next to the image and the overlay (spec entry 13).
	seedReserveBytes = 16 << 20
)

var hexDigest = regexp.MustCompile(`^[0-9a-f]{64}$`)

// imageSource is a resolved placement.image (spec entries 3 and 4): the
// catalogue entry, or a URL given directly with the catalogue defaults, plus
// the key its base volume is named after.
type imageSource struct {
	Entry string
	Key   string
	provider.KVMImage
}

func (s imageSource) volumeName() string { return baseVolumeName(s.Entry, s.Key, s.Format) }
func (s imageSource) markerName() string { return markerVolumeName(s.volumeName()) }

// resolveImage maps placement.image to a source. cfg is already defaulted,
// so a catalogue entry carries its format and firmware; a URL gets qcow2
// and uefi. entry is the provider entry's name, which the base volume is
// named after so two entries never share one file.
func resolveImage(entry string, cfg provider.KVMConfig, image string) (imageSource, error) {
	if img, ok := cfg.Images[image]; ok {
		return imageSource{Entry: entry, Key: image, KVMImage: img}, nil
	}
	if provider.IsHTTPURL(image) {
		return imageSource{
			Entry:    entry,
			Key:      urlImageKey(image),
			KVMImage: provider.KVMImage{URL: image, Format: provider.KVMFormatQcow2, Firmware: provider.KVMFirmwareUEFI},
		}, nil
	}
	return imageSource{}, permanentf("placement.image %q is neither a key in kvm2.images nor an http(s) URL", image)
}

// lockImage takes the per-(host, volume) lock of spec entry 19, or gives up
// with ctx: a caller whose ctx ends while another caller's download runs
// returns promptly (contract §1.5) instead of waiting out a 20-minute
// transfer.
func (p *Provider) lockImage(ctx context.Context, hostName, volume string) (func(), error) {
	release, err := p.imageLocks.acquire(ctx, hostName+"/"+volume)
	if err != nil {
		return nil, fmt.Errorf("%w: waiting for image %s on host %s: %w", provider.ErrTransient, volume, hostName, err)
	}
	return release, nil
}

// ensureImage makes the base image present and complete in the host's pool
// and returns its path (spec entry 4). be is the caller's pooled connection
// and carries the cheap checks; the upload itself goes over a connection of
// its own. Concurrent callers for the same (host, volume) queue on the
// lock: the first fetches, the rest find the marker.
func (p *Provider) ensureImage(ctx context.Context, host provider.KVMHost, pool libvirt.StoragePool, be backend, src imageSource) (string, error) {
	unlock, err := p.lockImage(ctx, host.Name, src.volumeName())
	if err != nil {
		return "", err
	}
	defer unlock()
	log := p.logger(ctx).WithValues("host", host.Name, "volume", src.volumeName())
	path, err := cachedImage(ctx, log, pool, be, src)
	if err != nil || path != "" {
		return path, err
	}
	vol, err := p.fetchImage(ctx, host, pool, be, src)
	if err != nil {
		return "", err
	}
	path, err = be.StorageVolGetPath(ctx, vol)
	if err != nil {
		return "", fmt.Errorf("reading the path of base volume %s: %w", vol.Name, err)
	}
	return path, nil
}

// cachedImage returns the base volume's path when the marker says the base
// is complete and its digest is the configured one (or none is configured).
// Anything else is cleared so the fetch starts from nothing, and "" is
// returned: a base without a marker is an upload that was cut off; a marker
// without a base, or one whose content is not a digest, is debris; a marker
// with a different digest means the catalogue entry changed.
func cachedImage(ctx context.Context, log logr.Logger, pool libvirt.StoragePool, be backend, src imageSource) (string, error) {
	name, markerName := src.volumeName(), src.markerName()
	marker, err := be.StorageVolLookupByName(ctx, pool, markerName)
	switch {
	case hasCode(err, libvirt.ErrNoStorageVol):
		return "", deleteVolumeIfPresent(ctx, be, pool, name)
	case err != nil:
		return "", fmt.Errorf("looking up marker %s: %w", markerName, err)
	}
	digest, err := readMarker(ctx, be, marker)
	if err != nil {
		return "", err
	}
	switch {
	case !hexDigest.MatchString(digest):
		log.V(1).Info("base image marker holds no digest; refetching")
	case src.SHA256 != "" && !strings.EqualFold(digest, src.SHA256):
		// Refusing rather than replacing: every overlay on this host names
		// the base as its backing file, and libvirt will unlink it under
		// running domains without a word.
		return "", permanentf("base image %s is cached with sha256 %s but %s is configured; add the image under a new key and move Proxies to it, or delete the Proxies using it and evict the base by hand",
			name, digest, strings.ToLower(src.SHA256))
	default:
		base, err := be.StorageVolLookupByName(ctx, pool, name)
		switch {
		case hasCode(err, libvirt.ErrNoStorageVol):
			return "", deleteVolumeIfPresent(ctx, be, pool, markerName)
		case err != nil:
			return "", fmt.Errorf("looking up base volume %s: %w", name, err)
		}
		path, err := be.StorageVolGetPath(ctx, base)
		if err != nil {
			return "", fmt.Errorf("reading the path of base volume %s: %w", name, err)
		}
		return path, nil
	}
	return "", errors.Join(deleteVolumeIfPresent(ctx, be, pool, markerName), deleteVolumeIfPresent(ctx, be, pool, name))
}

// readMarker returns the digest a marker volume holds, lowercased and
// stripped of padding. The volume is 64 bytes; the read is capped anyway.
func readMarker(ctx context.Context, be backend, marker libvirt.StorageVol) (string, error) {
	w := &cappedWriter{limit: markerReadCap}
	if err := be.StorageVolDownload(ctx, marker, w, 0, 0, 0); err != nil {
		return "", fmt.Errorf("reading marker %s: %w", marker.Name, err)
	}
	return strings.ToLower(strings.TrimRight(string(w.buf), "\x00 \r\n\t")), nil
}

// cappedWriter keeps the first limit bytes and accepts the rest without
// storing them, so a download into it is bounded whatever its length.
type cappedWriter struct {
	limit int
	buf   []byte
}

func (w *cappedWriter) Write(b []byte) (int, error) {
	if room := w.limit - len(w.buf); room > 0 {
		w.buf = append(w.buf, b[:min(room, len(b))]...)
	}
	return len(b), nil
}

// deleteVolumeIfPresent removes the named volume; one that is already gone
// is success.
func deleteVolumeIfPresent(ctx context.Context, be backend, pool libvirt.StoragePool, name string) error {
	vol, err := be.StorageVolLookupByName(ctx, pool, name)
	if hasCode(err, libvirt.ErrNoStorageVol) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("looking up volume %s: %w", name, err)
	}
	if err := be.StorageVolDelete(ctx, vol, 0); err != nil && !hasCode(err, libvirt.ErrNoStorageVol) {
		return fmt.Errorf("deleting volume %s: %w", name, err)
	}
	return nil
}

// fetchImage streams the image from its URL into a new base volume and then
// writes the marker, under the kvm2.download span (spec entry 17). The
// upload goes over a connection of its own (spec entry 18): abandoning a
// call kills the connection it ran on, and a 20-minute upload must not die
// because a poll for another Proxy was cancelled on the pooled one. The
// marker is written last, on the pooled connection, after
// StoragePoolRefresh made libvirt read the file's real format; its presence
// is what marks the base complete.
//
// Every exit between the base volume's creation and the marker deletes the
// base volume (best effort): the next attempt would delete it anyway on
// finding no marker, and removing it now frees the space at once.
func (p *Provider) fetchImage(ctx context.Context, host provider.KVMHost, pool libvirt.StoragePool, be backend, src imageSource) (libvirt.StorageVol, error) {
	name := src.volumeName()
	ctx, cancel := context.WithTimeout(ctx, defaultBudgets.transfer)
	defer cancel()
	ctx, span := tracing.StartSpan(ctx, tracing.Tracer(p.tracing...), "kvm2.download",
		trace.WithSpanKind(trace.SpanKindClient),
		trace.WithAttributes(
			attribute.String("kvm2.host", host.Name),
			attribute.String("kvm2.volume", name),
			attribute.String("kvm2.url_host", urlHost(src.URL)),
			attribute.Bool("kvm2.sha256_configured", src.SHA256 != ""),
		))
	defer span.End()
	// Taken after the span starts, so the lines carry its ID.
	log := p.logger(ctx).WithValues("host", host.Name, "volume", name)
	start := time.Now()

	// The body runs out of time before the surrounding budget does, so a
	// mirror that stalls mid-image ends the upload through its reader
	// rather than through the killed socket the budget would leave.
	httpCtx, cancelHTTP := context.WithTimeout(ctx, defaultBudgets.transfer-transferGrace)
	defer cancelHTTP()
	resp, err := p.do(httpCtx, http.MethodGet, src.URL)
	if err == nil {
		err = checkResponse(src.URL, resp)
	}
	if err != nil {
		span.SetStatus(codes.Error, err.Error())
		return libvirt.StorageVol{}, err
	}
	defer func() { _ = resp.Body.Close() }()
	size := resp.ContentLength
	span.SetAttributes(attribute.Int64("kvm2.bytes", size))
	log.V(1).Info("fetching base image into host pool", "url", src.URL, "bytes", size)

	vol, digest, err := p.uploadImage(ctx, host, pool, src, resp.Body, size)
	if err != nil {
		// uploadImage has returned, so its dedicated connection is closed
		// and libvirtd has let go of the stream before the file goes.
		p.discardVolume(ctx, log, be, pool, name)
		span.SetStatus(codes.Error, err.Error())
		return libvirt.StorageVol{}, err
	}
	if err := be.StoragePoolRefresh(ctx, pool, 0); err != nil {
		err = fmt.Errorf("refreshing storage pool %s after the upload: %w", pool.Name, err)
		p.discardVolume(ctx, log, be, pool, name)
		span.SetStatus(codes.Error, err.Error())
		return libvirt.StorageVol{}, err
	}
	if err := writeMarker(ctx, be, pool, src.markerName(), digest); err != nil {
		p.discardVolume(ctx, log, be, pool, name)
		span.SetStatus(codes.Error, err.Error())
		return libvirt.StorageVol{}, err
	}
	log.V(1).Info("base image cached", "bytes", size, "duration", time.Since(start), "sha256", digest)
	return vol, nil
}

// uploadImage creates the base volume and streams body into it over a
// dedicated connection with a running sha256, then checks length and
// digest. The error says why: a body that ended early is Transient, a
// digest that differs from the configured one is Permanent. Discarding
// the volume is the caller's, once this connection is closed.
func (p *Provider) uploadImage(ctx context.Context, host provider.KVMHost, pool libvirt.StoragePool, src imageSource, body io.Reader, size int64) (libvirt.StorageVol, string, error) {
	name := src.volumeName()
	xmlDoc, err := buildVolumeXML(volumeSpec{Name: name, CapacityBytes: size, Format: src.Format})
	if err != nil {
		return libvirt.StorageVol{}, "", err
	}
	fresh, err := p.hosts.dialFresh(ctx, host)
	if err != nil {
		return libvirt.StorageVol{}, "", fmt.Errorf("dialling host %s for the image upload: %w", host.Name, err)
	}
	defer func() { _ = fresh.Close() }()
	vol, err := fresh.StorageVolCreateXML(ctx, pool, xmlDoc, 0)
	if err != nil {
		return libvirt.StorageVol{}, "", fmt.Errorf("creating base volume %s: %w", name, err)
	}

	hash := sha256.New()
	counted := &countingReader{r: io.TeeReader(body, hash)}
	uploadErr := fresh.StorageVolUpload(ctx, vol, counted, 0, uint64(size), 0)
	read, readErr := counted.state()
	var digest string
	if uploadErr == nil {
		// Only then has the upload stopped reading; an abandoned call's
		// goroutine may still be writing into the hash.
		digest = hex.EncodeToString(hash.Sum(nil))
	}
	switch {
	case readErr != nil:
		err = fmt.Errorf("%w: image %s: the response ended after %d of %d bytes: %w", provider.ErrTransient, src.URL, read, size, readErr)
	case uploadErr != nil:
		err = fmt.Errorf("uploading base volume %s: %w", name, uploadErr)
	case read != size:
		err = transientf("image %s: short read, %d of %d bytes", src.URL, read, size)
	case src.SHA256 != "" && !strings.EqualFold(digest, src.SHA256):
		err = permanentf("image %s: sha256 %s does not match the configured %s", src.URL, digest, strings.ToLower(src.SHA256))
	}
	if err != nil {
		return libvirt.StorageVol{}, "", err
	}
	return vol, digest, nil
}

// countingReader counts what the upload consumed and keeps the first
// non-EOF error the source returned. The adapter can return from an
// abandoned upload while go-libvirt still reads, so the fields are guarded.
type countingReader struct {
	r   io.Reader
	mu  sync.Mutex
	n   int64
	err error
}

func (c *countingReader) Read(b []byte) (int, error) {
	n, err := c.r.Read(b)
	c.mu.Lock()
	c.n += int64(n)
	if err != nil && err != io.EOF && c.err == nil {
		c.err = err
	}
	c.mu.Unlock()
	return n, err
}

func (c *countingReader) state() (int64, error) {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.n, c.err
}

// discardVolume removes a base volume a failed fetch left behind. It runs
// with ctx's values but not its cancellation: the fetch may have failed
// because ctx ended, and the adapter refuses calls on a finished ctx. The
// budget bounds it on its own. Failure is logged, not returned: the next
// attempt finds the volume without a marker and deletes it then.
func (p *Provider) discardVolume(ctx context.Context, log logr.Logger, be backend, pool libvirt.StoragePool, name string) {
	ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), defaultBudgets.rpc)
	defer cancel()
	if err := deleteVolumeIfPresent(ctx, be, pool, name); err != nil {
		log.V(1).Info("incomplete base volume left behind; the next attempt removes it", "error", err.Error())
	}
}

// writeMarker creates the raw marker volume and uploads the digest into it.
// A marker that could not be filled is removed again: an empty marker next
// to a complete base would read as debris and cost a refetch.
func writeMarker(ctx context.Context, be backend, pool libvirt.StoragePool, name, digest string) error {
	xmlDoc, err := buildVolumeXML(volumeSpec{Name: name, CapacityBytes: markerCapacityBytes, Format: provider.KVMFormatRaw})
	if err != nil {
		return err
	}
	vol, err := be.StorageVolCreateXML(ctx, pool, xmlDoc, 0)
	if err != nil {
		return fmt.Errorf("creating marker %s: %w", name, err)
	}
	if err := be.StorageVolUpload(ctx, vol, strings.NewReader(digest), 0, uint64(len(digest)), 0); err != nil {
		_ = deleteVolumeIfPresent(ctx, be, pool, name)
		return fmt.Errorf("writing marker %s: %w", name, err)
	}
	return nil
}

// do issues one request bound to ctx. Accept-Encoding is fixed to identity
// because a transparently decompressed body reports no Content-Length. A
// transport error is Transient.
func (p *Provider) do(ctx context.Context, method, rawURL string) (*http.Response, error) {
	req, err := http.NewRequestWithContext(ctx, method, rawURL, nil)
	if err != nil {
		return nil, permanentf("image %s: %v", rawURL, err)
	}
	req.Header.Set("Accept-Encoding", "identity")
	resp, err := p.http.Do(req)
	if err != nil {
		return nil, fmt.Errorf("%w: fetching image %s: %w", provider.ErrTransient, rawURL, err)
	}
	return resp, nil
}

// checkResponse classifies the response at the source (spec entry 13) and
// closes its body when it is unusable. 404, 403 and 410, and every other
// 4xx, mean the catalogue entry is wrong; 5xx and 429 are worth a retry. A
// missing Content-Length is Permanent: the volume's capacity must be
// declared before the upload, and a mirror that only chunks is one the
// catalogue must not point at.
func checkResponse(rawURL string, resp *http.Response) error {
	var err error
	code := resp.StatusCode
	switch {
	case code >= 200 && code < 300:
	case code == http.StatusTooManyRequests || code >= 500:
		err = transientf("image %s: HTTP %s", rawURL, resp.Status)
	case code >= 400:
		err = permanentf("image %s: HTTP %s", rawURL, resp.Status)
	default:
		err = transientf("image %s: unexpected HTTP %s", rawURL, resp.Status)
	}
	if err == nil && resp.ContentLength < 0 {
		err = permanentf("image %s: response has no Content-Length; the volume capacity must be known before upload", rawURL)
	}
	if err == nil && resp.ContentLength == 0 {
		err = permanentf("image %s: response is empty (Content-Length 0)", rawURL)
	}
	if err != nil {
		_ = resp.Body.Close()
	}
	return err
}

// urlHost is the only part of the URL a span carries (contract §9).
func urlHost(rawURL string) string {
	u, err := url.Parse(rawURL)
	if err != nil {
		return ""
	}
	return u.Host
}

// imageSizeIfUncached returns the image's Content-Length when the base is
// not cached on the host, and 0 when the marker is present, for the quota
// precheck of spec entry 13. It never downloads: HEAD first, and a GET whose
// body is closed at once when HEAD answers anything but 2xx, since servers
// refuse HEAD with 405, 403 or 501 alike; the GET's status is what gets
// classified. A marker is taken at face value here; ensureImage decides
// whether it still matches.
func (p *Provider) imageSizeIfUncached(ctx context.Context, be backend, pool libvirt.StoragePool, src imageSource) (int64, error) {
	_, err := be.StorageVolLookupByName(ctx, pool, src.markerName())
	switch {
	case err == nil:
		return 0, nil
	case !hasCode(err, libvirt.ErrNoStorageVol):
		return 0, fmt.Errorf("looking up marker %s: %w", src.markerName(), err)
	}
	ctx, cancel := context.WithTimeout(ctx, defaultBudgets.rpc)
	defer cancel()
	resp, err := p.do(ctx, http.MethodHead, src.URL)
	if err != nil {
		return 0, err
	}
	_ = resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		if resp, err = p.do(ctx, http.MethodGet, src.URL); err != nil {
			return 0, err
		}
		_ = resp.Body.Close()
	}
	if err := checkResponse(src.URL, resp); err != nil {
		return 0, err
	}
	return resp.ContentLength, nil
}

// checkPoolSpace is the precheck of spec entry 13: the pool must be running
// and have needed bytes available, or Create returns ErrQuotaExceeded and
// the reconciler waits QuotaRetry instead of failing the volume creation.
func checkPoolSpace(ctx context.Context, be backend, pool libvirt.StoragePool, poolName string, needed uint64) error {
	state, _, _, available, err := be.StoragePoolGetInfo(ctx, pool)
	if err != nil {
		return fmt.Errorf("reading storage pool %s: %w", poolName, err)
	}
	if s := libvirt.StoragePoolState(state); s != libvirt.StoragePoolRunning {
		return permanentf("storage pool %s is not running (state %s)", poolName, s)
	}
	if available < needed {
		return fmt.Errorf("%w: storage pool %s has %d bytes available, %d needed", provider.ErrQuotaExceeded, poolName, available, needed)
	}
	return nil
}

// neededBytes is what one Create commits to the pool: the image when it is
// not cached yet (0 otherwise), the overlay's capacity and the seed reserve.
func neededBytes(imageBytes int64, diskSizeGB int64) uint64 {
	return uint64(max(imageBytes, 0)) + uint64(max(diskSizeGB, 0))<<30 + seedReserveBytes
}