package controller

import (
	"context"
	"encoding/json"
	"fmt"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/tidbcloud/tiinf/errs"
	"google.golang.org/grpc"

	"dataservices/commons/config"
	"dataservices/commons/constants"
	"dataservices/commons/logger"
	"dataservices/commons/metrics"
	"dataservices/commons/sessioncache"
	"dataservices/commons/utils"
	"dataservices/internal/dto/common"
	dto "dataservices/internal/dto/dataapi_gateway"
	"dataservices/internal/entity"

	"github.com/pingcap/errors"
	"github.com/tidbcloud/dataservices/commons/bizerr"
	"github.com/tidbcloud/dataservices/idl/pbgen/proto/dataapi"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
)

var (
	GlobalDataAPIService *dataAPIGatewayService
	sqlCommentsRegexp    = regexp.MustCompile(constants.SqlCommentsRegexp)
)

type dataAPIGatewayService struct {
	dataapi.UnimplementedDataAPIGatewayServiceServer
	cfg config.GrpcConfig

	executorClient   *ExecutorClient
	dispatchStrategy ExecutorDispatchStrategy
}

func CreateDataAPIGatewayService(cfg config.GrpcConfig) (dataapi.DataAPIGatewayServiceServer, error) {
	server := &dataAPIGatewayService{
		cfg: cfg,
	}
	err := server.init()
	if err != nil {
		return nil, err
	}
	GlobalDataAPIService = server
	return GlobalDataAPIService, nil
}

func (c *dataAPIGatewayService) init() error {
	var err error
	// AliCloud V2 temporary assumption: ExecutorSvcAddr points to a service with
	// exactly one executor instance behind it, so the default connection is enough.
	c.executorClient, err = newExecutorClient(c.cfg.ExecutorSvcAddr)
	if err != nil {
		return errors.New(fmt.Sprintf("connect to executor failed: %v", err))
	}
	c.dispatchStrategy = NewExecutorDispatchStrategy()
	return nil
}

func (c *dataAPIGatewayService) TestEndpointV2(ctx context.Context, in *dataapi.TestEndpointReqV2) (*dataapi.ExecSqlRespV2, error) {
	logRequest(ctx, in, "TestEndpointV2 begin")
	entityEndpoint, err := GetTestEndpointById(ctx, in.GetEndpointId())
	if err != nil {
		metrics.ExecEndpointFailedCounter.WithLabelValues("get_endpoint_failed").Inc()
		tilog.Error(ctx, "TestEndpointV2 get endpoint failed", zap.Error(err))
		return &dataapi.ExecSqlRespV2{BaseResp: common.NewBaseRespByError(err)}, nil
	}

	_, argValues, err := buildEndpointArgs("", in.ArgValues)
	if err != nil {
		metrics.ExecEndpointFailedCounter.WithLabelValues("parse_endpoint_failed").Inc()
		tilog.Error(ctx, "TestEndpointV2 failed", zap.Error(err))
		return &dataapi.ExecSqlRespV2{BaseResp: common.NewBaseRespByError(err)}, nil
	}

	if err := parseParamsWithEndpoint(ctx, entityEndpoint, argValues); err != nil {
		metrics.ExecEndpointFailedCounter.WithLabelValues("parse_args_failed").Inc()
		tilog.Error(ctx, "TestEndpointV2 failed", zap.Error(err))
		return &dataapi.ExecSqlRespV2{BaseResp: common.NewBaseRespByError(err)}, nil
	}

	endpoint, err := parseEntityEndpointToExec(entityEndpoint, in.ProjectId, "", in.ReturnAllStmtResult)
	if err != nil {
		metrics.ExecEndpointFailedCounter.WithLabelValues("parse_endpoint_failed").Inc()
		tilog.Error(ctx, "TestEndpointV2 parse endpoint failed", zap.Error(err))
		return &dataapi.ExecSqlRespV2{BaseResp: common.NewBaseRespByError(err)}, nil
	}

	_, resp, err1 := ExecEndpoints(ctx, endpoint, argValues, in.Credential, &ExecEndpointOption{
		RequestMethod: endpoint.method,
		EndpointType:  "draft",
	})
	if err1 != nil {
		tilog.Error(ctx, "TestEndpointV2 failed", zap.Error(err1))
		return &dataapi.ExecSqlRespV2{BaseResp: common.NewBaseRespByError(err1)}, nil
	}
	if resp != nil && resp.BaseResp != nil && resp.BaseResp.ErrCode != 0 && resp.BaseResp.ErrMsg != "" {
		tilog.Error(ctx, "TestEndpointV2 failed", zap.Int64("err-code", resp.BaseResp.ErrCode), zap.String("error", resp.BaseResp.ErrMsg))
	} else {
		tilog.Info(ctx, "TestEndpointV2 success")
	}
	return resp, nil
}

