package byoc

import (
	"bufio"
	"bytes"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"slices"
	"strings"

	distributionreference "github.com/distribution/reference"
	"gopkg.in/yaml.v3"
)

const defaultCoreImageRegistry = "385595570414.dkr.ecr.us-west-2.amazonaws.com"

// These Config Center images are the only exact seed references whose
// manifests are not present in the canonical staging ECR. Keep their public
// source explicit; every other Config Center image is mirrored from the
// access-controlled source registry used by the production image pipeline.
var configCenterPublicOnlySources = map[string]struct{}{
	"gcr.io/pingcap-public/dbaas/pd:release-nextgen-20251011-next-gen":      {},
	"gcr.io/pingcap-public/dbaas/tidb:release-nextgen-20251011-next-gen":    {},
	"gcr.io/pingcap-public/dbaas/tiflash:release-nextgen-20251011-next-gen": {},
}

// CoreImage is one immutable source reference and its path in the customer ECR
// registry. Tag is retained so the product configuration can keep using the
// same human-readable immutable tag after the exact manifest is copied.
type CoreImage struct {
	Source     string `json:"source"`
	Repository string `json:"repository"`
	Tag        string `json:"tag"`
}

// ImageClosure is the deterministic customer-dataplane image set bound into
// the activation fingerprint. Digest changes require destroy/recreate instead
// of silently changing an already-ready Region.
type ImageClosure struct {
	Digest string      `json:"digest"`
	Images []CoreImage `json:"images"`
}

var quotedSQLValuePattern = regexp.MustCompile(`'([^']*)'`)

// DiscoverCoreImageClosure derives the exact image closure from the same
// Premium dataplane templates and Config Center seed consumed by the local
// BYOC Region. It intentionally excludes control-plane runtime images: those
// continue to run in the control-plane account.
func DiscoverCoreImageClosure(repoRoot string) (*ImageClosure, error) {
	files := []string{
		"personas/app-delivery/envs/premium/dataplane-manager/k8s-infra-templates.yaml",
		"personas/app-delivery/envs/premium/dataplane-manager/k8s-dedicated-templates.yaml",
		"personas/app-delivery/envs/premium/dataplane-manager/k8s-dp-templates.yaml",
		"personas/app-delivery/envs/premium/dataplane-manager/k8s-addon-templates.yaml",
		"personas/app-delivery/envs/premium/cluster-ng-server/config.yaml",
		"personas/app-delivery/envs/premium/cluster-ng-executor/config.yaml",
	}
	images := make(map[string]CoreImage)
	for _, name := range files {
		filePath := filepath.Join(repoRoot, filepath.FromSlash(name))
		data, err := os.ReadFile(filePath)
		if err != nil {
			return nil, fmt.Errorf("read BYOC image source %s: %w", name, err)
		}
		if err := discoverYAMLImages(data, images); err != nil {
			return nil, fmt.Errorf("discover BYOC images in %s: %w", name, err)
		}
	}

	ticdcConfigName := "personas/app-delivery/envs/premium/ticdc-service/config.yaml"
	ticdcConfig, err := os.ReadFile(filepath.Join(repoRoot, filepath.FromSlash(ticdcConfigName)))
	if err != nil {
		return nil, fmt.Errorf("read BYOC image source %s: %w", ticdcConfigName, err)
	}
	ticdcCatalogName := "personas/app-delivery/envs/premium/premium-metadb/61-dedicated-config-center-20260625.sql"
	ticdcCatalog, err := os.ReadFile(filepath.Join(repoRoot, filepath.FromSlash(ticdcCatalogName)))
	if err != nil {
		return nil, fmt.Errorf("read BYOC image source %s: %w", ticdcCatalogName, err)
	}
	if err := discoverTiCDCServiceImages(ticdcConfig, ticdcCatalog, images); err != nil {
		return nil, fmt.Errorf("discover BYOC images in %s: %w", ticdcConfigName, err)
	}

	for _, sqlName := range []string{
		"personas/app-delivery/envs/premium/premium-metadb/30-config-center-regional-20260521.sql",
		"personas/app-delivery/envs/premium/premium-metadb/68-premium-config-center-staging-defaults-20260723.sql",
	} {
		sqlData, err := os.ReadFile(filepath.Join(repoRoot, filepath.FromSlash(sqlName)))
		if err != nil {
			return nil, fmt.Errorf("read BYOC image source %s: %w", sqlName, err)
		}
		if err := discoverConfigCenterImages(sqlData, images); err != nil {
			return nil, fmt.Errorf("discover BYOC images in %s: %w", sqlName, err)
		}
	}

	closureImages := make([]CoreImage, 0, len(images))
	for _, image := range images {
		if !imageReferenceIsPinned(image.Source) {
			return nil, fmt.Errorf("BYOC core image is not pinned: %s", image.Source)
		}
		closureImages = append(closureImages, image)
	}
	slices.SortFunc(closureImages, func(a, b CoreImage) int {
		return strings.Compare(a.Source, b.Source)
	})
	if len(closureImages) == 0 {
		return nil, errors.New("BYOC core image closure is empty")
	}
	return &ImageClosure{Digest: closureDigest(closureImages), Images: closureImages}, nil
}

