// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package cli

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"log/slog"
	"os"
	"slices"
	"strings"
	"text/template"

	"gopkg.in/yaml.v3"

	groupval "github.com/stacklok/toolhive-core/validation/group"
	"github.com/stacklok/toolhive/pkg/fileutils"
	vmcpconfig "github.com/stacklok/toolhive/pkg/vmcp/config"
	"github.com/stacklok/toolhive/pkg/vmcp/workloads"
)

// InitConfig holds all parameters for the Init command.
// Discoverer is injected rather than constructed internally to enable unit testing.
type InitConfig struct {
	// GroupName is the ToolHive group whose workloads are enumerated.
	GroupName string

	// OutputPath is the file path to write the generated config.
	// If empty or "-", content is written to Writer.
	OutputPath string

	// Writer is used when OutputPath is empty or "-".
	// Defaults to os.Stdout when nil.
	Writer io.Writer

	// Discoverer resolves running workloads in the group.
	// In production, callers pass a *pkg/workloads.DiscovererAdapter (CLI) or
	// the Kubernetes discoverer from pkg/vmcp/workloads/k8s.go.
	Discoverer workloads.Discoverer
}

// initTemplateData holds the data for the config template.
type initTemplateData struct {
	ServerName string
	GroupName  string
	Backends   []vmcpconfig.StaticBackendConfig
}

// configTemplate is the starter vMCP YAML template with inline comments.
// Text/template delimiters ({{...}}) do not conflict with YAML's {workload}_
// placeholder because Go templates require double braces.
const configTemplate = "# Generated by `thv vmcp init`. Review and customize before use.\n" +
	`
# name: unique identifier for this vMCP server instance.
name: {{.ServerName}}

# groupRef: the ToolHive group whose workloads are aggregated.
groupRef: {{.GroupName}}

# incomingAuth: controls how clients authenticate to this vMCP server.
# type: anonymous disables client auth (suitable for local development).
# Change to "oidc" for production deployments.
incomingAuth:
  type: anonymous

# outgoingAuth: controls how this vMCP server authenticates to backends.
# source: inline means auth config is fully specified here.
outgoingAuth:
  source: inline

# aggregation: controls how tools from multiple backends are combined.
# conflictResolution: prefix prepends the backend name to each tool name.
aggregation:
  conflictResolution: prefix
  conflictResolutionConfig:
    prefixFormat: "{workload}_"

# backends: one entry per running workload discovered in the group.
backends:{{if .Backends}}
{{range .Backends}}  - name: {{yamlStr .Name}}
    url: {{yamlStr .URL}}
    transport: {{yamlStr .Transport}}
{{end}}{{else}} []
{{end}}`

// parsedConfigTemplate is the configTemplate parsed once at package init with
// the yamlStr function registered, avoiding repeated parsing on every call.
var parsedConfigTemplate = template.Must(
	template.New("vmcp-init").Funcs(template.FuncMap{
		"yamlStr": yamlScalar,
	}).Parse(configTemplate),
)

// Init discovers workloads in cfg.GroupName, renders a starter vMCP YAML
// config file with one backend entry per accessible workload, and writes
// the result to cfg.OutputPath or cfg.Writer.
func Init(ctx context.Context, cfg InitConfig) error {
	if cfg.Discoverer == nil {
		return fmt.Errorf("discoverer is required")
	}
	if err := groupval.ValidateName(cfg.GroupName); err != nil {
		return fmt.Errorf("invalid group name: %w", err)
	}

	workloadList, err := cfg.Discoverer.ListWorkloadsInGroup(ctx, cfg.GroupName)
	if err != nil {
		return fmt.Errorf("failed to list workloads in group %q: %w", cfg.GroupName, err)
	}

	backends, err := resolveBackends(ctx, cfg.Discoverer, workloadList)
	if err != nil {
		return err
	}

	rendered, err := renderConfig(initTemplateData{
		ServerName: cfg.GroupName + "-vmcp",
		GroupName:  cfg.GroupName,
		Backends:   backends,
	})
	if err != nil {
		return err
	}

	return writeOutput(cfg, rendered)
}

// yamlScalar marshals a string to a properly quoted/escaped YAML scalar value
// using gopkg.in/yaml.v3, ensuring characters like '#' or ':' in names and
// URLs do not corrupt the rendered YAML.
func yamlScalar(v string) (string, error) {
	b, err := yaml.Marshal(v)
	if err != nil {
		return "", err
	}
	return strings.TrimRight(string(b), "\n"), nil
}

// normalizeTransport maps known transport aliases to the canonical values accepted
// by the static backend config validator. Returns the canonical value and true,
// or ("", false) for unsupported transports.
func normalizeTransport(t string) (string, bool) {
	switch t {
	case vmcpconfig.TransportSSE:
		return vmcpconfig.TransportSSE, true
	case vmcpconfig.TransportStreamableHTTP, "streamable":
		return vmcpconfig.TransportStreamableHTTP, true
	default:
		return "", false
	}
}

// resolveBackends calls GetWorkloadAsVMCPBackend for each workload and collects
// accessible backends, skipping those that return nil.
func resolveBackends(
	ctx context.Context,
	disc workloads.Discoverer,
	workloadList []workloads.TypedWorkload,
) ([]vmcpconfig.StaticBackendConfig, error) {
	var backends []vmcpconfig.StaticBackendConfig
	for _, wl := range workloadList {
		backend, err := disc.GetWorkloadAsVMCPBackend(ctx, wl)
		if err != nil {
			return nil, fmt.Errorf("failed to get backend for workload %q: %w", wl.Name, err)
		}
		if backend == nil {
			slog.Debug("skipping workload: not yet accessible", "workload", wl.Name)
			continue
		}
		transport, ok := normalizeTransport(backend.TransportType)
		if !ok {
			slog.Warn("skipping workload: unsupported transport type for static config",
				"workload", wl.Name, "transport", backend.TransportType)
			continue
		}
		backends = append(backends, vmcpconfig.StaticBackendConfig{
			Name:      backend.Name,
			URL:       backend.BaseURL,
			Transport: transport,
		})
	}
	slices.SortFunc(backends, func(a, b vmcpconfig.StaticBackendConfig) int {
		return strings.Compare(a.Name, b.Name)
	})
	return backends, nil
}

// renderConfig executes the pre-parsed configTemplate with the given data.
func renderConfig(data initTemplateData) ([]byte, error) {
	var buf bytes.Buffer
	if err := parsedConfigTemplate.Execute(&buf, data); err != nil {
		return nil, fmt.Errorf("failed to render config template: %w", err)
	}
	return buf.Bytes(), nil
}

// writeOutput writes the rendered config to the configured destination.
func writeOutput(cfg InitConfig, content []byte) error {
	if cfg.OutputPath != "" && cfg.OutputPath != "-" {
		if err := fileutils.AtomicWriteFile(cfg.OutputPath, content, 0o600); err != nil {
			return fmt.Errorf("failed to write config to %q: %w", cfg.OutputPath, err)
		}
		slog.Info("vMCP configuration written", "path", cfg.OutputPath)
		return nil
	}

	w := cfg.Writer
	if w == nil {
		w = os.Stdout
	}
	if _, err := w.Write(content); err != nil {
		return fmt.Errorf("failed to write config: %w", err)
	}
	return nil
}
