package graph

import (
	"context"
	"fmt"
	"io"
	"os"
	"strings"

	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/types"
	"oss.terrastruct.com/d2/d2format"
	"sigs.k8s.io/kustomize/api/resource"

	wh "edge-infra.dev/pkg/f8n/warehouse"
	"edge-infra.dev/pkg/f8n/warehouse/lift"
	"edge-infra.dev/pkg/f8n/warehouse/lift/cmd/internal"
	"edge-infra.dev/pkg/f8n/warehouse/oci"
	"edge-infra.dev/pkg/f8n/warehouse/oci/layer"
	whname "edge-infra.dev/pkg/f8n/warehouse/oci/name"
	"edge-infra.dev/pkg/f8n/warehouse/oci/walk"
	"edge-infra.dev/pkg/lib/cli/rags"
	"edge-infra.dev/pkg/lib/cli/sink"
)

const (
	outputDot  = "dot"
	outputD2   = "d2"
	outputTree = "tree"
)

var outputFmts = []string{outputTree, outputDot, outputD2}

func New(cfg lift.Config) *sink.Command {
	var (
		packer         = internal.NewPacker(cfg)
		outputType     string
		providerLabels bool
		depth          int
		outpath        string
	)

	cmd := &sink.Command{
		Use:   "graph [flags] <artifact>",
		Short: "render a graph for a Warehouse artifact in multiple formats",
		Flags: []*rags.Rag{
			{
				Name:  "type",
				Short: "t",
				Usage: fmt.Sprintf("output format. one of: %s", outputFmts),
				Value: &rags.StringEnum{
					Var:     &outputType,
					Default: outputTree,
					Valid:   []string{outputTree, outputDot, outputD2},
				},
			},
			{
				Name:  "add-provider-labels",
				Usage: "add labels to edges for provider variants. only applicable for --depth=1 or higher.",
				Value: &rags.Bool{Var: &providerLabels},
			},
			{
				Name:  "depth",
				Short: "d",
				Usage: "depth to draw the graph. only renders package nodes by default, increasing this value will render provider variant images, layers, and layer contents in that order.",
				Value: &rags.Int{Var: &depth},
			},
			{
				Name:  "out",
				Short: "o",
				Usage: "write output to file instead of stdout",
				Value: &rags.String{Var: &outpath},
			},
		},
		Extensions: []sink.Extension{packer},
		Exec: func(ctx context.Context, r sink.Run) error {
			switch len(r.Args()) {
			case 0:
				return fmt.Errorf("error: must provide an artifact")
			case 1:
			default:
				return fmt.Errorf("error: only one artifact can be provided")
			}

			if depth == 0 {
				f := false
				providerLabels = f
			}

			a, err := internal.ResolveArtifact(packer, r.Args()[0])
			if err != nil {
				return err
			}
			g, err := buildGraph(a, depth)
			if err != nil {
				return err
			}

			var out string
			switch outputType {
			case outputD2:
				d2, err := g.toD2(ctx, providerLabels)
				if err != nil {
					return err
				}
				out = d2format.Format(d2.g.AST)
			case outputTree:
				tree, err := g.toTree(providerLabels)
				if err != nil {
					return err
				}
				out = tree.String()
			case outputDot:
				dot, err := g.toDot()
				if err != nil {
					return err
				}
				out = dot.String()
			}

			if outpath != "" {
				return os.WriteFile(outpath, []byte(out), 0600)
			}
			fmt.Fprintln(r.Out(), out)

			return nil
		},
	}
	return cmd
}

type graph struct {
	nodes map[string]*node
	keys  []string // ordered list of visited keys so that we produce consistent results
}

const (
	idxNode     = "idx"
	imgNode     = "img"
	layerNode   = "layer"
	objectsNode = "objects"
)

type node struct {
	nodeType string
	name     string
	data     string // node types may optionally take in extra data
	digest   v1.Hash
	labels   map[string]string
	deps     []*node
}

func (n *node) label() string {
	return fmt.Sprintf("%s@sha256:%s", n.name, n.digest.Hex[:7])
}

