package cluster_svc

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"math"
	"strings"
	"time"

	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/serverless-service/idl/pbgen/proto/console"
	"github.com/tidbcloud/serverless-service/internal/domain/facade"
	"github.com/tidbcloud/serverless-service/internal/errorcode"
	"go.uber.org/zap"
	"google.golang.org/protobuf/types/known/timestamppb"
)

type IndexAdvisorResult struct {
	ID                 uint64
	CreatedAt          time.Time
	UpdatedAt          time.Time
	SchemaName         string
	TableName          string
	IndexName          string
	IndexColumns       []string
	IndexDetails       *IndexAdvisorResultIndexDetails
	TopImpactedQueries []*IndexAdvisorResultTopImpactedQuery
	WorkloadImpact     *IndexAdvisorResultWorkloadImpact
	Extra              *IndexAdvisorResultExtra
}

type IndexAdvisorResultIndexDetails struct {
	IndexSize int64
	Reason    string
}

type IndexAdvisorResultTopImpactedQuery struct {
	Query                string
	Improvement          float64
	CostSavingPer1MQuery float64
	CostSavingPerMon     float64
}

type IndexAdvisorResultWorkloadImpact struct {
	WorkloadImprovement float64
	CostSavingPerMon    float64
}

type IndexAdvisorResultExtra struct {
	Status string // open, closed, applied, pending
}

type IndexAdvisorFilter struct {
	State    string
	NameLike string
	Page     uint32
	PageSize uint32
	OrderBy  string
	Desc     bool
}

type indexAdvisorRepo interface {
	// Advise trigger the index advisor in TiDB kernel.
	// SQLs can be nil if not specified.
	Advise(ctx context.Context, SQLs []string) (DDLs []string, err error)

	// List returns the index advisor results by the filter.
	List(ctx context.Context, filter *IndexAdvisorFilter) ([]*IndexAdvisorResult, error)

	// Count returns the total count of the index advisor results by the filter.
	Count(ctx context.Context, filter *IndexAdvisorFilter) (uint64, error)

	// GetByID returns the index advisor result by the adviseID.
	GetByID(ctx context.Context, adviseID uint64) (*IndexAdvisorResult, error)

	// GetByIndexName returns the index advisor result by the database, table and index name.
	GetByIndexName(ctx context.Context, database, table, index string) (*IndexAdvisorResult, error)

	// ApplyByID applies the index advisor result by the adviseID.
	ApplyByID(ctx context.Context, adviseID uint64) error

	// UpdateStatusByID updates the status of the index advisor result by the adviseID.
	UpdateStatusByID(ctx context.Context, adviseID uint64, newStatus string) error

	// SyncApplyingAdvices updates the status of advises to 'applied' if the index is created.
	SyncApplyingAdvices(ctx context.Context) error
}

type indexAdvisorRepoImpl struct {
	proxy  facade.ClusterProxyFacade
	logger *zap.Logger
}