func (c *dataAPIGatewayService) ExecSqlTemplateV2(ctx context.Context, in *dataapi.ExecSqlTemplateReqV2) (*dataapi.ExecSqlRespV2, error) {
	logRequest(ctx, in, "ExecSqlTemplateV2 begin")

	args, argValues, err := buildEndpointArgs(in.Args, in.ArgValues)
	if err != nil {
		metrics.ExecEndpointFailedCounter.WithLabelValues("parse_endpoint_failed").Inc()
		logger.LogWithError(ctx, err, "ExecSqlTemplateV2 failed", zap.String("region", config.Global().AwsConf.Region))
		return &dataapi.ExecSqlRespV2{BaseResp: common.NewBaseRespByError(err)}, nil
	}

	endpoint := &EndpointExec{
		projectId:           in.ProjectId,
		clusterId:           in.ClusterId,
		sessionId:           in.SessionId,
		database:            in.Database,
		sqlTemplate:         in.SqlTemplate,
		argsList:            args,
		setting:             dto.ConvertEndpointSettingFromProto(in.Settings),
		returnAllStmtResult: in.ReturnAllStmtResult,
		batch:               len(argValues) > 1,
	}
	_, resp, err1 := ExecEndpoints(ctx, endpoint, argValues, in.Credential, &ExecEndpointOption{
		EndpointType: "draft",
		Source:       in.Source,
	})
	if err1 != nil {
		logger.LogWithError(ctx, err1, "ExecSqlTemplateV2 failed", zap.String("region", config.Global().AwsConf.Region))
		return &dataapi.ExecSqlRespV2{BaseResp: common.NewBaseRespByError(err1)}, nil
	}
	if resp != nil && resp.BaseResp != nil && resp.BaseResp.ErrCode != 0 && resp.BaseResp.ErrMsg != "" {
		logger.Warn(ctx, "ExecSqlTemplateV2 failed", zap.Int64("err-code", resp.BaseResp.ErrCode), zap.String("error", resp.BaseResp.ErrMsg), zap.String("region", config.Global().AwsConf.Region))
	} else {
		logger.Info(ctx, "ExecSqlTemplateV2 success")
	}
	return resp, nil
}