func buildGraph(a oci.Artifact, depth int) (*graph, error) {
	g := &graph{nodes: make(map[string]*node)}
	visited := make(map[v1.Hash]bool)

	walker := &walk.Fns{
		Index: func(idx v1.ImageIndex, _ v1.ImageIndex) error {
			n, err := nodeFromArtifact(g, idx)
			if err != nil {
				return err
			}
			if visited[n.digest] {
				return nil
			}

			m, err := idx.IndexManifest()
			if err != nil {
				return err
			}
			for _, manifest := range m.Manifests {
				if depth == 0 && descRefsIdx(m, manifest) {
					continue
				}
				child := g.nodes[manifest.Digest.Hex]
				if child == nil {
					child = &node{
						nodeType: nodeTypeFromMediaType(manifest.MediaType),
						digest:   manifest.Digest,
						name:     manifest.Annotations[wh.AnnotationRefName],
					}
					g.nodes[manifest.Digest.Hex] = child
					g.keys = append(g.keys, manifest.Digest.Hex)
				}
				child.labels = map[string]string{
					"providers": manifest.Annotations[wh.AnnotationClusterProviders],
				}
				n.deps = append(n.deps, child)
			}
			visited[n.digest] = true

			return nil
		},
		Image: func(img v1.Image, parent v1.ImageIndex) error {
			isPkg, err := oci.IsPackage(img, parent)
			switch {
			case err != nil:
				return err
			case !isPkg && depth < 2:
				return nil
			}

			n, err := nodeFromArtifact(g, img)
			if err != nil {
				return err
			}
			if visited[n.digest] {
				return nil
			}

			layers, err := layer.FromImage(img)
			if err != nil {
				return err
			}
			for _, l := range layers {
				digest, err := l.Digest()
				if err != nil {
					return err
				}
				annos := l.Annotations()

				child := g.nodes[digest.Hex]
				if child == nil {
					name := annos[wh.AnnotationLayerType]
					switch {
					case annos[wh.AnnotationLayerRuntimeCapability] != "":
						name = fmt.Sprintf("%s/%s", name,
							annos[wh.AnnotationLayerRuntimeCapability])
					case name == layer.Infra.String():
						name = name[:5]
					}

					child = &node{
						nodeType: layerNode,
						digest:   digest,
						name:     name,
					}
					g.nodes[digest.Hex] = child
					g.keys = append(g.keys, digest.Hex)
				}
				n.deps = append(n.deps, child)

				if depth < 3 {
					continue
				}

				objsChild := g.nodes[digest.Hex+"-objs"]
				if objsChild == nil {
					objsChild = &node{
						nodeType: objectsNode,
						digest:   digest,
						name:     child.name,
					}

					reader, err := l.Uncompressed()
					if err != nil {
						return err
					}
					content, err := io.ReadAll(reader)
					if err != nil {
						return err
					}
					f := resource.Factory{}
					resources, err := f.SliceFromBytes(content)
					if err != nil {
						return err
					}
					objIDs := make([]string, len(resources))
					for i, r := range resources {
						objIDs[i] = fmt.Sprintf("%s/%s", r.GetKind(), r.GetName())
					}
					objsChild.data = strings.Join(objIDs, "\n")

					g.nodes[digest.Hex+"-objs"] = objsChild
					child.deps = append(child.deps, objsChild)
				}
			}
			visited[n.digest] = true
			return nil
		},
	}

	if err := walk.Walk(a, walker); err != nil {
		return nil, err
	}

	return g, nil
}

func nodeTypeFromMediaType(mt types.MediaType) string {
	switch {
	case mt.IsIndex():
		return idxNode
	case mt.IsImage():
		return imgNode
	case mt.IsLayer():
		return layerNode
	default:
		return ""
	}
}

func nodeFromArtifact(g *graph, a oci.Artifact) (*node, error) {
	digest, err := a.Digest()
	if err != nil {
		return nil, err
	}
	if g.nodes[digest.Hex] != nil {
		return g.nodes[digest.Hex], nil
	}

	mt, err := a.MediaType()
	if err != nil {
		return nil, err
	}

	name, err := whname.FromArtifact(a)
	if err != nil {
		return nil, err
	}

	n := &node{
		nodeType: nodeTypeFromMediaType(mt),
		digest:   digest,
		name:     name,
	}
	g.nodes[digest.Hex] = n
	g.keys = append(g.keys, digest.Hex)
	return n, nil
}

// descRefsIdx returns true if the input descriptor is referencing the input
// image manifest, e.g., provider variants of a package
func descRefsIdx(idx *v1.IndexManifest, d v1.Descriptor) bool {
	return idx.Annotations[wh.AnnotationName] == d.Annotations[wh.AnnotationRefName]
}