func discoverTiCDCServiceImages(configData, catalogData []byte, images map[string]CoreImage) error {
	var config struct {
		BaseImage    string `yaml:"ticdcBaseImage"`
		PreStopImage string `yaml:"ticdcPreStopImage"`
	}
	if err := yaml.Unmarshal(configData, &config); err != nil {
		return fmt.Errorf("parse TiCDC service config: %w", err)
	}
	if config.BaseImage == "" || config.PreStopImage == "" {
		return errors.New("TiCDC service config requires base and pre-stop images")
	}
	if imageReferenceIsPinned(config.BaseImage) || looksLikeExplicitImageReference(config.BaseImage) {
		return fmt.Errorf("TiCDC base image must be an untagged repository because the service appends the catalog tag: %s", config.BaseImage)
	}
	if err := addCoreImage(images, config.PreStopImage); err != nil {
		return fmt.Errorf("TiCDC pre-stop image: %w", err)
	}

	scanner := bufio.NewScanner(bytes.NewReader(catalogData))
	found := false
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.Contains(line, "'ticdc'") {
			continue
		}
		matches := quotedSQLValuePattern.FindAllStringSubmatch(line, -1)
		if len(matches) < 5 || matches[1][1] != "ticdc" {
			continue
		}
		if found {
			return errors.New("TiCDC catalog contains more than one image pin")
		}
		if path.Base(config.BaseImage) != path.Base(matches[3][1]) {
			return fmt.Errorf("TiCDC catalog repository %s does not match service repository %s", matches[3][1], config.BaseImage)
		}
		if err := addCoreImage(images, config.BaseImage+":"+matches[4][1]); err != nil {
			return fmt.Errorf("TiCDC base image: %w", err)
		}
		found = true
	}
	if err := scanner.Err(); err != nil {
		return err
	}
	if !found {
		return errors.New("TiCDC catalog image pin is missing")
	}
	return nil
}

func discoverYAMLImages(data []byte, images map[string]CoreImage) error {
	decoder := yaml.NewDecoder(bytes.NewReader(data))
	for {
		var document yaml.Node
		err := decoder.Decode(&document)
		if errors.Is(err, io.EOF) {
			return nil
		}
		if err != nil {
			return err
		}
		if err := collectYAMLImages(&document, "", "", images); err != nil {
			return err
		}
	}
}

