package llm

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

const (
	promptKitProfileFingerprintName      = "promptkit_profile_source"
	promptKitLocalBackendFingerprintName = "promptkit_local_backend_target"
	promptKitLocalBackendMarker          = "notarius:promptkit-local-backend:v1"
	// The built-in profile catalog is compiled into this pinned PromptKit
	// release. Update this identity when the dependency is upgraded.
	promptKitBuiltinProfileCatalogID = "promptkit:v0.8.0:builtin-profiles"
)

func promptKitProfileFingerprint(profileDir, profileFile, fallbackProfileDigest string) (CheckpointFingerprint, error) {
	hasher := sha256.New()
	writeFingerprintPart(hasher, []byte(promptKitBuiltinProfileCatalogID))
	writeFingerprintPart(hasher, []byte(strings.TrimSpace(fallbackProfileDigest)))

	switch {
	case strings.TrimSpace(profileFile) != "":
		cleanProfileFile := strings.TrimSpace(profileFile)
		data, err := os.ReadFile(cleanProfileFile)
		if err != nil {
			return CheckpointFingerprint{}, fmt.Errorf("read PromptKit profile file for checkpoint identity")
		}
		writeFingerprintPart(hasher, []byte(filepath.ToSlash(filepath.Base(cleanProfileFile))))
		writeFingerprintPart(hasher, data)
	case strings.TrimSpace(profileDir) != "":
		files, err := promptKitProfileFiles(strings.TrimSpace(profileDir))
		if err != nil {
			return CheckpointFingerprint{}, err
		}
		for _, file := range files {
			writeFingerprintPart(hasher, []byte(file.path))
			writeFingerprintPart(hasher, file.digest)
		}
	}

	return CheckpointFingerprint{
		Name:  promptKitProfileFingerprintName,
		Value: "sha256:" + hex.EncodeToString(hasher.Sum(nil)),
	}, nil
}

func promptKitLocalBackendFingerprint(endpoint string) CheckpointFingerprint {
	hasher := sha256.New()
	writeFingerprintPart(hasher, []byte(promptKitLocalBackendMarker))
	writeFingerprintPart(hasher, []byte(strings.TrimSpace(endpoint)))
	return CheckpointFingerprint{
		Name:  promptKitLocalBackendFingerprintName,
		Value: "sha256:" + hex.EncodeToString(hasher.Sum(nil)),
	}
}

type promptKitProfileFile struct {
	path   string
	digest []byte
}

func promptKitProfileFiles(root string) ([]promptKitProfileFile, error) {
	var files []promptKitProfileFile
	err := filepath.WalkDir(root, func(name string, entry fs.DirEntry, walkErr error) error {
		if walkErr != nil {
			return walkErr
		}
		if entry.IsDir() {
			return nil
		}
		extension := filepath.Ext(entry.Name())
		if extension != ".yaml" && extension != ".yml" {
			return nil
		}
		data, err := os.ReadFile(name)
		if err != nil {
			return err
		}
		relativePath, err := filepath.Rel(root, name)
		if err != nil {
			return err
		}
		sum := sha256.Sum256(data)
		files = append(files, promptKitProfileFile{
			path:   filepath.ToSlash(relativePath),
			digest: append([]byte(nil), sum[:]...),
		})
		return nil
	})
	if err != nil {
		return nil, fmt.Errorf("read PromptKit profile directory for checkpoint identity")
	}
	sort.Slice(files, func(i, j int) bool {
		return files[i].path < files[j].path
	})
	return files, nil
}

func writeFingerprintPart(hasher interface{ Write([]byte) (int, error) }, value []byte) {
	length := []byte(fmt.Sprintf("%d:", len(value)))
	_, _ = hasher.Write(length)
	_, _ = hasher.Write(value)
}