package llm

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

	"gitea.maximumdirect.net/eric/scriptorium"
)

type AssetSource struct {
	FS   fs.FS
	Root string
}

type AssetRegistry struct {
	prompts []AssetSource
	schemas []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
}

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) ScriptoriumOptions() ([]scriptorium.Option, 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)
	}
	return []scriptorium.Option{
		scriptorium.WithPromptFS(promptFS, "."),
		scriptorium.WithSchemaFS(schemaFS, "."),
	}, 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 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 := assetMapFS{}
	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] = append([]byte(nil), data...)
			return nil
		}); err != nil {
			return nil, fmt.Errorf("walk asset root %s: %w", source.Root, err)
		}
	}
	return out, nil
}

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
}

type assetMapFS map[string][]byte

func (m assetMapFS) Open(name string) (fs.File, error) {
	cleaned, err := cleanOpenPath(name)
	if err != nil {
		return nil, &fs.PathError{Op: "open", Path: name, Err: err}
	}
	if data, ok := m[cleaned]; ok {
		return &assetFile{
			reader: bytes.NewReader(data),
			info:   assetFileInfo{name: path.Base(cleaned), size: int64(len(data))},
		}, nil
	}
	entries := m.dirEntries(cleaned)
	if entries != nil {
		return &assetDir{name: path.Base(cleaned), entries: entries}, nil
	}
	return nil, &fs.PathError{Op: "open", Path: name, Err: fs.ErrNotExist}
}

func (m assetMapFS) ReadFile(name string) ([]byte, error) {
	cleaned, err := cleanOpenPath(name)
	if err != nil {
		return nil, &fs.PathError{Op: "readfile", Path: name, Err: err}
	}
	data, ok := m[cleaned]
	if !ok {
		return nil, &fs.PathError{Op: "readfile", Path: name, Err: fs.ErrNotExist}
	}
	return append([]byte(nil), data...), nil
}

func (m assetMapFS) ReadDir(name string) ([]fs.DirEntry, error) {
	cleaned, err := cleanOpenPath(name)
	if err != nil {
		return nil, &fs.PathError{Op: "readdir", Path: name, Err: err}
	}
	entries := m.dirEntries(cleaned)
	if entries == nil {
		return nil, &fs.PathError{Op: "readdir", Path: name, Err: fs.ErrNotExist}
	}
	return entries, nil
}

func (m assetMapFS) dirEntries(dir string) []fs.DirEntry {
	children := map[string]assetDirEntry{}
	prefix := ""
	if dir != "." {
		prefix = dir + "/"
	}
	for name, data := range m {
		if !strings.HasPrefix(name, prefix) {
			continue
		}
		rest := strings.TrimPrefix(name, prefix)
		if rest == "" {
			continue
		}
		childName, _, hasSlash := strings.Cut(rest, "/")
		entry := assetDirEntry{name: childName, dir: hasSlash}
		if !hasSlash {
			entry.size = int64(len(data))
		}
		children[childName] = entry
	}
	if len(children) == 0 {
		return nil
	}
	names := make([]string, 0, len(children))
	for name := range children {
		names = append(names, name)
	}
	sort.Strings(names)
	entries := make([]fs.DirEntry, 0, len(names))
	for _, name := range names {
		entries = append(entries, children[name])
	}
	return entries
}

func cleanOpenPath(name string) (string, error) {
	if name == "." {
		return ".", nil
	}
	return cleanAssetPath(name)
}

type assetFile struct {
	reader *bytes.Reader
	info   assetFileInfo
}

func (f *assetFile) Stat() (fs.FileInfo, error) { return f.info, nil }
func (f *assetFile) Read(p []byte) (int, error) { return f.reader.Read(p) }
func (f *assetFile) Close() error               { return nil }

type assetDir struct {
	name    string
	offset  int
	entries []fs.DirEntry
}

func (d *assetDir) Stat() (fs.FileInfo, error) { return assetFileInfo{name: d.name, dir: true}, nil }
func (d *assetDir) Read([]byte) (int, error)   { return 0, fmt.Errorf("cannot read directory") }
func (d *assetDir) Close() error               { return nil }
func (d *assetDir) ReadDir(n int) ([]fs.DirEntry, error) {
	if d.offset >= len(d.entries) {
		return nil, io.EOF
	}
	end := len(d.entries)
	if n > 0 && d.offset+n < end {
		end = d.offset + n
	}
	out := append([]fs.DirEntry(nil), d.entries[d.offset:end]...)
	d.offset = end
	return out, nil
}

type assetDirEntry struct {
	name string
	dir  bool
	size int64
}

func (e assetDirEntry) Name() string      { return e.name }
func (e assetDirEntry) IsDir() bool       { return e.dir }
func (e assetDirEntry) Type() fs.FileMode { return e.InfoMode().Type() }
func (e assetDirEntry) Info() (fs.FileInfo, error) {
	return assetFileInfo{name: e.name, dir: e.dir, size: e.size}, nil
}
func (e assetDirEntry) InfoMode() fs.FileMode {
	if e.dir {
		return fs.ModeDir | 0o555
	}
	return 0o444
}

type assetFileInfo struct {
	name string
	dir  bool
	size int64
}

func (i assetFileInfo) Name() string       { return i.name }
func (i assetFileInfo) Size() int64        { return i.size }
func (i assetFileInfo) Mode() fs.FileMode  { return assetDirEntry{dir: i.dir}.InfoMode() }
func (i assetFileInfo) ModTime() time.Time { return time.Time{} }
func (i assetFileInfo) IsDir() bool        { return i.dir }
func (i assetFileInfo) Sys() any           { return nil }