// Package modules describes the optional AI modules (image, video, speech) that a machine can install: what each one
// needs on disk and in memory, which files to download, what the user can ask for, and how a request becomes a ComfyUI
// workflow. The catalog is embedded in the server (for the pages), the agent (to install and run) and the runner (the only
// place that decides which containers may run), so the three cannot disagree.
package modules

import (
	"encoding/json"
	"fmt"
	"sort"
	"strings"

	"wxch.org/wxch/catalog"
)

// Text holds a string per language ("th", "en").
type Text map[string]string

func (t Text) Get(lang string) string {
	if s := t[lang]; s != "" {
		return s
	}
	if s := t["en"]; s != "" {
		return s
	}
	for _, s := range t {
		return s
	}
	return ""
}

// Engine is a container image plus how to run it. Several capabilities share one engine.
type Engine struct {
	ID         string   `json:"id"`
	Title      Text     `json:"title"`
	Image      string   `json:"image"`       // repository, e.g. nvcr.io/nvidia/pytorch
	Digest     string   `json:"digest"`      // sha256:...  the image is pulled by digest, so it cannot be swapped
	Tag        string   `json:"tag"`         // for people only
	ImageBytes int64    `json:"image_bytes"` // size on disk as `docker images` reports it (the download is smaller, about 12 GB)
	WorkBytes  int64    `json:"work_bytes"`  // estimate: source code and Python packages kept in the engine folder
	HostBytes  int64    `json:"host_bytes"`  // estimate: the Python environment of an engine that runs directly on the machine (Python, PyTorch, CUDA libraries)
	SourceRepo string   `json:"source_repo"` // git repository the engine code comes from (empty: the agent ships it)
	SourceTag  string   `json:"source_tag"`  // pinned tag
	ModelsRoot string   `json:"models_root"` // models folder inside the engine folder
	Port       int      `json:"port"`
	Health     string   `json:"health"`  // path that answers 200 when the engine is ready
	Workdir    string   `json:"workdir"` // inside the container; the engine folder is mounted at /work
	Cmd        []string `json:"cmd"`     // run inside the container, in Workdir
	License    string   `json:"license"`
	Note       Text     `json:"note"`
}

// HostPortBase is where the engines are published on 127.0.0.1, one port per engine in catalog order. It lies outside
// the range the runner gives vLLM (8100-8199).
const HostPortBase = 8300

// HostPort is the port on 127.0.0.1 that this engine answers on.
func (c *Catalog) HostPort(engineID string) int {
	for i, e := range c.Engines {
		if e.ID == engineID {
			return HostPortBase + i
		}
	}
	return 0
}

// ImageRef is what is passed to docker: the repository pinned by digest.
func (e Engine) ImageRef() string { return e.Image + "@" + e.Digest }

// ModelFile is one file to download into the engine's models folder.
type ModelFile struct {
	Dest    string `json:"dest"` // relative to the engine's models folder, e.g. checkpoints/x.safetensors
	URL     string `json:"url"`
	Bytes   int64  `json:"bytes"`
	SHA256  string `json:"sha256,omitempty"` // of the whole file; checked after the download (empty for small files that have none)
	License string `json:"license"`
	Source  string `json:"source"` // repository the file comes from
}

type Option struct {
	Value string `json:"value"`
	Label Text   `json:"label"`
}

// Field is one thing the user can fill in. Type: text, int, float, bool, choice, scenes (a list of {prompt, seconds} for a
// story; Min and Max are the number of scenes), or a file (image, video, audio).
type Field struct {
	Name     string   `json:"name"`
	Type     string   `json:"type"`
	Label    Text     `json:"label"`
	Help     Text     `json:"help"`
	Required bool     `json:"required"`
	Default  any      `json:"default"`
	Min      *float64 `json:"min"`
	Max      *float64 `json:"max"`
	Options  []Option `json:"options"`
	// ShowWhen lists the values of the field "operation" for which this field matters (empty = always)
	ShowWhen []string `json:"show_when"`
}

func (f Field) isFile() bool { return f.Type == "image" || f.Type == "video" || f.Type == "audio" }

// Capability is one thing a machine can do for people, installed on its own.
type Capability struct {
	ID        string            `json:"id"` // image.generate, image.edit, video.generate, video.edit, audio.transcribe
	Title     Text              `json:"title"`
	Desc      Text              `json:"desc"`
	Engine    string            `json:"engine"`
	Models    []ModelFile       `json:"models"`
	MemBytes  int64             `json:"mem_bytes"` // estimate of the memory a job needs
	Inputs    []Field           `json:"inputs"`
	Workflows map[string]string `json:"workflows"` // operation (or "default") -> workflow file; empty for engines that are not ComfyUI
	Outputs   []string          `json:"outputs"`   // image, video, text
	MaxInput  int64             `json:"max_input_bytes"`
	Note      Text              `json:"note"`
	Measured  bool              `json:"measured"` // false while the memory figure is an estimate
}

