package oci

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/opencontainers/go-digest"
	specs "github.com/opencontainers/image-spec/specs-go"
	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
)

// fakeRegistry implements the minimum OCI distribution-spec surface our
// Client touches: /v2/ probe, manifest fetch by tag or digest, blob fetch
// by digest. Backed by an in-memory blob+manifest store.
type fakeRegistry struct {
	t          *testing.T
	srv        *httptest.Server
	blobs      map[digest.Digest][]byte // keyed by digest
	manifests  map[string][]byte        // keyed by digest string (raw form)
	tags       map[string]digest.Digest // tag -> manifest digest
	mediaTypes map[digest.Digest]string // descriptor.MediaType per stored object
}

func newFakeRegistry(t *testing.T) *fakeRegistry {
	t.Helper()
	r := &fakeRegistry{
		t:          t,
		blobs:      map[digest.Digest][]byte{},
		manifests:  map[string][]byte{},
		tags:       map[string]digest.Digest{},
		mediaTypes: map[digest.Digest]string{},
	}
	r.srv = httptest.NewServer(http.HandlerFunc(r.handle))
	t.Cleanup(r.srv.Close)
	return r
}

func (r *fakeRegistry) putBlob(media string, data []byte) digest.Digest {
	h := sha256.Sum256(data)
	d := digest.NewDigestFromBytes(digest.SHA256, h[:])
	r.blobs[d] = data
	r.mediaTypes[d] = media
	return d
}

// putManifest stores a manifest/index document under both its digest and the
// given tag, returning the digest the caller can embed in indexes.
func (r *fakeRegistry) putManifest(tag string, media string, doc []byte) digest.Digest {
	h := sha256.Sum256(doc)
	d := digest.NewDigestFromBytes(digest.SHA256, h[:])
	r.manifests[d.String()] = doc
	r.mediaTypes[d] = media
	if tag != "" {
		r.tags[tag] = d
	}
	return d
}

// repoRef returns the "host:port/repo" string for use with NewClient.
func (r *fakeRegistry) repoRef() string {
	u, _ := url.Parse(r.srv.URL)
	return u.Host + "/test/kubesolo-os"
}

func (r *fakeRegistry) handle(w http.ResponseWriter, req *http.Request) {
	// Routes we implement:
	//   GET /v2/                          -> 200 "{}"
	//   GET /v2/test/kubesolo-os/manifests/<tag-or-digest> -> manifest
	//   HEAD same                         -> same headers, no body
	//   GET /v2/test/kubesolo-os/blobs/<digest> -> blob

	path := req.URL.Path
	if path == "/v2/" || path == "/v2" {
		w.Header().Set("Docker-Distribution-API-Version", "registry/2.0")
		w.WriteHeader(http.StatusOK)
		_, _ = io.WriteString(w, "{}")
		return
	}

	const prefix = "/v2/test/kubesolo-os/"
	if !strings.HasPrefix(path, prefix) {
		http.NotFound(w, req)
		return
	}
	rest := strings.TrimPrefix(path, prefix)

	switch {
	case strings.HasPrefix(rest, "manifests/"):
		ref := strings.TrimPrefix(rest, "manifests/")
		var d digest.Digest
		var data []byte
		if td, ok := r.tags[ref]; ok {
			d = td
			data = r.manifests[d.String()]
		} else if md, ok := r.manifests[ref]; ok {
			d = digest.Digest(ref)
			data = md
		} else {
			http.NotFound(w, req)
			return
		}
		media := r.mediaTypes[d]
		w.Header().Set("Content-Type", media)
		w.Header().Set("Docker-Content-Digest", d.String())
		w.Header().Set("Content-Length", fmt.Sprintf("%d", len(data)))
		if req.Method == http.MethodHead {
			return
		}
		_, _ = w.Write(data)

	case strings.HasPrefix(rest, "blobs/"):
		ref := strings.TrimPrefix(rest, "blobs/")
		d := digest.Digest(ref)
		blob, ok := r.blobs[d]
		if !ok {
			http.NotFound(w, req)
			return
		}
		media := r.mediaTypes[d]
		if media == "" {
			media = "application/octet-stream"
		}
		w.Header().Set("Content-Type", media)
		w.Header().Set("Docker-Content-Digest", d.String())
		w.Header().Set("Content-Length", fmt.Sprintf("%d", len(blob)))
		if req.Method == http.MethodHead {
			return
		}
		_, _ = w.Write(blob)

	default:
		http.NotFound(w, req)
	}
}

