package vulndb

import (
	"archive/tar"
	"compress/gzip"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
	"oras.land/oras-go/v2"
	"oras.land/oras-go/v2/registry"
	"oras.land/oras-go/v2/registry/remote"
)

// DefaultRef is the published trivy-db OCI artifact, rebuilt every six hours.
const DefaultRef = "ghcr.io/aquasecurity/trivy-db:2"

// SupportedSchema is the trivy-db schema version this code understands.
//
// A different version is refused rather than parsed on the assumption it is
// close enough. Mis-reading the schema would not fail loudly — it would return
// no advisories, which is indistinguishable from a clean fleet.
const SupportedSchema = 2

// dbFileName and metaFileName are the two files inside the artifact layer.
const (
	dbFileName   = "trivy.db"
	metaFileName = "metadata.json"
)

// Ref returns the artifact reference, honouring VANTAGE_TRIVY_DB_REF so an
// air-gapped deployment can mirror the artifact into its own registry, and so
// a busy deployment can avoid the anonymous ghcr rate limit.
func Ref() string {
	if v := os.Getenv("VANTAGE_TRIVY_DB_REF"); v != "" {
		return v
	}
	return DefaultRef
}

// Disabled reports whether the puller and scheduler are switched off entirely.
// Findings already written are still served, and still marked stale.
func Disabled() bool {
	return strings.EqualFold(os.Getenv("VANTAGE_VULNDB_DISABLED"), "true")
}

// DebugEnabled turns on per-package and per-advisory tracing.
//
// It is a switch rather than always-on because a single scan asks the store one
// question per installed package — ~2000 lines per server, per tick — which
// would bury every other subsystem's logs on a fleet of any size. The lifecycle
// logs (pull, tick, per-server totals) are unconditional; only the inner loop
// is gated.
func DebugEnabled() bool {
	return strings.EqualFold(os.Getenv("VANTAGE_VULN_DEBUG"), "true")
}

// Debugf logs only when VANTAGE_VULN_DEBUG=true.
func Debugf(format string, args ...any) {
	if DebugEnabled() {
		log.Printf("vulndb[debug]: "+format, args...)
	}
}

// dbMetadata is the subset of trivy-db's metadata.json we read.
type dbMetadata struct {
	Version int `json:"Version"`
}

// Pull fetches the trivy-db artifact into dir and returns its schema version.
//
// It extracts into a staging directory and only moves the files into place once
// both are present and the schema has been accepted. A pull that fails partway
// therefore leaves the previous database untouched rather than a half-written
// one that Open would happily accept and scan against.
func Pull(ctx context.Context, dir string) (int, error) {
	ref := Ref()
	started := time.Now()
	log.Printf("vulndb: pull starting ref=%s dir=%s", ref, dir)

	parsed, err := registry.ParseReference(ref)
	if err != nil {
		return 0, fmt.Errorf("parse reference %q: %w", ref, err)
	}

	repo, err := remote.NewRepository(ref)
	if err != nil {
		return 0, fmt.Errorf("open repository %q: %w", ref, err)
	}

	// The tag or digest half of the reference; the repository already knows the
	// registry and path.
	target := parsed.Reference
	if target == "" {
		target = "latest"
	}

	_, manifestBytes, err := oras.FetchBytes(ctx, repo, target, oras.DefaultFetchBytesOptions)
	if err != nil {
		return 0, fmt.Errorf("fetch manifest %s: %w", ref, err)
	}

	var man ocispec.Manifest
	if err := json.Unmarshal(manifestBytes, &man); err != nil {
		return 0, fmt.Errorf("decode manifest %s: %w", ref, err)
	}
	if len(man.Layers) == 0 {
		return 0, fmt.Errorf("artifact %s has no layers", ref)
	}
	log.Printf("vulndb: manifest resolved layers=%d digest=%s size=%dB",
		len(man.Layers), man.Layers[0].Digest, man.Layers[0].Size)

	// Streamed rather than buffered: the layer is ~50MB and there is no reason
	// to hold it in memory on the way to disk.
	rc, err := repo.Blobs().Fetch(ctx, man.Layers[0])
	if err != nil {
		return 0, fmt.Errorf("fetch layer: %w", err)
	}
	defer rc.Close()

	staging, err := os.MkdirTemp(dir, ".staging-")
	if err != nil {
		return 0, fmt.Errorf("staging dir: %w", err)
	}
	defer os.RemoveAll(staging)

	if err := extractTarGz(rc, staging); err != nil {
		return 0, fmt.Errorf("extract layer: %w", err)
	}
	log.Printf("vulndb: layer extracted into %s after %s", staging, time.Since(started).Round(time.Millisecond))

	metaBytes, err := os.ReadFile(filepath.Join(staging, metaFileName))
	if err != nil {
		return 0, fmt.Errorf("read %s: %w", metaFileName, err)
	}
	var meta dbMetadata
	if err := json.Unmarshal(metaBytes, &meta); err != nil {
		return 0, fmt.Errorf("decode %s: %w", metaFileName, err)
	}
	if meta.Version != SupportedSchema {
		return 0, fmt.Errorf("trivy-db schema %d is not supported (want %d)", meta.Version, SupportedSchema)
	}

	fi, err := os.Stat(filepath.Join(staging, dbFileName))
	if err != nil {
		return 0, fmt.Errorf("artifact has no %s: %w", dbFileName, err)
	}
	log.Printf("vulndb: %s is %dB, schema %d accepted", dbFileName, fi.Size(), meta.Version)

	// Both files present and the schema accepted, so it is safe to replace.
	for _, name := range []string{dbFileName, metaFileName} {
		src := filepath.Join(staging, name)
		dst := filepath.Join(dir, name)
		if err := os.Remove(dst); err != nil && !os.IsNotExist(err) {
			return 0, fmt.Errorf("remove old %s: %w", name, err)
		}
		if err := os.Rename(src, dst); err != nil {
			return 0, fmt.Errorf("install %s: %w", name, err)
		}
	}

	log.Printf("vulndb: pull complete ref=%s schema=%d in %s", ref, meta.Version, time.Since(started).Round(time.Millisecond))
	return meta.Version, nil
}

// extractTarGz writes the artifact layer into dir. Paths are flattened and
// checked so a crafted archive cannot write outside dir.
func extractTarGz(r io.Reader, dir string) error {
	gz, err := gzip.NewReader(r)
	if err != nil {
		return err
	}
	defer gz.Close()

	tr := tar.NewReader(gz)
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
		if hdr.Typeflag != tar.TypeReg {
			continue
		}
		name := filepath.Base(hdr.Name) // flatten; the archive is two files
		if name == "." || name == ".." || name == "" {
			continue
		}
		dst := filepath.Join(dir, name)
		if !strings.HasPrefix(dst, filepath.Clean(dir)+string(os.PathSeparator)) {
			return fmt.Errorf("archive entry escapes destination: %q", hdr.Name)
		}
		f, err := os.Create(dst)
		if err != nil {
			return err
		}
		if _, err := io.Copy(f, tr); err != nil {
			f.Close()
			return err
		}
		if err := f.Close(); err != nil {
			return err
		}
	}
}
