package llm

import (
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"io/fs"
	"path"
	"strings"

	"gitea.maximumdirect.net/eric/notarius/internal/framework/readonlyfs"
	"gitea.maximumdirect.net/eric/promptkit"
)

type AssetSource struct {
	FS   fs.FS
	Root string
}

type AssetRegistry struct {
	prompts          []AssetSource
	schemas          []AssetSource
	fallbackProfiles []AssetSource
}

type AssetHashPart struct {
	FS   fs.FS
	Path string
}

func NewAssetRegistry() *AssetRegistry {
	return &AssetRegistry{}
}

func (r *AssetRegistry) RegisterPromptFS(fsys fs.FS, root string) error {
	if r == nil {
		return fmt.Errorf("asset registry must not be nil")
	}
	source, err := newAssetSource(fsys, root)
	if err != nil {
		return fmt.Errorf("register prompt assets: %w", err)
	}
	r.prompts = append(r.prompts, source)
	return nil
}

func (r *AssetRegistry) RegisterSchemaFS(fsys fs.FS, root string) error {
	if r == nil {
		return fmt.Errorf("asset registry must not be nil")
	}
	source, err := newAssetSource(fsys, root)
	if err != nil {
		return fmt.Errorf("register schema assets: %w", err)
	}
	r.schemas = append(r.schemas, source)
	return nil
}

// RegisterFallbackProfileFS registers profile assets that PromptKit uses only
// when an operator-configured source does not provide a matching profile.
func (r *AssetRegistry) RegisterFallbackProfileFS(fsys fs.FS, root string) error {
	if r == nil {
		return fmt.Errorf("asset registry must not be nil")
	}
	source, err := newAssetSource(fsys, root)
	if err != nil {
		return fmt.Errorf("register fallback profile assets: %w", err)
	}
	r.fallbackProfiles = append(r.fallbackProfiles, source)
	return nil
}

func (r *AssetRegistry) PromptFS() (fs.FS, error) {
	if r == nil {
		return nil, fmt.Errorf("asset registry must not be nil")
	}
	return flattenAssetSources(r.prompts)
}

func (r *AssetRegistry) SchemaFS() (fs.FS, error) {
	if r == nil {
		return nil, fmt.Errorf("asset registry must not be nil")
	}
	return flattenAssetSources(r.schemas)
}

func (r *AssetRegistry) FallbackProfileFS() (fs.FS, error) {
	if r == nil {
		return nil, fmt.Errorf("asset registry must not be nil")
	}
	return flattenAssetSources(r.fallbackProfiles)
}

// FallbackProfileDigest returns a deterministic, non-secret identity for the
// flattened fallback profile assets.
func (r *AssetRegistry) FallbackProfileDigest() (string, error) {
	_, digest, _, err := r.fallbackProfileAssets()
	return digest, err
}

func (r *AssetRegistry) PromptKitOptions() ([]promptkit.Option, error) {
	options, _, err := r.promptKitOptions()
	return options, err
}

func (r *AssetRegistry) promptKitOptions() ([]promptkit.Option, string, error) {
	promptFS, err := r.PromptFS()
	if err != nil {
		return nil, "", fmt.Errorf("prepare prompt assets: %w", err)
	}
	schemaFS, err := r.SchemaFS()
	if err != nil {
		return nil, "", fmt.Errorf("prepare schema assets: %w", err)
	}
	options := []promptkit.Option{
		promptkit.WithPromptFS(promptFS, "."),
		promptkit.WithSchemaFS(schemaFS, "."),
	}
	fallbackFS, fallbackDigest, hasFallback, err := r.fallbackProfileAssets()
	if err != nil {
		return nil, "", err
	}
	if hasFallback {
		options = append(options, promptkit.WithFallbackProfileFS(fallbackFS, "."))
	}
	return options, fallbackDigest, nil
}

func (r *AssetRegistry) promptKitFallbackProfileOption() (promptkit.Option, bool, error) {
	fallbackFS, _, hasFallback, err := r.fallbackProfileAssets()
	if err != nil {
		return nil, false, err
	}
	if !hasFallback {
		return nil, false, nil
	}
	return promptkit.WithFallbackProfileFS(fallbackFS, "."), true, nil
}

