package plugin

import (
	"context"
	"moss-dev/datahub/api"
	"moss-dev/gpt"
	"moss-dev/utils"
	vector "moss-dev/vector_engine"
	"net/url"

	"github.com/juju/errors"
	"github.com/pingcap/log"
	gogpt "github.com/sashabaranov/go-openai"
	"github.com/spf13/cobra"
	"go.uber.org/zap"
)

func NewModelEmbeddingCmd() *cobra.Command {
	cmdEmbed := &cobra.Command{
		Use:   "embed",
		Short: "create the embedding vector and store it into vector engine",
		//Args:  cobra.MinimumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := utils.InitLogger(LogLevel, ""); err != nil {
				log.Fatal("Failed to initialize log", zap.Error(err))
			}

			if len(DB) == 0 {
				return errors.New("db is must required")
			}
			if len(Table) == 0 {
				return errors.New("table is must required")
			}
			if len(DatahubToken) == 0 {
				return errors.New("datahub access token is must required")
			}
			if len(GPTToken) == 0 {
				return errors.New("gpt access token is must required")
			}
			if len(VectorEngineToken) == 0 {
				return errors.New("vector engine token is must required")
			}

			cli := api.NewDatahubClient(DatahubEndpoint, DatahubToken)
			_, tbl, err := fetchTableInfo(cli, ProjectID, DB, Table)
			if err != nil {
				log.Fatal("fail to fetch table info", zap.Error(err))
			}

			gptCli := gpt.NewGPT(gogpt.NewClient(GPTToken))
			vectorCli, err := createvectorEngine(VectorEngineEndpoint, VectorEngineToken, gptCli)
			if err != nil {
				log.Fatal("fail to create qdrant vector engine", zap.Error(err))
			}

			pointID, err := vectorCli.UpsertTableSchema(context.Background(), VectorCollection, ProjectID, DB, Table, tbl.CreateTableStatement(true))
			if err != nil {
				log.Fatal("fail to insert point", zap.Error(err))
			}

			log.Info("add point", zap.String("point id", pointID), zap.String("create table statement", tbl.CreateTableStatement(true)))

			return nil
		},
	}

	cmdEmbed.Flags().StringVar(&ProjectID, "project-id", "gcp-digital-bi", "whichi project that table model belongs to, e.g. bigquery project ID")
	cmdEmbed.Flags().StringVar(&DB, "db", "", "the database name of the table model, e.g. salesforce")
	cmdEmbed.Flags().StringVar(&Table, "table", "", "the table name of the table model, e.g. crm_account")
	cmdEmbed.Flags().StringVar(&DatahubEndpoint, "dh-endpoint", "https://datahub.pingcap.net/api/graphql", "the endpoint of datahub")
	cmdEmbed.Flags().StringVar(&DatahubToken, "dh-token", "", "the access token of datahub")
	cmdEmbed.Flags().StringVar(&GPTToken, "gpt-token", "", "the access token of gpt")
	cmdEmbed.Flags().StringVar(&VectorEngineEndpoint, "vector-endpoint", "qdrant://833f0a92-662a-414b-85c0-11cc2bf70416.us-east-1-0.aws.cloud.qdrant.io:6334", "the endpoint of vector engine")
	cmdEmbed.Flags().StringVar(&VectorEngineToken, "vector-token", "", "the access token of vector engine")
	cmdEmbed.Flags().StringVar(&VectorCollection, "vector-collection", "model-schema-meta", "the collection used to store data warehouse schema meta vectors")
	cmdEmbed.Flags().StringVar(&LogLevel, "log-level", "info", "the output log level: info, debug, warn, error")
	cmdEmbed.MarkFlagRequired("db")
	cmdEmbed.MarkFlagRequired("table")
	cmdEmbed.MarkFlagRequired("dh-token")
	cmdEmbed.MarkFlagRequired("gpt-token")
	cmdEmbed.MarkFlagRequired("vector-token")
	return cmdEmbed
}

func createvectorEngine(endpoint string, token string, gptCli *gpt.GPT) (*vector.QdrantVectorEngine, error) {
	if len(endpoint) == 0 {
		return nil, errors.New("vector engine endpoint is must required")
	}

	u, err := url.Parse(endpoint)
	if err != nil {
		return nil, errors.Trace(err)
	}

	switch u.Scheme {
	case "qdrant":
		return vector.CreateQdrantVectorEngine(u.Host, token, gptCli)
	default:
		return nil, errors.NotSupportedf("vector engine endpoint %s, now we only support qdrant://{endpoint}", endpoint)
	}
}