func collectYAMLImages(node *yaml.Node, inheritedRegistry, key string, images map[string]CoreImage) error {
	if node == nil {
		return nil
	}
	switch node.Kind {
	case yaml.DocumentNode:
		for _, child := range node.Content {
			if err := collectYAMLImages(child, inheritedRegistry, key, images); err != nil {
				return err
			}
		}
	case yaml.MappingNode:
		fields := make(map[string]*yaml.Node, len(node.Content)/2)
		for index := 0; index+1 < len(node.Content); index += 2 {
			fields[node.Content[index].Value] = node.Content[index+1]
		}
		registry := inheritedRegistry
		if value := scalarValue(fields["registry"]); value != "" {
			registry = strings.TrimSuffix(value, "/")
		}
		repository := scalarValue(fields["repository"])
		tag := scalarValue(fields["tag"])
		if repository != "" && tag != "" {
			if err := addCoreImage(images, qualifyImageRepository(repository, registry)+":"+tag); err != nil {
				return err
			}
		}
		for index := 0; index+1 < len(node.Content); index += 2 {
			if err := collectYAMLImages(node.Content[index+1], registry, node.Content[index].Value, images); err != nil {
				return err
			}
		}
	case yaml.SequenceNode:
		for _, child := range node.Content {
			if err := collectYAMLImages(child, inheritedRegistry, key, images); err != nil {
				return err
			}
		}
	case yaml.ScalarNode:
		value := strings.TrimSpace(node.Value)
		if value == "" {
			if isImageReferenceField(key) {
				return fmt.Errorf("image field %s requires a non-empty pinned image reference", key)
			}
			return nil
		}
		source, imageFlag, err := imageReferenceFromFlag(value)
		if err != nil {
			return err
		}
		if imageFlag {
			if err := addCoreImage(images, source); err != nil {
				return fmt.Errorf("image flag %s: %w", value, err)
			}
			return nil
		}
		if isImageReferenceField(key) {
			source = value
			if !looksLikeExplicitImageReference(source) {
				source = qualifyImageRepository(source, inheritedRegistry)
			}
			if err := addCoreImage(images, source); err != nil {
				return fmt.Errorf("image field %s: %w", key, err)
			}
			return nil
		}
		if looksLikeExplicitImageReference(value) {
			return addCoreImage(images, value)
		}
	}
	return nil
}

func discoverConfigCenterImages(data []byte, images map[string]CoreImage) error {
	scanner := bufio.NewScanner(bytes.NewReader(data))
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.Contains(line, "component_images") || !strings.Contains(line, "VALUES") {
			continue
		}
		matches := quotedSQLValuePattern.FindAllStringSubmatch(line, -1)
		if len(matches) < 5 {
			continue
		}
		baseImage := matches[3][1]
		tag := matches[4][1]
		if tag == "not-used" {
			continue
		}
		if baseImage != "" && tag != "" {
			repository := "tidbcloud/" + path.Base(baseImage)
			if err := addCoreImageAs(images, configCenterImageSource(baseImage, tag, repository), repository); err != nil {
				return err
			}
		}
	}
	return scanner.Err()
}

func configCenterImageSource(baseImage, tag, repository string) string {
	publicSource := baseImage + ":" + tag
	if _, ok := configCenterPublicOnlySources[publicSource]; ok {
		return publicSource
	}
	if strings.HasPrefix(baseImage, "gcr.io/pingcap-public/dbaas/") ||
		strings.HasPrefix(baseImage, "us.gcr.io/pingcap-public/tidbx/") {
		return defaultCoreImageRegistry + "/" + repository + ":" + tag
	}
	return publicSource
}

func addCoreImage(images map[string]CoreImage, source string) error {
	return addCoreImageAs(images, source, "")
}

func addCoreImageAs(images map[string]CoreImage, source, repository string) error {
	image, err := parseCoreImage(source)
	if err != nil {
		return err
	}
	if repository != "" {
		image.Repository = repository
	}
	for existingSource, existing := range images {
		if existingSource != image.Source && existing.Repository == image.Repository && existing.Tag == image.Tag {
			return fmt.Errorf(
				"customer image target %s:%s has conflicting sources %s and %s",
				image.Repository,
				image.Tag,
				existing.Source,
				image.Source,
			)
		}
	}
	images[image.Source] = image
	return nil
}

