package dataset

import (
	"context"
	_ "embed"
	"encoding/json"
	"strings"

	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"

	"github.com/tidbcloud/aitestsuit/consts"
	"github.com/tidbcloud/aitestsuit/internal/dao"
	"github.com/tidbcloud/aitestsuit/internal/entity"
)

var (
	//go:embed questions520.json
	questions520 []byte
	//go:embed spider.json
	spiderData []byte
)

type QuestionItem struct {
	Database  string   `json:"database"`
	Table     string   `json:"table"`
	SQL       string   `json:"sql"`
	ColNames  []string `json:"col_names"`
	Questions []string `json:"questions"`
}

func GenQuestion(ctx context.Context, questionSetName string) error {
	var items []QuestionItem
	var data []byte
	if questionSetName == consts.DataSetQuestion520 {
		data = questions520
	} else if questionSetName == consts.DataSetSpider {
		data = spiderData
	}

	if err := json.Unmarshal(data, &items); err != nil {
		tilog.Error(ctx, "failed to unmarshal data", zap.Error(err))
		return err
	}

	plan := entity.TestPlan{
		Type:        "SQL Generator",
		Creator:     "Brian",
		Name:        questionSetName,
		Description: "",
	}
	err := dao.CreateTestPlan(ctx, &plan)
	if err != nil {
		tilog.Error(ctx, "failed to create test plan", zap.Error(err))
		return err
	}

	for _, item := range items {
		for _, question := range item.Questions {
			testCase := entity.TestCase{
				PlanID:      plan.ID,
				Question:    question,
				Database:    item.Database,
				Table:       item.Table,
				SQL:         item.SQL,
				ColumnNames: strings.Join(item.ColNames, ","),
			}
			if err = dao.CreateTestCase(ctx, &testCase); err != nil {
				tilog.Error(ctx, "failed to create test case", zap.Error(err))
			}
		}
	}

	return nil
}