package main

import (
	"context"
	"errors"
	"fmt"
	"os"

	commandkit "gitea.maximumdirect.net/eric/CommandKit"
	configkit "gitea.maximumdirect.net/eric/ConfigKit"
	"gitea.maximumdirect.net/eric/promptkit"
)

type campaign struct {
	PromptID      string `yaml:"prompt_id"`
	PromptVersion string `yaml:"prompt_version"`
	ProfileID     string `yaml:"profile_id"`
	Model         string `yaml:"model"`
}

type appConfiguration struct {
	DefaultCampaign string              `yaml:"default_campaign"`
	Campaigns       map[string]campaign `yaml:"campaigns"`
}

type invocationCache struct {
	application *configkit.Result[appConfiguration]
}

type configContext struct {
	Cache    *invocationCache
	Campaign commandkit.Supplied[string]
}

type definitionContext struct {
	Cache    *invocationCache
	Campaign commandkit.Supplied[string]
}

type restoreInput struct {
	Campaign commandkit.Supplied[string]
	Session  string
}

type session struct {
	ID       string
	Campaign string
	Summary  string
}

type sessionStore struct {
	active  string
	records map[string]session
}

func loadApplicationConfig(_ context.Context, _ commandkit.Selection, cache *invocationCache) (*configkit.Result[appConfiguration], error) {
	if cache != nil && cache.application != nil {
		return cache.application, nil
	}
	result, err := configkit.Load[appConfiguration](configkit.Mapping("composed app config", map[string]any{
		"default_campaign": "spring",
		"campaigns": map[string]any{
			"spring": map[string]any{"prompt_id": "briefing", "prompt_version": "1", "profile_id": "spring", "model": "spring-model"},
			"autumn": map[string]any{"prompt_id": "briefing", "prompt_version": "1", "profile_id": "autumn", "model": "autumn-model"},
		},
	}))
	if err == nil && cache != nil {
		cache.application = result
	}
	return result, err
}

func selectedCampaign(configuration appConfiguration, supplied commandkit.Supplied[string]) (string, campaign, error) {
	name := configuration.DefaultCampaign
	if supplied.Set {
		name = supplied.Value
	}
	value, ok := configuration.Campaigns[name]
	if !ok {
		return "", campaign{}, fmt.Errorf("unknown campaign %q", name)
	}
	return name, value, nil
}

func loadPromptEngine(ctx context.Context, selection commandkit.Selection, cache *invocationCache, selected commandkit.Supplied[string]) (*promptkit.Engine, error) {
	application, err := loadApplicationConfig(ctx, selection, cache)
	if err != nil {
		return nil, err
	}
	name, chosen, err := selectedCampaign(application.Config, selected)
	if err != nil {
		return nil, err
	}
	engineConfig := configkit.Mapping("campaign "+name, map[string]any{
		"profiles": map[string]any{chosen.ProfileID: map[string]any{
			"endpoint": "https://example.invalid/v1",
			"model":    chosen.Model,
			"extra_params": map[string]any{
				"campaign": name,
			},
		}},
		"prompts": map[string]any{chosen.PromptID: map[string]any{
			chosen.PromptVersion: map[string]any{
				"default_profile": chosen.ProfileID,
				"inputs":          []any{map[string]any{"name": "topic", "required": true, "content_type": "text/plain", "description": "Topic to brief"}},
				"messages":        []any{map[string]any{"role": "user", "content": "Write a concise briefing about {{.topic}}."}},
				"output":          map[string]any{"format": "text", "validation_mode": "none"},
			},
		}},
	})
	return promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{Main: &configkit.Input{Source: engineConfig}}})
}

func campaignSelector[C any](set func(*C, commandkit.Supplied[string])) commandkit.Selector[C] {
	return commandkit.Selector[C]{
		Binding: commandkit.ValueFlag[C](commandkit.FlagMeta{Name: "campaign", Usage: "Select a configured campaign."}, nil, func(value string) (string, error) { return value, nil }, set),
		Scopes:  []string{"app", "promptkit"},
	}
}