func parseCoreImage(source string) (CoreImage, error) {
	source = strings.TrimSpace(strings.Trim(source, "'\""))
	if !imageReferenceIsPinned(source) {
		return CoreImage{}, fmt.Errorf("image reference is not pinned: %s", source)
	}
	parsed, err := distributionreference.Parse(source)
	if err != nil {
		return CoreImage{}, fmt.Errorf("image reference is malformed: %s: %w", source, err)
	}
	named, ok := parsed.(distributionreference.Named)
	if !ok || !isRegistryHost(distributionreference.Domain(named)) || parsed.String() != source {
		return CoreImage{}, fmt.Errorf("image reference has no explicit valid registry: %s", source)
	}
	lastSlash := strings.LastIndex(source, "/")
	var repositoryWithRegistry, tag string
	if digestIndex := strings.LastIndex(source, "@sha256:"); digestIndex > lastSlash {
		digest := source[digestIndex+1:]
		if !validSHA256Digest(digest) {
			return CoreImage{}, fmt.Errorf("image reference has an invalid digest: %s", source)
		}
		repositoryWithRegistry = source[:digestIndex]
		tag = "cp9-digest-" + strings.TrimPrefix(digest, "sha256:")[:16]
	} else {
		lastColon := strings.LastIndex(source, ":")
		if lastColon <= lastSlash {
			return CoreImage{}, fmt.Errorf("image reference has no tag: %s", source)
		}
		tag = source[lastColon+1:]
		repositoryWithRegistry = source[:lastColon]
	}
	firstSlash := strings.Index(repositoryWithRegistry, "/")
	if firstSlash <= 0 || !isRegistryHost(repositoryWithRegistry[:firstSlash]) {
		return CoreImage{}, fmt.Errorf("image reference has no registry: %s", source)
	}
	repository := repositoryWithRegistry[firstSlash+1:]
	if repository == "" || tag == "" {
		return CoreImage{}, fmt.Errorf("incomplete image reference: %s", source)
	}
	return CoreImage{Source: source, Repository: repository, Tag: tag}, nil
}

func imageReferenceIsPinned(source string) bool {
	source = strings.TrimSpace(source)
	lastSlash := strings.LastIndex(source, "/")
	if digestIndex := strings.LastIndex(source, "@sha256:"); digestIndex > lastSlash {
		return validSHA256Digest(source[digestIndex+1:])
	}
	lastColon := strings.LastIndex(source, ":")
	if lastColon <= lastSlash || lastColon == len(source)-1 {
		return false
	}
	return !strings.EqualFold(source[lastColon+1:], "latest")
}

func looksLikeExplicitImageReference(value string) bool {
	parsed, err := distributionreference.Parse(strings.TrimSpace(value))
	if err != nil {
		return false
	}
	named, ok := parsed.(distributionreference.Named)
	if !ok || !isRegistryHost(distributionreference.Domain(named)) {
		return false
	}
	_, tagged := parsed.(distributionreference.Tagged)
	_, digested := parsed.(distributionreference.Digested)
	return tagged || digested
}

func imageReferenceFromFlag(value string) (string, bool, error) {
	candidate := strings.TrimSpace(value)
	if !strings.HasPrefix(candidate, "--") {
		return "", false, nil
	}
	flag, assignment, ok := strings.Cut(candidate, "=")
	flagName := strings.ToLower(strings.TrimPrefix(flag, "--"))
	if flagName != "image" && !strings.HasSuffix(flagName, "-image") {
		return "", false, nil
	}
	if !ok || strings.TrimSpace(assignment) == "" {
		return "", false, fmt.Errorf("image flag %s requires an explicit image reference assignment", flag)
	}
	candidate = strings.TrimSpace(assignment)
	if !looksLikeExplicitImageReference(candidate) {
		return "", false, fmt.Errorf("image flag %s has a malformed image reference: %s", flag, candidate)
	}
	return candidate, true, nil
}

func isImageReferenceField(key string) bool {
	key = strings.ToLower(strings.TrimSpace(key))
	return key == "image" || key == "images" || strings.HasSuffix(key, "image")
}

func qualifyImageRepository(value, registry string) string {
	value = strings.TrimSpace(value)
	firstSlash := strings.Index(value, "/")
	if firstSlash > 0 && isRegistryHost(value[:firstSlash]) {
		return value
	}
	if registry == "" {
		registry = defaultCoreImageRegistry
	}
	return strings.TrimSuffix(registry, "/") + "/" + strings.TrimPrefix(value, "/")
}

func isRegistryHost(value string) bool {
	if strings.HasPrefix(value, "-") || strings.Contains(value, "=") {
		return false
	}
	return value == "localhost" || strings.Contains(value, ".") || strings.Contains(value, ":")
}

func scalarValue(node *yaml.Node) string {
	if node == nil || node.Kind != yaml.ScalarNode {
		return ""
	}
	return strings.TrimSpace(node.Value)
}

func closureDigest(images []CoreImage) string {
	canonical := append([]CoreImage(nil), images...)
	slices.SortFunc(canonical, func(a, b CoreImage) int {
		return strings.Compare(a.Source, b.Source)
	})
	data, _ := json.Marshal(canonical)
	digest := sha256.Sum256(data)
	return "sha256:" + hex.EncodeToString(digest[:])
}