package entity

import (
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"sync"
	"time"

	"gorm.io/driver/mysql"
	"gorm.io/gorm"

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

	"github.com/tidbcloud/aitestsuit/configs"
	"github.com/tidbcloud/aitestsuit/internal/service"
)

var (
	db         *gorm.DB
	initDBOnce sync.Once
)

func InitDB(cfg *configs.Config, migrate bool) error {
	var err error

	initDBOnce.Do(func() {
		db, err = gorm.Open(mysql.Open(cfg.DBDSN), &gorm.Config{})
		if err != nil {
			return
		}

		// run migration
		if migrate {
			err = db.AutoMigrate(
				&TestPlan{}, &TestCase{}, &PromptTemplate{}, &TestRun{}, &TestReport{}, &TestCaseResult{}, &AIModel{},
			)
		}
	})

	return err
}

func GetDB() *gorm.DB {
	return db
}

type TestPlan struct {
	gorm.Model

	AIModelID   uint   `json:"ai_model_id"`
	Type        string `json:"type"` // 1. SQL Generator 2. DBName Selector   3. Gen Question
	Creator     string `json:"creator"`
	Name        string `json:"name"`
	Description string `json:"description"`
}

type Chart struct {
	ChartName string                 `json:"chart_name"`
	Title     string                 `json:"title"`
	Options   map[string]interface{} `json:"options"`
}

type CompareArgs struct {
	Chart Chart `json:"chart"`
}

// Scan scan value into Jsonb, implements sql.Scanner interface
func (j *CompareArgs) Scan(value interface{}) error {
	bytes, ok := value.([]byte)
	if !ok {
		return errors.New(fmt.Sprint("Failed to unmarshal JSONB value:", value))
	}

	var result = CompareArgs{}
	err := json.Unmarshal(bytes, &result)
	*j = result
	return err
}

// Value return json value, implement driver.Valuer interface
func (j CompareArgs) Value() (driver.Value, error) {
	return json.Marshal(j)
}

type TestCase struct {
	gorm.Model

	AIModelID   uint        `json:"ai_model_id"`
	PlanID      uint        `json:"plan_id" gorm:"index"`
	Question    string      `json:"question"`
	Database    string      `json:"database"`
	Table       string      `json:"table"`
	SQL         string      `json:"sql"`
	ColumnNames string      `json:"column_names"`
	CompareArgs CompareArgs `json:"compare_args"`
}

type AIModel struct {
	gorm.Model

	Name        string `json:"name"`
	Description string `json:"description"`
	ModelType   string `json:"model_type"`
}

type PromptTemplate struct {
	gorm.Model

	Version string `json:"version" gorm:"uniqueIndex,size:32"`
	Content string `json:"content"`
}

type Params struct {
	MaxTokens   int      `json:"max_tokens"`
	ModelType   string   `json:"model_type"`
	TopP        float32  `json:"top_p"`
	Temperature float32  `json:"temperature"`
	N           int      `json:"n"`
	Stop        []string `json:"stop"`
}

// Scan scan value into Jsonb, implements sql.Scanner interface
func (j *Params) Scan(value interface{}) error {
	bytes, ok := value.([]byte)
	if !ok {
		return errors.New(fmt.Sprint("Failed to unmarshal JSONB value:", value))
	}

	var result = Params{}
	err := json.Unmarshal(bytes, &result)
	*j = result
	return err
}

// Value return json value, implement driver.Valuer interface
func (j Params) Value() (driver.Value, error) {
	return json.Marshal(j)
}

type TestRun struct {
	gorm.Model

	Name   string `json:"name"`
	Params Params `json:"params"`
}

type Stats struct {
	GeneratingPassCount int `json:"generating_pass_count"`
	ExecutingPassCount  int `json:"executing_pass_count"`
	ComparingPassCount  int `json:"comparing_pass_count"`
	ChartTypeMatchCount int `json:"chart_type_match_count"`
	NonTableChartCount  int `json:"non_table_chart_count"`
	TotalCount          int `json:"total_count"`
}

// Scan scan value into Jsonb, implements sql.Scanner interface
func (j *Stats) Scan(value interface{}) error {
	bytes, ok := value.([]byte)
	if !ok {
		return errors.New(fmt.Sprint("Failed to unmarshal JSONB value:", value))
	}

	var result = Stats{}
	err := json.Unmarshal(bytes, &result)
	*j = result
	return err
}

// Value return json value, implement driver.Valuer interface
func (j Stats) Value() (driver.Value, error) {
	return json.Marshal(j)
}

type TestReport struct {
	gorm.Model

	StartAt          time.Time `json:"start_at"`
	EndAt            time.Time `json:"end_at"`
	Status           int       `json:"status"`
	TokenUsed        int       `json:"token_used"`
	PlanID           uint      `json:"plan_id" gorm:"index"`
	TestRunID        uint      `json:"test_run_id" gorm:"index"`
	PromptTemplateID uint      `json:"prompt_template_id"`
	PassRate         float64   `json:"pass_rate"`
	Prompt           string    `json:"prompt"`
	Stats            Stats     `json:"stats"`
}

type SQLResp struct {
	DataAPIExecSQLResp *dataapi.ExecSqlResp   `json:"dataapi_exec_sql_resp"`
	DataAPPResp        *service.Chat2DataResp `json:"dataapp_resp"`
}

// Scan scan value into Jsonb, implements sql.Scanner interface
func (j *SQLResp) Scan(value interface{}) error {
	bytes, ok := value.([]byte)
	if !ok {
		return errors.New(fmt.Sprint("Failed to unmarshal JSONB value:", value))
	}

	var result = SQLResp{}
	err := json.Unmarshal(bytes, &result)
	*j = result
	return err
}

// Value return json value, implement driver.Valuer interface
func (j SQLResp) Value() (driver.Value, error) {
	return json.Marshal(j)
}

type TestCaseResult struct {
	gorm.Model

	TestReportID   uint    `json:"test_report_id" gorm:"index"`
	TestCaseID     uint    `json:"test_case_id"`
	Status         int     `json:"status"`
	TokenUsed      int     `json:"token_used"`
	Input          string  `json:"input"`
	Output         string  `json:"output"`
	ReturnSum      int64   `json:"return_sum"`       // rows count
	ReturnColCount int64   `json:"return_col_count"` // column count
	SQLResp        SQLResp `json:"sql_resp"`         // data api response
	ResultA        int     `json:"result_a"`
	ResultH        int     `json:"result_h"`
	ErrorLog       string  `json:"error_log"`
	Duration       int64   `json:"duration"`
}