// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package server

import (
	"context"
	"fmt"
	"strings"

	yaml "github.com/goccy/go-yaml"
	"github.com/googleapis/genai-toolbox/internal/auth"
	"github.com/googleapis/genai-toolbox/internal/auth/google"
	"github.com/googleapis/genai-toolbox/internal/sources"
	"github.com/googleapis/genai-toolbox/internal/tools"
	"github.com/googleapis/genai-toolbox/internal/util"
)

type ServerConfig struct {
	// Server version
	Version string
	// Address is the address of the interface the server will listen on.
	Address string
	// Port is the port the server will listen on.
	Port int
	// SourceConfigs defines what sources of data are available for tools.
	SourceConfigs SourceConfigs
	// AuthServiceConfigs defines what sources of authentication are available for tools.
	AuthServiceConfigs AuthServiceConfigs
	// ToolConfigs defines what tools are available.
	ToolConfigs ToolConfigs
	// ToolsetConfigs defines what tools are available.
	ToolsetConfigs ToolsetConfigs
	// LoggingFormat defines whether structured loggings are used.
	LoggingFormat logFormat
	// LogLevel defines the levels to log.
	LogLevel StringLevel
	// TelemetryGCP defines whether GCP exporter is used.
	TelemetryGCP bool
	// TelemetryOTLP defines OTLP collector url for telemetry exports.
	TelemetryOTLP string
	// TelemetryServiceName defines the value of service.name resource attribute.
	TelemetryServiceName string
	// Stdio indicates if Toolbox is listening via MCP stdio.
	Stdio bool
	// DisableReload indicates if the user has disabled dynamic reloading for Toolbox.
	DisableReload bool
	// UI indicates if Toolbox UI endpoints (/ui) are available
	UI bool
}

type logFormat string

// String is used by both fmt.Print and by Cobra in help text
func (f *logFormat) String() string {
	if string(*f) != "" {
		return strings.ToLower(string(*f))
	}
	return "standard"
}

// validate logging format flag
func (f *logFormat) Set(v string) error {
	switch strings.ToLower(v) {
	case "standard", "json":
		*f = logFormat(v)
		return nil
	default:
		return fmt.Errorf(`log format must be one of "standard", or "json"`)
	}
}

// Type is used in Cobra help text
func (f *logFormat) Type() string {
	return "logFormat"
}

type StringLevel string

// String is used by both fmt.Print and by Cobra in help text
func (s *StringLevel) String() string {
	if string(*s) != "" {
		return strings.ToLower(string(*s))
	}
	return "info"
}

// validate log level flag
func (s *StringLevel) Set(v string) error {
	switch strings.ToLower(v) {
	case "debug", "info", "warn", "error":
		*s = StringLevel(v)
		return nil
	default:
		return fmt.Errorf(`log level must be one of "debug", "info", "warn", or "error"`)
	}
}

// Type is used in Cobra help text
func (s *StringLevel) Type() string {
	return "stringLevel"
}

// SourceConfigs is a type used to allow unmarshal of the data source config map
type SourceConfigs map[string]sources.SourceConfig

// validate interface
var _ yaml.InterfaceUnmarshalerContext = &SourceConfigs{}

func (c *SourceConfigs) UnmarshalYAML(ctx context.Context, unmarshal func(interface{}) error) error {
	*c = make(SourceConfigs)
	// Parse the 'kind' fields for each source
	var raw map[string]util.DelayedUnmarshaler
	if err := unmarshal(&raw); err != nil {
		return err
	}

	for name, u := range raw {
		// Unmarshal to a general type that ensure it capture all fields
		var v map[string]any
		if err := u.Unmarshal(&v); err != nil {
			return fmt.Errorf("unable to unmarshal %q: %w", name, err)
		}

		kind, ok := v["kind"]
		if !ok {
			return fmt.Errorf("missing 'kind' field for source %q", name)
		}
		kindStr, ok := kind.(string)
		if !ok {
			return fmt.Errorf("invalid 'kind' field for source %q (must be a string)", name)
		}

		yamlDecoder, err := util.NewStrictDecoder(v)
		if err != nil {
			return fmt.Errorf("error creating YAML decoder for source %q: %w", name, err)
		}

		sourceConfig, err := sources.DecodeConfig(ctx, kindStr, name, yamlDecoder)
		if err != nil {
			return err
		}
		(*c)[name] = sourceConfig
	}
	return nil
}