func (r *indexAdvisorRepoImpl) Advise(ctx context.Context, sqls []string) (ddls []string, err error) {
	ws, err := r.proxy.OpenWebsocket(ctx)
	if err != nil {
		r.logger.Error("failed to connect to database", zap.Error(err))
		return nil, err
	}
	defer ws.Close()

	if err = ws.Exec(ctx, "recommend index set timeout='300s'"); err != nil {
		r.logger.Error("failed to set index advisor timeout", zap.Error(err))
		return nil, fmt.Errorf("failed to set index advisor timeout: %w", err)
	}

	stmt := "recommend index run"
	if len(sqls) > 0 { // only for test
		stmt = fmt.Sprintf(`recommend index run for "%q"`, strings.Join(sqls, ";"))
	}
	r.logger.Info("start index advisor", zap.String("sql", stmt))
	begin := time.Now()
	data, bizErr := ws.Query(ctx, stmt)
	if bizErr != nil {
		r.logger.Error("failed to recommend index", zap.Error(bizErr))
		return nil, fmt.Errorf("failed to advise index: %w", bizErr)
	}
	r.logger.Info("proxy response", zap.String("response", string(data)))

	type ProxyResult struct {
		Database     []string `json:"database"`
		Table        []string `json:"table"`
		IndexName    []string `json:"index_name"`
		IndexColumns []string `json:"index_columns"`
	}
	var result ProxyResult
	if err := json.Unmarshal(data, &result); err != nil {
		r.logger.Error("failed to unmarshal json", zap.Error(err))
		return nil, bizerr.New(bizerr.ErrCodeJsonMarshal, "failed to unmarshal json")
	}
	r.logger.Info("index advisor finished", zap.Duration("duration", time.Since(begin)))

	d := func(str string) (string, error) {
		decoded, err := base64.StdEncoding.DecodeString(str)
		return string(decoded), err
	}
	for i := range result.Database {
		indexName, err := d(result.IndexName[i])
		if err != nil {
			return nil, err
		}
		database, err := d(result.Database[i])
		if err != nil {
			return nil, err
		}
		table, err := d(result.Table[i])
		if err != nil {
			return nil, err
		}
		indexColumns, err := d(result.IndexColumns[i])
		if err != nil {
			return nil, err
		}

		ddls = append(ddls, fmt.Sprintf("CREATE INDEX %s ON %s.%s (%s)",
			indexName, database, table, indexColumns))
	}
	return ddls, nil
}

func (r *indexAdvisorRepoImpl) List(ctx context.Context, filter *IndexAdvisorFilter) ([]*IndexAdvisorResult, error) {
	var wherePredicates []string
	var orderByClause string
	var offset, limit string
	if filter != nil {
		if filter.State != "" {
			wherePredicates = append(wherePredicates, filter.State)
		}
		if filter.NameLike != "" {
			nameLike := fmt.Sprintf("%%%s%%", filter.NameLike)
			wherePredicates = append(wherePredicates,
				fmt.Sprintf("`schema_name` like '%s' or `table_name` like '%s'", nameLike, nameLike))
		}
		if filter.OrderBy != "" {
			orderByClause = filter.OrderBy
			if filter.Desc {
				orderByClause += " DESC"
			}
			orderByClause += ", id DESC"
		} else {
			orderByClause = "id DESC"
		}
		if filter.Page >= 1 && filter.Page > 0 {
			offset = fmt.Sprintf("%d", (filter.Page-1)*filter.PageSize)
			limit = fmt.Sprintf("%d", filter.PageSize)
		}
	}

	stmt := "SELECT * from mysql.index_advisor_results"
	if len(wherePredicates) > 0 {
		stmt += " WHERE " + strings.Join(wherePredicates, " AND ")
	}
	if orderByClause != "" {
		stmt += " ORDER BY " + orderByClause
	}
	if offset != "" && limit != "" {
		stmt += fmt.Sprintf(" LIMIT %v, %v", offset, limit)
	}

	r.logger.Info("list index advisor result statement.", zap.String("sql", stmt))
	res, err := r.proxy.QuerySQL(ctx, stmt)
	if err != nil {
		r.logger.Error("failed to query sql", zap.Error(err))
		return nil, err
	}
	r.logger.Info("proxy response", zap.String("response", string(res)))

	results, derr := decodeIndexAdvisorResultFromProxy(res)
	if derr != nil {
		r.logger.Error("failed to decode index advisor result", zap.Error(derr))
		return nil, bizerr.New(errorcode.ErrCodeSvcClusterProxy, "failed to decode index advisor result")
	}
	return results, nil
}