func (c *dataAPIGatewayService) ExecSql(ctx context.Context, in *dataapi.BasicEndpointReqV2) (resp *dataapi.ExecSqlRespV2, err error) {
	start := time.Now()
	defer func() {
		cost := time.Since(start)
		metrics.ExecutionLinkDuration.WithLabelValues("endpoint_exec", "exec_sql").Observe(cost.Seconds())
		logger.LogWithError(ctx, err, "ExecSql completed", zap.Duration("cost", time.Since(start)))
	}()
	logRequest(ctx, in, "ExecSql begin")
	if err := c.dispatchStrategy.PrepareExecReq(ctx, in); err != nil {
		return nil, err
	}
	callOpts := c.dispatchStrategy.ExecCallOptions(in)
	c.executorClient.Exec(ctx, func(ctx context.Context, conn *grpc.ClientConn) error {
		executorClient := dataapi.NewDataAPIExecutorServiceClient(conn)
		resp, err = executorClient.ExecSql(ctx, in)
		//logger.Info(ctx, "ExecSql bytes_count((c *dataAPIGatewayService) ExecSql)", zap.Any("resp", resp))
		return err
	}, callOpts...)
	if err != nil {
		logger.LogWithError(ctx, err, "ExecSql failed")
		if strings.Contains(err.Error(), "context deadline") {
			return nil, errs.NewBadRequest("Execute endpoint sql timeout, "+err.Error(), err)
		}
		if strings.Contains(err.Error(), "received message larger than max") {
			return nil, errs.NewBadRequest("The response message exceeds the maximum size limit of 8MB.", err)
		}
		if _, ok := err.(errs.Error); !ok {
			return nil, errs.NewInternalGRPC("failed to execute sql", err)
		}
	} else if resp != nil && resp.BaseResp != nil && resp.BaseResp.ErrCode != 0 && resp.BaseResp.ErrMsg != "" {
		logger.Warn(ctx, "ExecSql failed", zap.Any("base_resp", resp.BaseResp))
	} else {
		logger.Info(ctx, "ExecSql success")
	}
	return resp, err
}

func (c *dataAPIGatewayService) GetDatabaseMeta(ctx context.Context, in *dataapi.GetDatabaseMetaReq) (*dataapi.GetDatabaseMetaResp, error) {
	logRequest(ctx, in, "GetDatabaseMeta begin")
	executorClient := dataapi.NewDataAPIExecutorServiceClient(c.executorClient.Get())
	resp, err := executorClient.GetDatabaseMeta(ctx, in)
	if err != nil {
		tilog.Error(ctx, "GetDatabaseMeta failed", zap.Error(err))
	} else if resp != nil && resp.BaseResp != nil && resp.BaseResp.ErrCode != 0 && resp.BaseResp.ErrMsg != "" {
		tilog.Error(ctx, "GetDatabaseMeta failed", zap.Any("base_resp", resp.BaseResp))
	} else {
		tilog.Info(ctx, "GetDatabaseMeta success")
	}
	return resp, err
}

func (c *dataAPIGatewayService) CreateSession(ctx context.Context, in *dataapi.CreateSessionReq) (*dataapi.CreateSessionResp, error) {
	logRequest(ctx, in, "CreateSession begin")
	executorClient := dataapi.NewDataAPIExecutorServiceClient(c.executorClient.Get())
	resp, err := executorClient.CreateSession(ctx, in)
	if err != nil {
		tilog.Error(ctx, "CreateSession failed", zap.Error(err))
	} else if resp != nil && resp.BaseResp != nil && resp.BaseResp.ErrCode != 0 && resp.BaseResp.ErrMsg != "" {
		tilog.Error(ctx, "CreateSession failed", zap.Any("base_resp", resp.BaseResp))
	} else {
		tilog.Info(ctx, "CreateSession success")
	}
	return resp, err
}

func (c *dataAPIGatewayService) CloseSession(ctx context.Context, in *dataapi.CloseSessionReq) (resp *dataapi.CloseSessionResp, err error) {
	logRequest(ctx, in, "CloseSession begin")
	if err := c.dispatchStrategy.PrepareCloseReq(ctx, in); err != nil {
		return nil, err
	}
	callOpts := c.dispatchStrategy.CloseCallOptions(in)
	c.executorClient.Exec(ctx, func(ctx context.Context, conn *grpc.ClientConn) error {
		executorClient := dataapi.NewDataAPIExecutorServiceClient(conn)
		resp, err = executorClient.CloseSession(ctx, in)
		return err
	}, callOpts...)
	if err != nil {
		tilog.Error(ctx, "CloseSession failed", zap.Error(err))
	} else if resp != nil && resp.BaseResp != nil && resp.BaseResp.ErrCode != 0 && resp.BaseResp.ErrMsg != "" {
		tilog.Error(ctx, "CloseSession failed", zap.Any("base_resp", resp.BaseResp))
	} else {
		tilog.Info(ctx, "CloseSession success")
	}
	return resp, err
}