// seedSingleArchManifest puts kernel+initramfs blobs and a manifest with the
// given annotations into the registry, tagged as `tag`.
func (r *fakeRegistry) seedSingleArchManifest(t *testing.T, tag string, annot map[string]string) (kernelData, initramfsData []byte) {
	t.Helper()
	kernelData = []byte("FAKE-KERNEL-" + tag)
	initramfsData = []byte("FAKE-INITRAMFS-" + tag)

	kd := r.putBlob(MediaKernel, kernelData)
	id := r.putBlob(MediaInitramfs, initramfsData)

	// An empty config blob with sha256 of "{}" (the canonical "empty" body
	// per OCI). We don't actually fetch the config so any valid descriptor
	// works for the tests, but the digest still has to be syntactically valid.
	emptyConfigBody := []byte("{}")
	emptyConfigDigest := r.putBlob("application/vnd.oci.empty.v1+json", emptyConfigBody)

	manifest := ocispec.Manifest{
		Versioned: specs.Versioned{SchemaVersion: 2},
		MediaType: ocispec.MediaTypeImageManifest,
		Config: ocispec.Descriptor{
			MediaType: "application/vnd.oci.empty.v1+json",
			Size:      int64(len(emptyConfigBody)),
			Digest:    emptyConfigDigest,
		},
		Layers: []ocispec.Descriptor{
			{MediaType: MediaKernel, Digest: kd, Size: int64(len(kernelData))},
			{MediaType: MediaInitramfs, Digest: id, Size: int64(len(initramfsData))},
		},
		Annotations: annot,
	}
	manifestBytes, err := json.Marshal(manifest)
	if err != nil {
		t.Fatalf("marshal manifest: %v", err)
	}
	r.putManifest(tag, ocispec.MediaTypeImageManifest, manifestBytes)
	return
}

// seedIndex creates a manifest index pointing at per-arch manifests created
// via seedSingleArchManifest with arch-suffixed tags, then publishes the
// index under `tag`.
func (r *fakeRegistry) seedIndex(t *testing.T, tag string, perArchAnnots map[string]map[string]string) {
	t.Helper()
	var descriptors []ocispec.Descriptor
	for arch, annot := range perArchAnnots {
		// Reuse seedSingleArchManifest but under an internal arch-suffixed tag
		archTag := tag + "-" + arch
		r.seedSingleArchManifest(t, archTag, annot)
		d := r.tags[archTag]
		descriptors = append(descriptors, ocispec.Descriptor{
			MediaType: ocispec.MediaTypeImageManifest,
			Digest:    d,
			Size:      int64(len(r.manifests[d.String()])),
			Platform:  &ocispec.Platform{Architecture: arch, OS: "linux"},
		})
	}
	index := ocispec.Index{
		Versioned: specs.Versioned{SchemaVersion: 2},
		MediaType: ocispec.MediaTypeImageIndex,
		Manifests: descriptors,
	}
	indexBytes, _ := json.Marshal(index)
	r.putManifest(tag, ocispec.MediaTypeImageIndex, indexBytes)
}

// ---------------------------------------------------------------------------

func TestFetchMetadataSingleArchManifest(t *testing.T) {
	reg := newFakeRegistry(t)
	reg.seedSingleArchManifest(t, "v0.3.0", map[string]string{
		AnnotVersion: "v0.3.0",
		AnnotChannel: "stable",
		AnnotArch:    "amd64",
	})

	c, err := NewClient(reg.repoRef())
	if err != nil {
		t.Fatalf("NewClient: %v", err)
	}
	c.WithPlainHTTP(true)
	c.Arch = "amd64"

	meta, err := c.FetchMetadata(context.Background(), "v0.3.0")
	if err != nil {
		t.Fatalf("FetchMetadata: %v", err)
	}
	if meta.Version != "v0.3.0" {
		t.Errorf("version: got %q, want v0.3.0", meta.Version)
	}
	if meta.Channel != "stable" {
		t.Errorf("channel: got %q", meta.Channel)
	}
}

func TestFetchMetadataIndexSelectsArch(t *testing.T) {
	reg := newFakeRegistry(t)
	reg.seedIndex(t, "stable", map[string]map[string]string{
		"amd64": {AnnotVersion: "v0.3.0", AnnotChannel: "stable", AnnotArch: "amd64"},
		"arm64": {AnnotVersion: "v0.3.0", AnnotChannel: "stable", AnnotArch: "arm64"},
	})

	for _, arch := range []string{"amd64", "arm64"} {
		t.Run(arch, func(t *testing.T) {
			c, err := NewClient(reg.repoRef())
			if err != nil {
				t.Fatalf("NewClient: %v", err)
			}
			c.WithPlainHTTP(true)
			c.Arch = arch

			meta, err := c.FetchMetadata(context.Background(), "stable")
			if err != nil {
				t.Fatalf("FetchMetadata: %v", err)
			}
			if meta.Architecture != arch {
				t.Errorf("arch annotation: got %q, want %q", meta.Architecture, arch)
			}
			if meta.Version != "v0.3.0" {
				t.Errorf("version: got %q, want v0.3.0", meta.Version)
			}
		})
	}
}