// AuthServiceConfigs is a type used to allow unmarshal of the data authService config map
type AuthServiceConfigs map[string]auth.AuthServiceConfig

// validate interface
var _ yaml.InterfaceUnmarshalerContext = &AuthServiceConfigs{}

func (c *AuthServiceConfigs) UnmarshalYAML(ctx context.Context, unmarshal func(interface{}) error) error {
	*c = make(AuthServiceConfigs)
	// Parse the 'kind' fields for each authService
	var raw map[string]util.DelayedUnmarshaler
	if err := unmarshal(&raw); err != nil {
		return err
	}

	for name, u := range raw {
		var v map[string]any
		if err := u.Unmarshal(&v); err != nil {
			return fmt.Errorf("unable to unmarshal %q: %w", name, err)
		}

		kind, ok := v["kind"]
		if !ok {
			return fmt.Errorf("missing 'kind' field for %q", name)
		}

		dec, err := util.NewStrictDecoder(v)
		if err != nil {
			return fmt.Errorf("error creating decoder: %w", err)
		}
		switch kind {
		case google.AuthServiceKind:
			actual := google.Config{Name: name}
			if err := dec.DecodeContext(ctx, &actual); err != nil {
				return fmt.Errorf("unable to parse as %q: %w", kind, err)
			}
			(*c)[name] = actual
		default:
			return fmt.Errorf("%q is not a valid kind of auth source", kind)
		}
	}
	return nil
}

// ToolConfigs is a type used to allow unmarshal of the tool configs
type ToolConfigs map[string]tools.ToolConfig

// validate interface
var _ yaml.InterfaceUnmarshalerContext = &ToolConfigs{}

func (c *ToolConfigs) UnmarshalYAML(ctx context.Context, unmarshal func(interface{}) error) error {
	*c = make(ToolConfigs)
	// Parse the 'kind' fields for each source
	var raw map[string]util.DelayedUnmarshaler
	if err := unmarshal(&raw); err != nil {
		return err
	}

	for name, u := range raw {
		var v map[string]any
		if err := u.Unmarshal(&v); err != nil {
			return fmt.Errorf("unable to unmarshal %q: %w", name, err)
		}

		// `authRequired` and `useClientOAuth` cannot be specified together
		if v["authRequired"] != nil && v["useClientOAuth"] == true {
			return fmt.Errorf("`authRequired` and `useClientOAuth` are mutually exclusive. Choose only one authentication method")
		}

		// Make `authRequired` an empty list instead of nil for Tool manifest
		if v["authRequired"] == nil {
			v["authRequired"] = []string{}
		}

		kindVal, ok := v["kind"]
		if !ok {
			return fmt.Errorf("missing 'kind' field for tool %q", name)
		}
		kindStr, ok := kindVal.(string)
		if !ok {
			return fmt.Errorf("invalid 'kind' field for tool %q (must be a string)", name)
		}

		yamlDecoder, err := util.NewStrictDecoder(v)
		if err != nil {
			return fmt.Errorf("error creating YAML decoder for tool %q: %w", name, err)
		}

		toolCfg, err := tools.DecodeConfig(ctx, kindStr, name, yamlDecoder)
		if err != nil {
			return err
		}
		(*c)[name] = toolCfg
	}
	return nil
}

// ToolConfigs is a type used to allow unmarshal of the toolset configs
type ToolsetConfigs map[string]tools.ToolsetConfig

// validate interface
var _ yaml.InterfaceUnmarshalerContext = &ToolsetConfigs{}

func (c *ToolsetConfigs) UnmarshalYAML(ctx context.Context, unmarshal func(interface{}) error) error {
	*c = make(ToolsetConfigs)

	var raw map[string][]string
	if err := unmarshal(&raw); err != nil {
		return err
	}

	for name, toolList := range raw {
		(*c)[name] = tools.ToolsetConfig{Name: name, ToolNames: toolList}
	}
	return nil
}