func (c *dataAPIGatewayService) ParseOneStmt(ctx context.Context, in *dataapi.ParseOneStmtReq) (*dataapi.ParseOneStmtResp, error) {
	logRequest(ctx, in, "ParseOneStmt begin")
	executorClient := dataapi.NewDataAPIExecutorServiceClient(c.executorClient.Get())
	resp, err := executorClient.ParseOneStmt(ctx, in)
	if err != nil {
		tilog.Error(ctx, "ParseOneStmt failed", zap.Error(err))
	} else if resp != nil && resp.BaseResp != nil && resp.BaseResp.ErrCode != 0 && resp.BaseResp.ErrMsg != "" {
		tilog.Error(ctx, "ParseOneStmt failed", zap.Any("base_resp", resp.BaseResp))
	} else {
		tilog.Info(ctx, "ParseOneStmt success")
	}
	return resp, err
}

func parseEntityEndpointToExec(in *entity.Endpoint, projectId uint64, sessionId string, returnAllStmtResult bool) (*EndpointExec, errs.Error) {
	argsList := []*entity.EndpointArgs{}
	if len(in.Args) > 0 {
		err := json.Unmarshal([]byte(in.Args), &argsList)
		if err != nil {
			return nil, errs.NewBadRequest("Parsing endpoint arguments failed, Please check if the arguments are correct.", err)
		}
	}

	setting := entity.EndpointSettings{}
	if len(in.Settings) > 0 {
		err := json.Unmarshal([]byte(in.Settings), &setting)
		if err != nil {
			return nil, errs.NewBadRequest("Parsing endpoint settings failed, Please check if the settings are correct.", err)
		}
	}

	return &EndpointExec{
		id:                  in.ID,
		typ:                 in.Type,
		endpoint:            in.Endpoint,
		projectId:           projectId,
		clusterId:           in.ClusterId,
		dataappId:           in.DataappId,
		sessionId:           sessionId,
		database:            in.Database,
		sqlTemplate:         in.SqlTemplate,
		argsList:            argsList,
		setting:             setting,
		returnAllStmtResult: returnAllStmtResult,
		batch:               in.IsBatch(),
		tag:                 in.Tag,
		method:              in.Method,
	}, nil
}

func buildEndpointArgs(args, argValues string) ([]*entity.EndpointArgs, []map[string]any, errs.Error) {
	var argsTarget []*entity.EndpointArgs
	var valuesTarget []map[string]any
	if args != "" {
		if err := utils.FromJson(args, &argsTarget); err != nil {
			return nil, nil, errs.NewBadRequest("Invalid endpoint argument:"+err.Error(), nil)
		}
	}
	if argValues != "" {
		if err := utils.FromJson(argValues, &valuesTarget); err != nil {
			return nil, nil, errs.NewBadRequest("Invalid endpoint argument value:"+err.Error(), nil)
		}
	}
	return argsTarget, valuesTarget, nil
}

type EndpointExec struct {
	id                  uint64
	endpoint            string
	projectId           uint64
	clusterId           uint64
	dataappId           uint64
	sessionId           string
	database            string
	sqlTemplate         string
	argsList            []*entity.EndpointArgs
	setting             entity.EndpointSettings
	returnAllStmtResult bool
	tag                 string
	batch               bool
	method              string
	typ                 string
}

func (endpoint *EndpointExec) cacheKey(ctx context.Context, params map[string]any, externalAppId string) string {
	if len(params) == 0 || len(endpoint.argsList) == 0 {
		return fmt.Sprintf("%s:%s:%s", externalAppId, endpoint.method, endpoint.endpoint)
	}

	keys := make([]string, 0, len(endpoint.argsList))
	for _, arg := range endpoint.argsList {
		keys = append(keys, arg.Name)
	}

	if endpoint.setting.EnablePagination == 1 {
		keys = append(keys, constants.EndpointPageParam, constants.EndpointPageSizeParam)
	}
	sort.Strings(keys)

	queryStrings := make([]string, 0, len(endpoint.argsList))
	for _, key := range keys {
		if _, ok := params[key]; !ok {
			continue
		}
		queryStrings = append(queryStrings, fmt.Sprintf("%s=%v", key, params[key]))
	}
	path := fmt.Sprintf("%s?%s", endpoint.endpoint, strings.Join(queryStrings, "&"))
	return fmt.Sprintf("%s:%s:%s", externalAppId, endpoint.method, path)
}

