// Package controller add by blacknc <blacknc@163.com> at 2023/03/13 15:09:57
package controller

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"time"

	"dataservices/commons/constants"
	commonhttp "dataservices/commons/http"
	"dataservices/commons/metrics"
	"dataservices/commons/utils"
	"dataservices/dbhubsdk"
	dtodataapi "dataservices/internal/dto/dataapi"

	dto "dataservices/internal/dto/dataapi_gateway"
	"dataservices/internal/entity"
	service "dataservices/internal/service/dataapi_gateway"

	"github.com/gin-gonic/gin"
	"github.com/tidbcloud/tiinf/errs"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"

	"github.com/tidbcloud/dataservices/idl/pbgen/proto/dataapi"
)

const maxInstructionLength = 1024

func chat2queryEndpoint(c *gin.Context, appInfo *entity.DataApp, endpointInfo *entity.Endpoint, publicKey string) (*dto.AdvancedRes, bool, error) {
	ctx := c.Request.Context()
	returnAllStmtResult := c.GetHeader("return_all_stmt_result")
	endpoint, err := parseEntityEndpointToExec(endpointInfo, appInfo.ProjectId, "", strings.ToLower(returnAllStmtResult) == "true")
	if err != nil {
		metrics.ExecEndpointFailedCounter.WithLabelValues("parse_endpoint_failed").Inc()
		return nil, false, err
	}
	if endpoint.setting.TimeOut < 1 {
		endpoint.setting.TimeOut = constants.ExecSQLDefaultTimeoutMs
	}
	timeout := time.Millisecond * time.Duration(endpoint.setting.TimeOut)
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	tilog.Info(ctx, "handle chat2query endpoint", zap.Any("endpointInfo", endpointInfo))
	switch endpointInfo.Type {
	case entity.EndpointTypeChat2EDA:
		return Chat2QueryV2(c, ctx, appInfo, endpoint, publicKey)
	case entity.EndpointTypeChat2Data:
		return Chat2Data(c, ctx, appInfo, endpoint, publicKey)
	default:
		err := errs.NewBadRequest("Unsupported endpoint type", nil)
		tilog.Error(ctx, "chat2query endpoint failed",
			zap.Error(err),
			zap.String("external_app_id", appInfo.ExternalAppID))
		return nil, false, err
	}
}

func Chat2QueryV2(c *gin.Context, ctx context.Context, appInfo *entity.DataApp, endpointInfo *EndpointExec, publicKey string) (*dto.AdvancedRes, bool, error) {
	tilog.Info(ctx, "Chat2QueryV2", zap.Any("appInfo", appInfo), zap.Any("endpointInfo", endpointInfo), zap.String("publicKey", publicKey))

	timeout := time.Duration(endpointInfo.setting.TimeOut) * time.Millisecond
	resp, httpCode, err := commonhttp.ProxyToEDAService(ctx, appInfo, c.Request, endpointInfo.method, timeout)
	tilog.Info(ctx, "request global eda service", zap.String("method", endpointInfo.method), zap.String("url", c.Request.URL.Path), zap.Int("httpCode", httpCode), zap.Error(err))
	if err != nil {
		if utils.IsTimeoutError(err) {
			err = errs.NewBadRequest(fmt.Sprintf("Call endpoint timeout, timeout is %s", timeout.String()), err)
			return nil, false, err
		}
		return nil, false, err
	}

	c.Status(httpCode)
	c.Header("Content-Type", "application/json")
	c.Writer.WriteString(resp)

	return nil, false, nil
}