func (r *indexAdvisorRepoImpl) GetByID(ctx context.Context, adviseID uint64) (*IndexAdvisorResult, error) {
	stmt := fmt.Sprintf("select * from mysql.index_advisor_results where id = %d", adviseID)
	r.logger.Info("get index advisor result statement", zap.String("sql", stmt))
	res, err := r.proxy.QuerySQL(ctx, stmt)
	if err != nil {
		r.logger.Error("failed to query sql", zap.Error(err))
		return nil, err
	}
	r.logger.Info("proxy response", zap.String("response", string(res)))

	results, derr := decodeIndexAdvisorResultFromProxy(res)
	if derr != nil {
		r.logger.Error("failed to decode index advisor result", zap.Error(derr))
		return nil, bizerr.New(errorcode.ErrCodeSvcClusterProxy, "failed to decode index advisor result")
	}
	if len(results) == 0 {
		return nil, bizerr.New(errorcode.ErrCodeSvcClusterProxy, "index advisor result not found")
	}

	return results[0], nil
}

func (r *indexAdvisorRepoImpl) GetByIndexName(
	ctx context.Context,
	database, table, index string,
) (*IndexAdvisorResult, error) {
	stmt := fmt.Sprintf(
		"select * from mysql.index_advisor_results where `schema_name` = '%s' and `table_name` = '%s' and `index_name` = '%s'",
		database,
		table,
		index,
	)

	r.logger.Info("get index advisor result statement", zap.String("sql", stmt))
	res, err := r.proxy.QuerySQL(ctx, stmt)
	if err != nil {
		r.logger.Error("failed to query sql", zap.Error(err))
		return nil, err
	}
	r.logger.Info("proxy response", zap.String("response", string(res)))

	results, derr := decodeIndexAdvisorResultFromProxy(res)
	if derr != nil {
		r.logger.Error("failed to decode index advisor result", zap.Error(derr))
		return nil, bizerr.New(errorcode.ErrCodeSvcClusterProxy, "failed to decode index advisor result")
	}
	if len(results) == 0 {
		return nil, bizerr.New(errorcode.ErrCodeSvcClusterProxy, "index advisor result not found")
	}

	return results[0], nil
}

func (r *indexAdvisorRepoImpl) Count(ctx context.Context, filter *IndexAdvisorFilter) (uint64, error) {
	var wherePredicates []string
	if filter != nil {
		if filter.State != "" {
			wherePredicates = append(wherePredicates, filter.State)
		}
		if filter.NameLike != "" {
			nameLike := fmt.Sprintf("%%%s%%", filter.NameLike)
			wherePredicates = append(wherePredicates,
				fmt.Sprintf("`schema_name` like '%s' or `table_name` like '%s'", nameLike, nameLike))
		}
	}

	stmt := "SELECT count(*) as total from mysql.index_advisor_results"
	if len(wherePredicates) > 0 {
		stmt += " WHERE " + strings.Join(wherePredicates, " AND ")
	}

	r.logger.Info("count index advisor result statement", zap.String("sql", stmt))
	res, err := r.proxy.QuerySQL(ctx, stmt)
	if err != nil {
		r.logger.Error("failed to query sql", zap.Error(err))
		return 0, err
	}
	r.logger.Info("proxy response", zap.String("response", string(res)))

	type proxyResult struct {
		Total []int `json:"total"`
	}
	var result proxyResult

	if err := json.Unmarshal(res, &result); err != nil {
		r.logger.Error("failed to unmarshal json", zap.Error(err))
		return 0, bizerr.New(bizerr.ErrCodeJsonMarshal, "failed to unmarshal json")
	}

	return uint64(result.Total[0]), nil
}

func (r *indexAdvisorRepoImpl) UpdateStatusByID(ctx context.Context, adviseID uint64, newStatus string) error {
	if newStatus != STATEAPPLIED && newStatus != STATECLOSED && newStatus != STATEOPEN && newStatus != STATEAPPLYING {
		return errors.New("invalid status")
	}

	stmt := "update mysql.index_advisor_results set extra = json_set(ifnull(extra, '{}'), '$.status', '%v') where id = %v"
	stmt = fmt.Sprintf(stmt, newStatus, adviseID)
	r.logger.Info("update index advisor result statement", zap.String("sql", stmt))
	err := r.proxy.ExecuteSQL(ctx, stmt)
	if err != nil {
		r.logger.Error("failed to query sql", zap.Error(err))
		return err
	}

	return nil
}

