package plugin

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"moss-dev/datahub/api"
	"moss-dev/datamodels"
	"moss-dev/gpt"
	"moss-dev/schema_utils"
	"moss-dev/utils"

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

var (
	ProjectID            string
	DB                   string
	Table                string
	DatahubEndpoint      string
	DatahubToken         string
	GPTToken             string
	LogLevel             string
	MySQLConnect         string
	UpdateMeta           bool
	VectorEngineEndpoint string
	VectorEngineToken    string
	VectorCollection     string
)

const (
	metagenTag  = "metagen"
	analysisTag = "analysis available"
)

func NewGenMetaCmd() *cobra.Command {
	cmdGen := &cobra.Command{
		Use:   "gen",
		Short: "generate metadata(column description) for datahub model",
		//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")
			}

			cli := api.NewDatahubClient(DatahubEndpoint, DatahubToken)
			urn, tbl, err := fetchTableInfo(cli, ProjectID, DB, Table)
			if err != nil {
				log.Fatal("fail to fetch table info", zap.Error(err))
			}
			log.Debug("table info", zap.String("create table statement", tbl.CreateTableStatement(false)))
			if !filterTable(tbl) {
				log.Fatal("table is not available for analysis", zap.String("db", tbl.DB), zap.String("table", tbl.Table))
			}

			gptCli := gpt.NewGPT(gogpt.NewClient(GPTToken))

			modelAid := datamodels.NewModelAid()
			modelAid.SetGPTClient(gptCli)
			if UpdateMeta {
				modelAid.SetDatahubClient(cli)
			}
			cd, err := modelAid.GenMetaData(urn, tbl)
			if err != nil {
				log.Fatal("fail to generate meta", zap.Error(err))
			}

			if len(MySQLConnect) > 0 {
				mysqlCli, err := sql.Open("mysql", MySQLConnect)
				if err != nil {
					log.Fatal("fail to connect to mysql", zap.String("mysql", MySQLConnect), zap.Error(err))
				}
				defer mysqlCli.Close()

				if err := insertModelDefitionIntoMySQL(mysqlCli, urn, tbl, cd); err != nil {
					log.Fatal("fail to insert model into mysql", zap.Error(err))
				}
			} else {
				fmt.Printf("####### model metadata ########\n %s\n ##########", cd)
			}
			return nil
		},
	}

	cmdGen.Flags().StringVar(&ProjectID, "project-id", "gcp-digital-bi", "whichi project that table model belongs to, e.g. bigquery project ID")
	cmdGen.Flags().StringVar(&DB, "db", "", "the database name of the table model, e.g. salesforce")
	cmdGen.Flags().StringVar(&Table, "table", "", "the table name of the table model, e.g. crm_account")
	cmdGen.Flags().StringVar(&DatahubEndpoint, "dh-endpoint", "https://datahub.pingcap.net/api/graphql", "the endpoint of datahub")
	cmdGen.Flags().StringVar(&DatahubToken, "dh-token", "", "the access token of datahub")
	cmdGen.Flags().StringVar(&GPTToken, "gpt-token", "", "the access token of gpt")
	cmdGen.Flags().StringVar(&LogLevel, "log-level", "info", "the output log level: info, debug, warn, error")
	cmdGen.Flags().StringVar(&MySQLConnect, "mysql", "", "mysql connect string. e.g. user:password@tcp(host:port)/database")
	cmdGen.Flags().BoolVar(&UpdateMeta, "update-meta", true, "whether to update metadata in datahub")
	cmdGen.MarkFlagRequired("db")
	cmdGen.MarkFlagRequired("table")
	cmdGen.MarkFlagRequired("dh-token")
	cmdGen.MarkFlagRequired("gpt-token")
	return cmdGen
}

func filterTable(tbl *schema_utils.TableInfo) bool {
	available := false
	for _, tag := range tbl.Tags {
		if tag == metagenTag {
			log.Fatal("table has meta gen tag, skip", zap.String("tag", tag))
		}

		if tag == analysisTag {
			available = true
		}
	}
	return available
}

func fetchTableInfo(cli *api.DatahubClient, projectID, db, table string) (string, *schema_utils.TableInfo, error) {
	datasetUrn, err := cli.SearchDatasetUrn(projectID, fmt.Sprintf("%s.%s", db, table))
	if err != nil {
		log.Fatal("fail to search dataset", zap.Error(err))
	}
	log.Info("search table model", zap.String("urn", datasetUrn))

	dh := schema_utils.NewDatahubDB(cli)
	tbl, err := dh.FetchTableInfo(context.Background(), datasetUrn, db, table)
	if err != nil {
		log.Fatal("fail to fetch dataset", zap.Error(err))
	}

	return datasetUrn, tbl, nil
}

func insertModelDefitionIntoMySQL(db *sql.DB, urn string, tbl *schema_utils.TableInfo, colDef *datamodels.ColumnDenitions) error {
	if tbl == nil {
		return errors.New("table info is nil")
	}

	if colDef == nil {
		return nil
	}

	var (
		rawACD []byte
	)

	rawTableInfo, err := json.Marshal(tbl)
	if err != nil {
		return err
	}
	if colDef != nil {
		if rawACD, err = json.Marshal(colDef); err != nil {
			return err
		}
	}

	// Insert data
	insert, err := db.Query(`insert into model_def(urn, tableinfo, ACD) values(?, ?, ?);`, urn, string(rawTableInfo), string(rawACD))
	if err != nil {
		return err
	}
	defer insert.Close()

	return nil
}

func insertModelEvalIntoMySQL(db *sql.DB, urn string, tbl *schema_utils.TableInfo, ms *datamodels.ModelSuggestion) error {
	if tbl == nil {
		return errors.New("table info is nil")
	}

	if ms == nil {
		return nil
	}

	var (
		rawAOS []byte
	)

	rawTableInfo, err := json.Marshal(tbl)
	if err != nil {
		return err
	}
	if ms != nil {
		if rawAOS, err = json.Marshal(ms); err != nil {
			return err
		}
	}

	// Insert data
	insert, err := db.Query(`insert into model_eval(urn, tableinfo, AOS) values(?, ?, ?);`, urn, string(rawTableInfo), string(rawAOS))
	if err != nil {
		return err
	}
	defer insert.Close()

	return nil
}