func Chat2Data(c *gin.Context, ctx context.Context, appInfo *entity.DataApp, endpointInfo *EndpointExec, publicKey string) (*dto.AdvancedRes, bool, error) {
	startT := time.Now()
	param, err := parseChat2QueryParam(c, true, true)
	if err != nil {
		tilog.Error(ctx, "Chat2Data parse param failed", zap.Error(err))
		return nil, false, err
	}
	tilog.Info(ctx, "Chat2Data start", zap.Any("param", param))

	redirect, err := Redirect(ctx, c, appInfo.ProjectId, param.ClusterID, true)
	if err != nil {
		return nil, false, err
	}
	if redirect {
		return nil, true, nil
	}

	m := map[string]string{
		"cluster_id":  fmt.Sprintf("%d", param.ClusterID),
		"database":    param.Database,
		"instruction": param.Instruction,
	}
	reqBytes, err := json.Marshal(m)
	if err != nil {
		return nil, false, errs.NewInternalJsonMarshal("marshal chat2 query param failed", err)
	}

	timeout := time.Duration(endpointInfo.setting.TimeOut) * time.Millisecond
	c.Request.Body = ioutil.NopCloser(bytes.NewReader(reqBytes))
	respBody, status, err := commonhttp.ProxyToEDAService(ctx, appInfo, c.Request, endpointInfo.method, timeout)
	if err != nil {
		if utils.IsTimeoutError(err) {
			err = errs.NewBadRequest(fmt.Sprintf("Call endpoint timeout, timeout is %s", timeout.String()), err)
			return nil, false, err
		}
		return nil, false, err
	}
	if status != http.StatusOK {
		err = errs.NewInternalHTTP("Call endpoint failed", nil)
		return nil, false, err
	}

	var edaResp commonhttp.EDAJobResponse
	if err := json.Unmarshal([]byte(respBody), &edaResp); err != nil {
		return nil, false, errs.NewInternalJsonMarshal("parse eda job response failed", err)
	}

	respBody, status, err = commonhttp.LoopQueryEDAJob(ctx, appInfo, edaResp.Result.JobID, 120)
	if err != nil {
		return nil, false, err
	}
	if status != http.StatusOK {
		return nil, false, errs.NewInternalHTTP("Call endpoint failed", nil)
	}

	var edaChat2DataResp commonhttp.EDAChat2DataResp
	if err := json.Unmarshal([]byte(respBody), &edaChat2DataResp); err != nil {
		return nil, false, errs.NewInternalJsonMarshal("parse eda job response failed", err)
	}

	genRes := &dtodataapi.InteractResponse{
		Question: param.Instruction,
		Response: edaChat2DataResp.Result.Result.SQL,
		Database: param.Database,
		Code:     0,
	}

	aiCost := time.Since(startT)
	endpointInfo.sqlTemplate = genRes.Response
	endpointInfo.database = genRes.Database
	endpointInfo.clusterId = param.ClusterID

	// check sql read/write mode
	rwMode := dbhubsdk.RWMode(ctx, endpointInfo.sqlTemplate)
	if rwMode == constants.SqlModeINVALID {
		return nil, false, errs.NewBadRequest("Execute endpoint failed, invalid sql", nil)
	}
	terr := service.CheckApiKeyPermission(ctx, endpointInfo.dataappId, endpointInfo.clusterId, rwMode, publicKey)
	if terr != nil {
		if errs.IsError(terr, errs.ErrForbidden) && rwMode == constants.SqlModeReadWrite {
			return nil, false, errs.NewUnauthorized("The SQL generated by AI includes write operations, but the current API key only has read-only access. Please replace it with an API key that has write permissions.", terr)
		}
		return nil, false, terr
	}

	credential := &dataapi.Credential{
		Type: dataapi.CrtType_APIKEY,
		Key:  publicKey,
	}
	req, resp, err := endpointInfo.exec(ctx, nil, credential, appInfo.OrgId, constants.NeedMeasure, "chat2data")
	service.AddExecOplog(ctx, &entity.Chat2QueryParam{
		OrgID:         appInfo.OrgId,
		ProjectID:     appInfo.ProjectId,
		ClusterID:     param.ClusterID,
		Apikey:        publicKey,
		Chat2QuerySid: genRes.SessionID,
	}, endpointInfo.sqlTemplate, resp, err)
	if err != nil {
		return nil, false, err
	}
	totalCost := time.Since(startT)

	result := dto.ConvertDataAPIResToAdvanced(req, resp, constants.EndpointResTypeChat2Data)
	result.Data.Result.Sql = genRes.Response
	result.Data.Result.AiLatency = aiCost.String()
	result.Data.Result.StartMs = startT.UnixMilli()
	result.Data.Result.EndMs = time.Now().UnixMilli()
	result.Data.Result.Latency = totalCost.String()
	tilog.Info(ctx, "Chat2Data end",
		zap.Any("cost", time.Since(startT).String()),
		zap.Any("result", result),
		zap.Any("param", param))
	return result, false, nil
}

func parseChat2QueryParam(c *gin.Context, needInstruction bool, needDatabaseSchema bool) (*entity.Chat2QueryParam, error) {
	req := &entity.Chat2QueryReq{}
	if err := c.ShouldBindJSON(req); err != nil {
		tilog.Error(c.Request.Context(), "invalid params", zap.Any("req", req), zap.Error(err))
		return nil, errs.NewBadRequest("Bind params failed", err)
	}
	if needDatabaseSchema {
		if req.ClusterID == 0 {
			return nil, errs.NewBadRequest("Invalid params, cluster_id is required", nil)
		}
		database := req.Database
		if database == "" {
			return nil, errs.NewBadRequest("Invalid params, database is required", nil)
		}
	}
	if needInstruction {
		if req.Instruction == "" {
			return nil, errs.NewBadRequest("Invalid params, instruction is required", nil)
		}
		if len(req.Instruction) > maxInstructionLength {
			return nil, errs.NewBadRequest("Invalid params, instruction is too long", nil)
		}
	}
	// new added param
	var (
		sampleData     string
		databaseSchema string
	)
	if len(req.SampleData) != 0 {
		jsonBytes, err := json.Marshal(req.SampleData)
		if err != nil {
			tilog.Error(c.Request.Context(), "marshal req.SampleData error", zap.Error(err))
		}
		sampleData = string(jsonBytes)
	}
	if len(req.DatabaseSchema) != 0 {
		jsonBytes, err := json.Marshal(req.DatabaseSchema)
		if err == nil {
			tilog.Error(c.Request.Context(), "marshal req.DatabaseSchema error", zap.Error(err))
		}
		databaseSchema = string(jsonBytes)
	}
	param := &entity.Chat2QueryParam{
		Instruction:    req.Instruction,
		ClusterID:      req.ClusterID,
		Database:       req.Database,
		QuestionNum:    req.QuestionNum,
		Tables:         req.Tables,
		SampleData:     sampleData,
		DatabaseSchema: databaseSchema,
	}
	return param, nil
}