func (endpoint *EndpointExec) exec(
	ctx context.Context,
	params map[string]any,
	credential *dataapi.Credential,
	orgId uint64,
	measureType string,
	source string,
) (*dataapi.BasicEndpointReqV2, *dataapi.ExecSqlRespV2, errs.Error) {
	start := time.Now()
	defer func() {
		cost := time.Since(start)
		metrics.ExecutionLinkDuration.WithLabelValues("endpoint_exec", "exec_direct").Observe(cost.Seconds())
	}()

	var (
		sqlTemplate string
		args        []entity.SqlArgValue
		err         error
	)
	sqlTemplate = sqlCommentsRegexp.ReplaceAllString(endpoint.sqlTemplate, "")
	if !entity.IsSystemDataEndpoint(endpoint.typ) {
		sqlTemplate, args, err = BuildEndpoint(sqlTemplate, endpoint.argsList, params)
	}

	if err != nil {
		metrics.ExecEndpointFailedCounter.WithLabelValues("build_endpoint_failed").Inc()
		return nil, nil, errs.NewError(bizerr.ErrInvalidSqlParam, "Parse endpoint sql failed, "+err.Error(), err)
	}

	var page, pageSize uint64
	pageStr, _ := params[constants.EndpointPageParam].(string)
	pageSizeStr, _ := params[constants.EndpointPageSizeParam].(string)
	if pageStr != "" && pageSizeStr != "" && endpoint.setting.EnablePagination == 1 {
		var err1, err2 error
		page, err1 = strconv.ParseUint(pageStr, 10, 64)
		pageSize, err2 = strconv.ParseUint(pageSizeStr, 10, 64)
		if err1 != nil || err2 != nil {
			return nil, nil, errs.NewBadRequest("Failed to parse paging parameters", fmt.Errorf("err1: %v, err2: %v", err1, err2))
		}
		if page == 0 || pageSize == 0 {
			return nil, nil, errs.NewBadRequest("Invalid paging parameters", nil)
		}
		if pageSize > uint64(endpoint.setting.RowsLimit) {
			return nil, nil, errs.NewBadRequest("The page size has reached the maximum number of rows limit", nil)
		}
	}

	req := &dataapi.BasicEndpointReqV2{
		ProjectId:           endpoint.projectId,
		ClusterId:           endpoint.clusterId,
		SessionId:           endpoint.sessionId,
		Database:            endpoint.database,
		Sql:                 sqlTemplate,
		Args:                utils.ToJson(endpoint.argsList),
		ArgValues:           utils.ToJson(args),
		LimitRows:           uint64(endpoint.setting.RowsLimit),
		TimeoutMs:           uint64(endpoint.setting.TimeOut),
		Credential:          credential,
		ReturnAllStmtResult: endpoint.returnAllStmtResult,
		DataappId:           endpoint.dataappId,
		EndpointName:        endpoint.endpoint,
		AccessKey:           credential.Key,
		AppType:             measureType,
		OrgId:               orgId,
		Page:                page,
		PageSize:            pageSize,
		Source:              source,
		EndpointType:        endpoint.typ,
	}
	if req.TimeoutMs < 1 {
		req.TimeoutMs = constants.ExecSQLDefaultTimeoutMs
	}
	if req.LimitRows < 1 {
		req.LimitRows = constants.ExecSQLDefaultRowlimit
	}

	timeout := time.Millisecond * time.Duration(req.TimeoutMs)
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	start = time.Now()
	resp, err := GlobalDataAPIService.ExecSql(ctx, req)
	metrics.RpcCostHistogram.WithLabelValues("dataapi_gateway", "ExecSql").Observe(time.Since(start).Seconds())
	if err != nil {
		metrics.ExecEndpointFailedCounter.WithLabelValues("exec_endpoint_failed").Inc()
		if strings.Contains(err.Error(), "context deadline") {
			return nil, nil, errs.NewBadRequest(fmt.Sprintf("Execute endpoint sql timeout, timeout is %s", timeout.String()), err)
		}
		if tiErr, ok := errs.AsTiError(err); ok {
			return nil, nil, tiErr
		} else {
			return nil, nil, errs.NewBadRequest("Execute endpoint sql failed.", err)
		}
	} else if resp.BaseResp.ErrCode != 0 && strings.Contains(errs.GetOrigErrMsgFromBaseResp(resp.BaseResp), "context deadline") {
		metrics.ExecEndpointFailedCounter.WithLabelValues("exec_endpoint_failed").Inc()
		tilog.Warn(ctx, "Exec sql timeout", zap.Any("resp", resp))
		return nil, nil, errs.NewBadRequest(fmt.Sprintf("Execute endpoint sql timeout, timeout is %s", timeout.String()), nil)
	} else if resp.BaseResp.ErrCode != 0 && strings.Contains(errs.GetOrigErrMsgFromBaseResp(resp.BaseResp), "session has expired") {
		metrics.ExecEndpointFailedCounter.WithLabelValues("exec_endpoint_failed").Inc()
		tilog.Warn(ctx, "session has expired", zap.Any("resp", resp))
		return nil, nil, errs.NewBadRequest(fmt.Sprintf("session expired,sessionid:%s", req.SessionId), nil)
	}
	return req, resp, nil
}