func newApp() (*commandkit.App, error) {
	app, err := commandkit.New(commandkit.Identity{Name: "composed-example", Version: "1.0"})
	if err != nil {
		return nil, err
	}
	appScope, err := commandkit.NewConfigKitScope[configContext, appConfiguration]("app", "Application campaigns.", func(ctx context.Context, selection commandkit.Selection, input configContext) (*configkit.Result[appConfiguration], error) {
		return loadApplicationConfig(ctx, selection, input.Cache)
	})
	if err != nil {
		return nil, err
	}
	promptScope, err := commandkit.NewPromptKitScope(func(ctx context.Context, selection commandkit.Selection, input configContext) (*promptkit.Engine, error) {
		return loadPromptEngine(ctx, selection, input.Cache, input.Campaign)
	})
	if err != nil {
		return nil, err
	}
	configSuite := commandkit.ConfigSuite[configContext]{
		NewContext: func() configContext { return configContext{Cache: &invocationCache{}} },
		Selectors: []commandkit.Selector[configContext]{campaignSelector(func(input *configContext, value commandkit.Supplied[string]) {
			input.Campaign = value
		})},
		Scopes: []commandkit.ConfigScope[configContext]{appScope, promptScope},
	}
	if err := commandkit.RegisterConfig(app, configSuite); err != nil {
		return nil, err
	}

	definitionSuite, err := commandkit.NewPromptKitSuite(func(ctx context.Context, selection commandkit.Selection, input definitionContext) (*promptkit.Engine, error) {
		return loadPromptEngine(ctx, selection, input.Cache, input.Campaign)
	})
	if err != nil {
		return nil, err
	}
	definitionSuite.NewContext = func() definitionContext { return definitionContext{Cache: &invocationCache{}} }
	definitionSuite.Selectors = []commandkit.Selector[definitionContext]{campaignSelector(func(input *definitionContext, value commandkit.Supplied[string]) {
		input.Campaign = value
	})}
	if err := commandkit.RegisterPromptKit(app, definitionSuite); err != nil {
		return nil, err
	}

	store := &sessionStore{records: map[string]session{"session-42": {ID: "session-42", Campaign: "spring", Summary: "A synthetic planning session."}}}
	restore, err := commandkit.NewCommand(commandkit.CommandSpec[restoreInput]{
		Name:    "restore",
		Summary: "Restore one saved session.",
		Flags: []commandkit.Flag[restoreInput]{commandkit.ValueFlag[restoreInput](commandkit.FlagMeta{Name: "campaign", Usage: "Campaign owning the session."}, nil, func(value string) (string, error) { return value, nil }, func(input *restoreInput, value commandkit.Supplied[string]) {
			input.Campaign = value
		})},
		Arguments: []commandkit.Argument[restoreInput]{commandkit.ValueArgument[restoreInput](commandkit.ArgumentMeta{Name: "ID"}, true, func(value string) (string, error) { return value, nil }, func(input *restoreInput, value commandkit.Supplied[string]) {
			input.Session = value.Value
		})},
		Run: func(invocation commandkit.Invocation, input restoreInput) commandkit.Outcome {
			return restoreSession(invocation, store, input)
		},
	})
	if err != nil {
		return nil, err
	}
	sessionGroup, err := commandkit.Group("session", "Manage saved sessions.", restore)
	if err != nil {
		return nil, err
	}
	if err := app.Register(sessionGroup); err != nil {
		return nil, err
	}
	return app, nil
}

func restoreSession(invocation commandkit.Invocation, store *sessionStore, input restoreInput) commandkit.Outcome {
	application, err := loadApplicationConfig(invocation.Context, commandkit.Selection{}, nil)
	if err != nil {
		return commandkit.Failure(err)
	}
	name, _, err := selectedCampaign(application.Config, input.Campaign)
	if err != nil {
		return commandkit.Failure(err)
	}
	record, ok := store.records[input.Session]
	if !ok {
		return commandkit.Outcome{Kind: commandkit.OutcomeDomain, Status: 4, Domain: "session_not_found", Cause: errors.New("saved session was not found")}
	}
	if record.Campaign != name {
		return commandkit.Outcome{Kind: commandkit.OutcomeDomain, Status: 3, Domain: "campaign_mismatch", Cause: fmt.Errorf("session belongs to campaign %q", record.Campaign)}
	}
	store.active = record.ID
	if _, err := fmt.Fprintf(invocation.Streams.Out, "restored %s: %s\n", record.ID, record.Summary); err != nil {
		return commandkit.Failure(err)
	}
	return commandkit.Success()
}

func main() {
	app, err := newApp()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	outcome := app.Execute(context.Background(), os.Args[1:], commandkit.Streams{Out: os.Stdout, Err: os.Stderr})
	if outcome.Status != 0 {
		os.Exit(outcome.Status)
	}
}