package main

import (
	"encoding/json"
	"errors"
	"os"
	"reflect"
	"regexp"
	"strconv"
	"strings"

	"waifu_gallery/openrouter"
	"waifu_gallery/s3"
)

// STRONGLY recommended to use an uncensored LLM for funnier outputs
type Chatbots struct {
	Endpoint      string `json:"endpoint"`      // required, API endpoint
	Authorization string `json:"authorization"` // required, authorization tken in format `BEARER XXXXXXX`

	// generate the enhanced prompt, 1:1 copy of chatGPT/openrouter's chat API
	// the user input will be appended at the end with the "user" role
	PromptCompletion openrouter.ChatCompletions `json:"promptCompletion"`
	// generate the description for the artpiece
	// Title: "%s" with the user input and Content: "%s" will be appened at the end with the "user" role
	DescriptionCompletion openrouter.ChatCompletions `json:"descriptionCompletion"`
}

type Config struct {
	// Fooocus serverless api endpoint
	RunpodApi struct {
		Endpoint      string `json:"endpoint"`
		Authorization string `json:"authorization"`
	} `json:"runpodApi"`

	// Image storage
	S3 s3.Config `json:"s3"`

	// Optional,LLMs
	Chatbot Chatbots `json:"chatbot"`

	// Where our web server listens
	ListenAddress     string `json:"listenAddress"`
	MaxConcurrentJobs int    `json:"maxConcurrentJobs"` // set this to limit the number of generation jobs that can be set in parallel, 0 to disable

	// Content generation properties
	Loras          []Lora   `json:"loras"`
	AspectRatio    string   `json:"aspectRatio"`    // necessary! used to compute display properties. e.g: 512*512
	ImageSeed      int      `json:"imageSeed"`      // set it to make image generation deterministic
	RandomPrompt   []string `json:"randomPrompt"`   // prepended to base prompt, used to spice things up a bit (e.g multiple artist styles)
	BasePrompt     string   `json:"basePrompt"`     // user prompt will be appended to this
	NegativePrompt string   `json:"negativePrompt"` // self explanatory
	UserPromptBias float64  `json:"userPromptBias"` // how much do we favor the user's input
	GuidanceScale  float64  `json:"GuidanceScale"`  // The guidance scale, default 7.0
	Sharpness      float64  `json:"sharpness"`      // Sharpness, defaults to 2.0
	Steps          int      `json:"steps"`          // Number of generation steps, default 25
	Sampler        string   `json:"sampler"`        // Which sampler to use, default euler_ancestral
	BaseModel      string   `json:"baseModel"`      // Base model, leave it empty to use fooocus' default
	BannedWords    []string `json:"bannedWords"`    // user input will be blocked if it contains one of these

	// Internal app state
	state struct {
		width  int
		height int

		bannedWordsRegex *regexp.Regexp

		canarySet            bool
		chatbotPromptEnabled bool
		descriptionEnabled   bool
	}
}

func parseConfig(name string, cfg *Config) (err error) {
	// read file
	file, err := os.ReadFile(name)
	if err != nil {
		return err
	}

	// unmarshal
	err = json.Unmarshal(file, &cfg)
	if err != nil {
		return err
	}

	// set defaults
	if len(cfg.ListenAddress) == 0 {
		cfg.ListenAddress = "127.0.0.1:8080"
	}
	if len(cfg.AspectRatio) == 0 {
		cfg.AspectRatio = "1024*1024"
	}
	if len(cfg.Sampler) == 0 {
		cfg.Sampler = "euler_ancestral"
	}
	if cfg.UserPromptBias == 0 {
		cfg.UserPromptBias = 0.5
	}
	if cfg.GuidanceScale == 0 {
		cfg.GuidanceScale = 7.0
	}
	if cfg.Steps == 0 {
		cfg.Steps = 25
	}
	if cfg.Sharpness == 0 {
		cfg.Sharpness = 2.0
	}

	// initialize aspect ratio
	parts := strings.Split(cfg.AspectRatio, "*")
	if len(parts) != 2 {
		return errors.New("invalid aspect ratio format (must be something like 1024*1024)")
	}
	cfg.state.width, err = strconv.Atoi(parts[0])
	if err != nil {
		return err
	}
	cfg.state.height, err = strconv.Atoi(parts[1])
	if err != nil {
		return err
	}
	if cfg.state.width <= 0 || cfg.state.height <= 0 {
		return errors.New("invalid aspect ratio (must be something like 1024*1024)")
	}
	// initialize badword regex
	config.state.bannedWordsRegex = regexp.MustCompile(strings.Join(cfg.BannedWords, "|"))

	// check for a bunch of chatbot related states
	if !reflect.ValueOf(cfg.Chatbot.PromptCompletion).IsZero() {
		cfg.state.chatbotPromptEnabled = true
	}
	if !reflect.ValueOf(cfg.Chatbot.DescriptionCompletion).IsZero() {
		cfg.state.descriptionEnabled = true
	}

	return nil
}