func (c Capability) ModelBytes() int64 {
	var n int64
	for _, m := range c.Models {
		n += m.Bytes
	}
	return n
}

type Catalog struct {
	Updated      string       `json:"updated"`
	Engines      []Engine     `json:"engines"`
	Capabilities []Capability `json:"capabilities"`
}

func (c *Catalog) Engine(id string) (Engine, bool) {
	for _, e := range c.Engines {
		if e.ID == id {
			return e, true
		}
	}
	return Engine{}, false
}

func (c *Catalog) Capability(id string) (Capability, bool) {
	for _, x := range c.Capabilities {
		if x.ID == id {
			return x, true
		}
	}
	return Capability{}, false
}

// Load reads the embedded catalog and checks it for mistakes a typo could cause (a capability that names no engine, a
// workflow that does not exist, a field type that nobody handles).
func Load() (*Catalog, error) {
	var c Catalog
	if err := json.Unmarshal(catalog.ModulesJSON, &c); err != nil {
		return nil, fmt.Errorf("modules.json: %v", err)
	}
	if err := c.check(); err != nil {
		return nil, err
	}
	return &c, nil
}

func (c *Catalog) check() error {
	engines := map[string]bool{}
	for _, e := range c.Engines {
		if e.ID == "" || e.Image == "" || !strings.HasPrefix(e.Digest, "sha256:") || len(e.Cmd) == 0 || e.Port == 0 || !strings.HasPrefix(e.Workdir, "/work") {
			return fmt.Errorf("engine %q is incomplete", e.ID)
		}
		engines[e.ID] = true
	}
	seen := map[string]bool{}
	for _, cp := range c.Capabilities {
		if seen[cp.ID] {
			return fmt.Errorf("capability %q is listed twice", cp.ID)
		}
		seen[cp.ID] = true
		if !engines[cp.Engine] {
			return fmt.Errorf("capability %q names the unknown engine %q", cp.ID, cp.Engine)
		}
		dests := map[string]bool{}
		for _, m := range cp.Models {
			if m.Dest == "" || strings.Contains(m.Dest, "..") || strings.HasPrefix(m.Dest, "/") || !strings.HasPrefix(m.URL, "https://huggingface.co/") || m.Bytes <= 0 {
				return fmt.Errorf("capability %q: bad model entry %q", cp.ID, m.Dest)
			}
			if dests[m.Dest] {
				return fmt.Errorf("capability %q lists %q twice", cp.ID, m.Dest)
			}
			dests[m.Dest] = true
		}
		for op, wf := range cp.Workflows {
			if wf == BuiltinStory {
				continue // built in code (story.go), not a file
			}
			if _, err := catalog.Workflows.ReadFile("workflows/" + wf); err != nil {
				return fmt.Errorf("capability %q: workflow %q (%s) is missing", cp.ID, wf, op)
			}
		}
		names := map[string]bool{}
		for _, f := range cp.Inputs {
			switch f.Type {
			case "text", "int", "float", "bool", "choice", "scenes", "image", "video", "audio":
			default:
				return fmt.Errorf("capability %q: field %q has the unknown type %q", cp.ID, f.Name, f.Type)
			}
			if f.Name == "" || names[f.Name] {
				return fmt.Errorf("capability %q: field name %q is empty or repeated", cp.ID, f.Name)
			}
			names[f.Name] = true
			if f.Type == "choice" && len(f.Options) == 0 {
				return fmt.Errorf("capability %q: choice %q has no options", cp.ID, f.Name)
			}
		}
	}
	return nil
}

// Sizes says how much disk the chosen capabilities need: the model files (a file two capabilities share counts once) and
// the engines they run on (an image shared by two engines counts once). The engine figures are lower bounds: the image is
// counted compressed, the packages are an estimate.
type Sizes struct {
	Models  int64 `json:"models"`
	Engines int64 `json:"engines"`
}

func (s Sizes) Total() int64 { return s.Models + s.Engines }

func (c *Catalog) Sizes(capIDs []string) Sizes {
	var s Sizes
	files := map[string]bool{}
	images := map[string]bool{}
	works := map[string]bool{}
	for _, id := range capIDs {
		cp, ok := c.Capability(id)
		if !ok {
			continue
		}
		for _, m := range cp.Models {
			key := cp.Engine + "/" + m.Dest
			if !files[key] {
				files[key] = true
				s.Models += m.Bytes
			}
		}
		if e, ok := c.Engine(cp.Engine); ok {
			if !images[e.Digest] {
				images[e.Digest] = true
				s.Engines += e.ImageBytes
			}
			if !works[e.ID] {
				works[e.ID] = true
				s.Engines += e.WorkBytes
			}
		}
	}
	return s
}

// IDs returns every capability id, sorted.
func (c *Catalog) IDs() []string {
	var out []string
	for _, x := range c.Capabilities {
		out = append(out, x.ID)
	}
	sort.Strings(out)
	return out
}