/*
  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.

  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 plugins

import (
	"database/sql/driver"
	"fmt"

	"github.com/aws/aws-advanced-go-wrapper/awssql/driver_infrastructure"
	"github.com/aws/aws-advanced-go-wrapper/awssql/error_util"
	"github.com/aws/aws-advanced-go-wrapper/awssql/host_info_util"
	"github.com/aws/aws-advanced-go-wrapper/awssql/plugin_helpers"
	"github.com/aws/aws-advanced-go-wrapper/awssql/utils"
	"github.com/aws/aws-advanced-go-wrapper/awssql/utils/telemetry"
)

type DefaultPlugin struct {
	PluginService       driver_infrastructure.PluginService
	DefaultConnProvider driver_infrastructure.ConnectionProvider
	ConnProviderManager driver_infrastructure.ConnectionProviderManager
}

func (d *DefaultPlugin) GetPluginCode() string {
	return "default" // plugin code is not used
}

func (d *DefaultPlugin) InitHostProvider(
	_ *utils.RWMap[string, string],
	_ driver_infrastructure.HostListProviderService,
	_ func() error) error {
	// Do nothing.
	// It's guaranteed that this plugin is always the last in plugin chain so initHostProviderFunc can be omitted.
	return nil
}

func (d *DefaultPlugin) GetSubscribedMethods() []string {
	return []string{plugin_helpers.ALL_METHODS}
}

func (d *DefaultPlugin) Execute(
	connInvokedOn driver.Conn,
	methodName string,
	executeFunc driver_infrastructure.ExecuteFunc,
	methodArgs ...any) (wrappedReturnValue any, wrappedReturnValue2 any, wrappedOk bool, wrappedErr error) {
	wrappedReturnValue, wrappedReturnValue2, wrappedOk, wrappedErr = executeFunc()
	if wrappedErr != nil {
		return
	}

	if connInvokedOn != nil && connInvokedOn != d.PluginService.GetCurrentConnection() {
		// Is using an old connection, so transaction analysis should be skipped.
		// PluginManager blocks all methods invoked using old connections except for close/abort.
		return
	}

	if utils.DoesOpenTransaction(methodName, methodArgs...) {
		d.PluginService.SetInTransaction(true)
	} else if utils.DoesCloseTransaction(methodName, methodArgs...) {
		d.PluginService.SetInTransaction(false)
	}

	return
}

func (d *DefaultPlugin) Connect(
	hostInfo *host_info_util.HostInfo,
	props *utils.RWMap[string, string],
	isInitialConnection bool,
	_ driver_infrastructure.ConnectFunc) (driver.Conn, error) {
	// It's guaranteed that this plugin is always the last in plugin chain so connectFunc can be ignored.
	connProvider := d.ConnProviderManager.GetConnectionProvider(*hostInfo, props)
	return d.connectInternal(hostInfo, props, connProvider, isInitialConnection)
}

func (d *DefaultPlugin) ForceConnect(
	hostInfo *host_info_util.HostInfo,
	props *utils.RWMap[string, string],
	isInitialConnection bool,
	_ driver_infrastructure.ConnectFunc) (driver.Conn, error) {
	// It's guaranteed that this plugin is always the last in plugin chain so connectFunc can be ignored.
	return d.connectInternal(hostInfo, props, d.DefaultConnProvider, isInitialConnection)
}

func (d *DefaultPlugin) connectInternal(
	hostInfo *host_info_util.HostInfo,
	props *utils.RWMap[string, string],
	connProvider driver_infrastructure.ConnectionProvider,
	isInitialConnection bool) (driver.Conn, error) {
	parentCtx := d.PluginService.GetTelemetryContext()
	telemetryCtx, ctx := d.PluginService.GetTelemetryFactory().OpenTelemetryContext(
		fmt.Sprintf(telemetry.TELEMETRY_CONNECT_INTERNAL, utils.GetStructName(connProvider)), telemetry.NESTED, parentCtx)
	d.PluginService.SetTelemetryContext(ctx)
	defer func() {
		telemetryCtx.CloseContext()
		d.PluginService.SetTelemetryContext(parentCtx)
	}()

	conn, err := connProvider.Connect(hostInfo, props.GetAllEntries(), d.PluginService)
	if err == nil {
		d.PluginService.SetAvailability(hostInfo.AllAliases, host_info_util.AVAILABLE)
		if isInitialConnection {
			d.PluginService.UpdateDialect(conn)
		}
	}
	return conn, err
}

func (d *DefaultPlugin) AcceptsStrategy(strategy string) bool {
	return d.ConnProviderManager.AcceptsStrategy(strategy)
}

func (d *DefaultPlugin) GetHostInfoByStrategy(
	role host_info_util.HostRole,
	strategy string,
	hosts []*host_info_util.HostInfo) (*host_info_util.HostInfo, error) {
	if len(hosts) == 0 {
		return nil, error_util.NewGenericAwsWrapperError(error_util.GetMessage("DefaultConnectionPlugin.noHostsAvailable"))
	}
	return d.ConnProviderManager.GetHostInfoByStrategy(hosts, role, strategy, d.PluginService.GetProperties())
}

func (d *DefaultPlugin) GetHostSelectorStrategy(strategy string) (driver_infrastructure.HostSelector, error) {
	return d.ConnProviderManager.GetHostSelectorStrategy(strategy)
}

func (d *DefaultPlugin) NotifyConnectionChanged(_ map[driver_infrastructure.HostChangeOptions]bool) driver_infrastructure.OldConnectionSuggestedAction {
	return driver_infrastructure.NO_OPINION
}

func (d *DefaultPlugin) NotifyHostListChanged(_ map[string]map[driver_infrastructure.HostChangeOptions]bool) {
	// Do nothing.
}