func (r *indexAdvisorRepoImpl) ApplyByID(ctx context.Context, adviseID uint64) error {
	advise, err := r.GetByID(ctx, adviseID)
	if err != nil {
		return err
	}
	if advise.Extra.Status == STATEAPPLIED || advise.Extra.Status == STATEAPPLYING {
		return fmt.Errorf("index advisor result %d already applied", adviseID)
	}

	stmt := fmt.Sprintf(
		"alter table %s.%s add index %s (%s)",
		advise.SchemaName,
		advise.TableName,
		advise.IndexName,
		strings.Join(advise.IndexColumns, ","),
	)
	r.logger.Info("apply index advisor result statement", zap.String("sql", stmt))
	err = r.proxy.ExecuteMultiSQL(ctx, []string{"set tidb_enable_async_index_creation = 1", stmt})
	if err != nil {
		if strings.Contains(err.Error(), "index already exist") {
			// maybe created by the user manually, ignore this error
			r.logger.Info("index already exist", zap.Error(err))
		} else {
			r.logger.Error("failed to query sql", zap.Error(err))
			return errorcode.MarkClusterSQLErrorUserActionable(err)
		}
	}

	err = r.UpdateStatusByID(ctx, adviseID, STATEAPPLYING)
	if err != nil {
		r.logger.Error("failed to update status", zap.Error(err))
		return err
	}

	return nil
}

func (r *indexAdvisorRepoImpl) SyncApplyingAdvices(ctx context.Context) error {
	// update the advise status to 'applied' if the index is created.
	stmt := `update mysql.index_advisor_results res join information_schema.tidb_indexes idx
				on res.schema_name=idx.table_schema and res.table_name=idx.table_name and res.index_name=idx.key_name 
				set extra=json_set(extra, '$.status', 'applied')  where res.extra->'$.status'='applying'`
	return r.proxy.ExecuteSQL(ctx, stmt)
}

type TimeProxy struct {
	Time  string `json:"time"`
	Valid bool   `json:"valid"`
}