func logRequest(ctx context.Context, req interface{}, msg string) {
	tilog.Info(ctx, msg, logger.Redact("request", req))
}

func fillSessionForExecReq(ctx context.Context, req *dataapi.BasicEndpointReqV2) error {
	start := time.Now()
	defer func() {
		cost := time.Since(start)
		metrics.ExecutionLinkDuration.WithLabelValues("endpoint_exec", "load_session").Observe(cost.Seconds())
		tilog.Info(ctx, "fillSessionForExecReq completed", zap.String("session_id", req.SessionId), zap.Duration("cost", time.Since(start)))
	}()
	sessionItem, err := getSessionItem(ctx, req.SessionId)
	if err != nil {
		return err
	}
	req.SessionItem = sessionItem
	return nil
}

func fillSessionForCloseReq(ctx context.Context, req *dataapi.CloseSessionReq) error {
	sessionItem, err := getSessionItem(ctx, req.SessionId)
	if err != nil {
		return err
	}
	req.SessionItem = sessionItem
	return nil
}

func getSessionItem(ctx context.Context, sessionID string) (*dataapi.SessionItem, error) {
	if sessionID == "" {
		return nil, nil
	}
	var sessionItem sessioncache.SessionItem
	sessionItem.SessionID = sessionID
	if err := sessionItem.Load(ctx); err != nil {
		if config.Global().RegionalConf.SessionFailFast && !sessioncache.IsEndpointSessionID(sessionID) {
			return nil, err
		} else {
			tilog.Warn(ctx, "get session item by id failed, but will not fail fast", zap.String("sessionId", sessionID), zap.Error(err))
			return nil, nil
		}
	}
	cmds := make([]*dataapi.SessionCmd, 0, len(sessionItem.Cmds))
	for _, cmd := range sessionItem.Cmds {
		cmds = append(cmds, &dataapi.SessionCmd{
			Type:    cmd.Type,
			Content: cmd.Content,
		})
	}
	return &dataapi.SessionItem{
		SessionId:        sessionItem.SessionID,
		EndSec:           sessionItem.EndSec,
		Database:         sessionItem.Database,
		LastActiveTime:   sessionItem.LastActiveTime,
		CrtId:            sessionItem.CrtID,
		ClusterId:        sessionItem.ClusterID,
		ExpireSec:        sessionItem.ExpireSec,
		Cmds:             cmds,
		ExecutorEndpoint: sessionItem.ExecutorEndpoint,
	}, nil
}