func TestFetchMetadataIndexMissingArchErrors(t *testing.T) {
	reg := newFakeRegistry(t)
	reg.seedIndex(t, "stable", map[string]map[string]string{
		"amd64": {AnnotVersion: "v0.3.0", AnnotArch: "amd64"},
	})

	c, _ := NewClient(reg.repoRef())
	c.WithPlainHTTP(true)
	c.Arch = "arm64" // not in the index

	_, err := c.FetchMetadata(context.Background(), "stable")
	if err == nil {
		t.Fatal("expected error for missing arch, got nil")
	}
	if !strings.Contains(err.Error(), "arm64") {
		t.Errorf("expected error mentioning arm64, got: %v", err)
	}
}

func TestFetchMetadataSingleArchManifestRejectsCrossArch(t *testing.T) {
	// If the manifest declares an arch via annotation and it doesn't match
	// our runtime, Pull should refuse — defense in depth on top of the
	// channel/version gates in cmd/apply.go.
	reg := newFakeRegistry(t)
	reg.seedSingleArchManifest(t, "v0.3.0-arm64", map[string]string{
		AnnotArch: "arm64",
	})

	c, _ := NewClient(reg.repoRef())
	c.WithPlainHTTP(true)
	c.Arch = "amd64"

	_, err := c.FetchMetadata(context.Background(), "v0.3.0-arm64")
	if err == nil {
		t.Fatal("expected error pulling cross-arch single-arch manifest, got nil")
	}
}

func TestPullDownloadsBlobsAndVerifiesDigest(t *testing.T) {
	reg := newFakeRegistry(t)
	kernelData, initramfsData := reg.seedSingleArchManifest(t, "v0.3.0",
		map[string]string{AnnotVersion: "v0.3.0", AnnotArch: "amd64"})

	c, _ := NewClient(reg.repoRef())
	c.WithPlainHTTP(true)
	c.Arch = "amd64"

	stageDir := filepath.Join(t.TempDir(), "stage")
	staged, meta, err := c.Pull(context.Background(), "v0.3.0", stageDir)
	if err != nil {
		t.Fatalf("Pull: %v", err)
	}
	if meta.Version != "v0.3.0" {
		t.Errorf("meta version: got %q", meta.Version)
	}
	if staged.Version != "v0.3.0" {
		t.Errorf("staged version: got %q", staged.Version)
	}

	gotKernel, err := os.ReadFile(staged.VmlinuzPath)
	if err != nil {
		t.Fatalf("read kernel: %v", err)
	}
	if string(gotKernel) != string(kernelData) {
		t.Errorf("kernel mismatch:\n got %q\nwant %q", gotKernel, kernelData)
	}
	gotInit, err := os.ReadFile(staged.InitramfsPath)
	if err != nil {
		t.Fatalf("read initramfs: %v", err)
	}
	if string(gotInit) != string(initramfsData) {
		t.Errorf("initramfs mismatch")
	}
}

func TestPullRejectsTamperedBlob(t *testing.T) {
	// Mutate the kernel blob after it's been digested into the manifest.
	// Pull should refuse with a digest mismatch.
	reg := newFakeRegistry(t)
	_, _ = reg.seedSingleArchManifest(t, "v0.3.0",
		map[string]string{AnnotVersion: "v0.3.0", AnnotArch: "amd64"})

	// Corrupt every stored kernel blob in the registry by replacing its body.
	for d, m := range reg.mediaTypes {
		if m == MediaKernel {
			reg.blobs[d] = []byte("TAMPERED-KERNEL-WRONG-LENGTH-AND-DIGEST")
		}
	}

	c, _ := NewClient(reg.repoRef())
	c.WithPlainHTTP(true)
	c.Arch = "amd64"

	_, _, err := c.Pull(context.Background(), "v0.3.0", filepath.Join(t.TempDir(), "stage"))
	if err == nil {
		t.Fatal("expected digest mismatch error on tampered blob, got nil")
	}
	if !strings.Contains(err.Error(), "mismatch") {
		t.Errorf("expected mismatch in error, got: %v", err)
	}
}

func TestNewClientRejectsGarbageReference(t *testing.T) {
	_, err := NewClient("not a valid reference")
	if err == nil {
		t.Error("expected error on bad reference, got nil")
	}
}