func decodeIndexAdvisorResultFromProxy(raw []byte) ([]*IndexAdvisorResult, error) {
	type indexAdvisorResultProxy struct {
		ID                 []int       `json:"id"`
		CreatedAt          []TimeProxy `json:"created_at"`
		UpdatedAt          []TimeProxy `json:"updated_at"`
		SchemaNames        []string    `json:"schema_name"`
		TableNames         []string    `json:"table_name"`
		IndexNames         []string    `json:"index_name"`
		IndexColumns       []string    `json:"index_columns"`
		IndexDetails       []string    `json:"index_details"`
		TopImpactedQueries []string    `json:"top_impacted_queries"`
		WorkloadImpact     []string    `json:"workload_impact"`
		Extra              []string    `json:"extra"`
	}

	var proxyResult indexAdvisorResultProxy
	if err := json.Unmarshal(raw, &proxyResult); err != nil {
		return nil, err
	}

	d := func(str string) (string, error) {
		decoded, err := base64.StdEncoding.DecodeString(str)
		return string(decoded), err
	}
	results := make([]*IndexAdvisorResult, 0, len(proxyResult.ID))
	for i := range proxyResult.ID {
		var err error
		proxyResult.SchemaNames[i], err = d(proxyResult.SchemaNames[i])
		if err != nil {
			return nil, err
		}
		proxyResult.TableNames[i], err = d(proxyResult.TableNames[i])
		if err != nil {
			return nil, err
		}
		proxyResult.IndexNames[i], err = d(proxyResult.IndexNames[i])
		if err != nil {
			return nil, err
		}
		proxyResult.IndexColumns[i], err = d(proxyResult.IndexColumns[i])
		if err != nil {
			return nil, err
		}
		proxyResult.IndexDetails[i], err = d(proxyResult.IndexDetails[i])
		if err != nil {
			return nil, err
		}
		proxyResult.TopImpactedQueries[i], err = d(proxyResult.TopImpactedQueries[i])
		if err != nil {
			return nil, err
		}
		proxyResult.WorkloadImpact[i], err = d(proxyResult.WorkloadImpact[i])
		if err != nil {
			return nil, err
		}
		proxyResult.Extra[i], err = d(proxyResult.Extra[i])
		if err != nil {
			return nil, err
		}

		indexDetails := new(IndexAdvisorResultIndexDetails)
		if err := json.Unmarshal([]byte(proxyResult.IndexDetails[i]), indexDetails); err != nil {
			return nil, fmt.Errorf("failed to unmarshal index details: %w", err)
		}

		var topImpactedQueries []*IndexAdvisorResultTopImpactedQuery
		if err := json.Unmarshal([]byte(proxyResult.TopImpactedQueries[i]), &topImpactedQueries); err != nil {
			return nil, fmt.Errorf("failed to unmarshal top impacted queries: %w", err)
		}

		workloadImpact := new(IndexAdvisorResultWorkloadImpact)
		if err := json.Unmarshal([]byte(proxyResult.WorkloadImpact[i]), workloadImpact); err != nil {
			return nil, fmt.Errorf("failed to unmarshal workload impact: %w", err)
		}

		extra := new(IndexAdvisorResultExtra)
		if proxyResult.Extra[i] != "" { // not null
			if err := json.Unmarshal([]byte(proxyResult.Extra[i]), extra); err != nil {
				return nil, fmt.Errorf("failed to unmarshal extra: %w", err)
			}
		}
		if extra.Status == "" { // default is 'open'
			extra.Status = STATEOPEN
		}
		createdAt, err := time.Parse(time.RFC3339, proxyResult.CreatedAt[i].Time)
		if err != nil {
			return nil, err
		}
		updatedAt, err := time.Parse(time.RFC3339, proxyResult.UpdatedAt[i].Time)
		if err != nil {
			return nil, err
		}

		results = append(results, &IndexAdvisorResult{
			ID:                 uint64(proxyResult.ID[i]),
			CreatedAt:          createdAt,
			UpdatedAt:          updatedAt,
			SchemaName:         proxyResult.SchemaNames[i],
			TableName:          proxyResult.TableNames[i],
			IndexName:          proxyResult.IndexNames[i],
			IndexColumns:       strings.Split(proxyResult.IndexColumns[i], ","),
			IndexDetails:       indexDetails,
			TopImpactedQueries: topImpactedQueries,
			WorkloadImpact:     workloadImpact,
			Extra:              extra,
		})
	}
	return results, nil
}

const (
	STATEOPEN     = "open"
	STATEAPPLYING = "applying"
	STATEAPPLIED  = "applied"
	STATECLOSED   = "closed"
)

func stateToString(s console.IndexAdviceState) string {
	var result string
	switch s {
	case console.IndexAdviceState_OPEN:
		result = "extra->'$.status' is null"
	case console.IndexAdviceState_APPLYING:
		result = "extra->'$.status' = 'applying'"
	case console.IndexAdviceState_APPLIED:
		result = "extra->'$.status' = 'applied'"
	case console.IndexAdviceState_CLOSED:
		result = "extra->'$.status' = 'closed'"
	}
	return result
}