func (r *AssetRegistry) fallbackProfileAssets() (fs.FS, string, bool, error) {
	if r == nil {
		return nil, "", false, fmt.Errorf("asset registry must not be nil")
	}
	if len(r.fallbackProfiles) == 0 {
		empty := sha256.Sum256([]byte("notarius:fallback-profile-assets:empty"))
		return nil, "sha256:" + hex.EncodeToString(empty[:]), false, nil
	}
	fallbackFS, err := r.FallbackProfileFS()
	if err != nil {
		return nil, "", false, fmt.Errorf("prepare fallback profile assets: %w", err)
	}
	digest, err := hashAssetFilesystem(fallbackFS)
	if err != nil {
		return nil, "", false, err
	}
	return fallbackFS, digest, true, nil
}

func HashAssets(parts []AssetHashPart) (string, error) {
	if len(parts) == 0 {
		return "", fmt.Errorf("asset hash requires at least one part")
	}
	hash := sha256.New()
	for _, part := range parts {
		cleanPath, err := cleanAssetPath(part.Path)
		if err != nil {
			return "", fmt.Errorf("hash asset %q: %w", part.Path, err)
		}
		data, err := fs.ReadFile(part.FS, cleanPath)
		if err != nil {
			return "", fmt.Errorf("read hash asset %s: %w", cleanPath, err)
		}
		if _, err := io.WriteString(hash, cleanPath); err != nil {
			return "", err
		}
		if _, err := hash.Write([]byte{0}); err != nil {
			return "", err
		}
		if _, err := hash.Write(data); err != nil {
			return "", err
		}
		if _, err := hash.Write([]byte{0}); err != nil {
			return "", err
		}
	}
	return "sha256:" + hex.EncodeToString(hash.Sum(nil)), nil
}

func hashAssetFilesystem(fsys fs.FS) (string, error) {
	var parts []AssetHashPart
	err := fs.WalkDir(fsys, ".", func(name string, entry fs.DirEntry, walkErr error) error {
		if walkErr != nil {
			return walkErr
		}
		if entry.IsDir() {
			return nil
		}
		parts = append(parts, AssetHashPart{FS: fsys, Path: name})
		return nil
	})
	if err != nil {
		return "", fmt.Errorf("walk assets for digest: %w", err)
	}
	if len(parts) > 0 {
		return HashAssets(parts)
	}
	empty := sha256.Sum256([]byte("notarius:fallback-profile-assets:empty"))
	return "sha256:" + hex.EncodeToString(empty[:]), nil
}

func newAssetSource(fsys fs.FS, root string) (AssetSource, error) {
	if fsys == nil {
		return AssetSource{}, fmt.Errorf("filesystem must not be nil")
	}
	cleanRoot, err := cleanAssetRoot(root)
	if err != nil {
		return AssetSource{}, err
	}
	return AssetSource{FS: fsys, Root: cleanRoot}, nil
}

func flattenAssetSources(sources []AssetSource) (fs.FS, error) {
	out := make(map[string][]byte)
	for _, source := range sources {
		if err := fs.WalkDir(source.FS, source.Root, func(name string, entry fs.DirEntry, walkErr error) error {
			if walkErr != nil {
				return walkErr
			}
			if entry.IsDir() {
				return nil
			}
			rel := name
			if source.Root != "." {
				rel = strings.TrimPrefix(name, source.Root+"/")
			}
			rel, err := cleanAssetPath(rel)
			if err != nil {
				return err
			}
			if _, exists := out[rel]; exists {
				return fmt.Errorf("duplicate asset path %q", rel)
			}
			data, err := fs.ReadFile(source.FS, name)
			if err != nil {
				return err
			}
			out[rel] = data
			return nil
		}); err != nil {
			return nil, fmt.Errorf("walk asset root %s: %w", source.Root, err)
		}
	}
	return readonlyfs.New(out)
}

func cleanAssetRoot(root string) (string, error) {
	trimmed := strings.TrimSpace(root)
	if trimmed == "" || trimmed == "." {
		return ".", nil
	}
	return cleanAssetPath(trimmed)
}

func cleanAssetPath(name string) (string, error) {
	trimmed := strings.TrimSpace(name)
	if trimmed == "" {
		return "", fmt.Errorf("path must not be empty")
	}
	cleaned := path.Clean(strings.TrimPrefix(trimmed, "/"))
	if cleaned == "." || !fs.ValidPath(cleaned) {
		return "", fmt.Errorf("invalid path %q", name)
	}
	return cleaned, nil
}