func stringToState(s string) console.IndexAdviceState {
	var state console.IndexAdviceState
	switch s {
	case STATEOPEN:
		state = console.IndexAdviceState_OPEN
	case STATEAPPLYING:
		state = console.IndexAdviceState_APPLYING
	case STATEAPPLIED:
		state = console.IndexAdviceState_APPLIED
	case STATECLOSED:
		state = console.IndexAdviceState_CLOSED
	default:
		state = console.IndexAdviceState_INDEX_ADVICE_STATE_UNSPECIFIED
	}
	return state
}

func buildFilter(req *console.ListIndexAdvicesReq) *IndexAdvisorFilter {
	filter := new(IndexAdvisorFilter)
	if req.NameFilter != nil {
		filter.NameLike = *req.NameFilter
	}

	if req.Desc != nil {
		filter.Desc = *req.Desc
	}

	if req.StateFilter != nil && *req.StateFilter != console.IndexAdviceState_INDEX_ADVICE_STATE_UNSPECIFIED {
		filter.State = stateToString(*req.StateFilter)
	}

	if req.PageSize == nil || *req.PageSize <= 0 {
		filter.PageSize = 10
	} else {
		filter.PageSize = *req.PageSize
	}

	if req.PageToken == nil || *req.PageToken <= 0 {
		filter.Page = 1
	} else {
		filter.Page = *req.PageToken
	}

	if req.OrderBy != nil {
		switch *req.OrderBy {
		case "state":
			filter.OrderBy = "case when extra->'$.status' is null then 1 when extra->'$.status'='applying' then 2 when extra->'$.status'='applied' then 3 when extra->'$.status'='closed' then 10 else 20 end"
		case "index_size":
			filter.OrderBy = "index_details->'$.IndexSize'"
		case "improvement":
			filter.OrderBy = "workload_impact->'$.WorkloadImprovement'"
		default:
			filter.OrderBy = ""
		}
	}

	return filter
}

func roundFloat(val float64, precision uint) float64 {
	ratio := math.Pow(10, float64(precision))
	return math.Round(val*ratio) / ratio
}

func convertIndexAdvisorResult(result *IndexAdvisorResult) *console.IndexAdvice {
	state := stringToState(result.Extra.Status)
	stateReason := ""
	if state == console.IndexAdviceState_CLOSED {
		stateReason = "User closed it manually"
	}
	impactedQueriesResult := make([]*console.IndexAdviceImpact, 0, len(result.TopImpactedQueries))
	for _, query := range result.TopImpactedQueries {
		impactedQueriesResult = append(impactedQueriesResult, &console.IndexAdviceImpact{
			Query:       query.Query,
			Improvement: query.Improvement,
		})
	}

	var maxCostSavingPerQuery float64
	for _, query := range result.TopImpactedQueries {
		if query.CostSavingPer1MQuery/1e6 > maxCostSavingPerQuery {
			maxCostSavingPerQuery = query.CostSavingPer1MQuery / 1e6
		}
	}
	costSavingMonthly := roundFloat(result.WorkloadImpact.CostSavingPerMon, 2)
	if costSavingMonthly > 100 {
		costSavingMonthly = 100
	}

	return &console.IndexAdvice{
		Id:                result.ID,
		Name:              result.IndexName,
		Database:          result.SchemaName,
		Table:             result.TableName,
		LastRecommendTime: timestamppb.New(result.UpdatedAt),
		State:             state,
		IndexStatement: fmt.Sprintf("CREATE INDEX %s ON %s.%s (%s)",
			result.IndexName, result.SchemaName, result.TableName, strings.Join(result.IndexColumns, ", ")),
		Improvement:        roundFloat(result.WorkloadImpact.WorkloadImprovement, 4),
		IndexSize:          roundFloat(float64(result.IndexDetails.IndexSize)/1024/1024, 2),
		Reason:             result.IndexDetails.Reason,
		TopImpactedQueries: impactedQueriesResult,
		StateReason:        stateReason,
		CostSavingMonthly:  costSavingMonthly,
		CostSavingPerQuery: maxCostSavingPerQuery,
	}
}