package dto

import (
	"github.com/sashabaranov/go-openai"

	"github.com/tidbcloud/aipower/idl/pbgen/proto/api"
	"github.com/tidbcloud/aipower/internal/bean/entity"
)

type GPT35Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type OpenAIParams struct {
	Model            string  `json:"model"`
	Temperature      float32 `json:"temperature"`
	TopP             float32 `json:"top_p"`
	N                int     `json:"n"`
	Stream           bool    `json:"stream"`
	MaxTokens        int     `json:"max_tokens"`
	PresencePenalty  float32 `json:"presence_penalty"`
	FrequencyPenalty float32 `json:"frequency_penalty"`
}

type CreateGPT35TurboCompletionParam struct {
	OpenAIParams

	Model    string         `json:"model"`
	Messages []GPT35Message `json:"messages"`
}

type CreateCompletionParam struct {
	OpenAIParams

	Model    string   `json:"model"`
	Prompt   []string `json:"prompt"`
	Logprobs int32    `json:"logprobs"`
}

type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}
type CreateMessageParam struct {
	Model            string    `json:"model"`
	Messages         []Message `json:"messages"`
	MaxTokens        int32     `json:"max_tokens"`
	Temperature      float32   `json:"temperature"`
	TopP             float32   `json:"top_p"`
	Stream           bool      `json:"stream"`
	FrequencyPenalty int32     `json:"frequency_penalty"`
	PresencePenalty  int32     `json:"presence_penalty"`
}

type CreateMessageResult struct {
	Id      string `json:"id"`
	Object  string `json:"object"`
	Created int    `json:"created"`
	Model   string `json:"model"`
	Choices []struct {
		Index        int     `json:"index"`
		Message      Message `json:"message"`
		FinishReason string  `json:"finish_reason"`
	} `json:"choices"`
	Usage struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
		TotalTokens      int `json:"total_tokens"`
	} `json:"usage"`
}

type CreateCompletionError struct {
	Error struct {
		Message string `json:"message"`
		Type    string `json:"type"`
		Code    string `json:"code"`
	}
}

type CreateCompletionResult struct {
	JobID string `json:"job_id"`
	Err   error  `json:"err"`

	Object  string       `json:"object"`
	Created int64        `json:"created"`
	Model   string       `json:"model"`
	Content string       `json:"content"`
	Usage   openai.Usage `json:"usage"`
}

type InteractRequest struct {
	SessionID             string   `json:"session_id"`
	Question              string   `json:"question"`
	BotType               string   `json:"bot_type"`
	ProjectID             uint64   `json:"-"`
	ClusterID             uint64   `json:"cluster_id,omitempty,string"`
	Database              string   `json:"database,omitempty"`
	PromptTemplate        string   `json:"prompt_template"`
	PromptTemplateVersion string   `json:"prompt_template_version"`
	QuestionNums          int      `json:"question_nums"`
	Tables                []string `json:"tables"`
	Model                 string   `json:"model"`
	SampleData            string   `json:"sample_data"`
	DatabaseSchema        string   `json:"database_schema"`
	Async                 bool     `json:"async"`
	Caller                string   `json:"caller"`
	Timeout               int      `json:"timeout"`
	OrgID                 uint64   `json:"org_id,string"`
	AppID                 uint64   `json:"app_id"`
	AppKey                string   `json:"app_key"`
	Endpoint              string   `json:"endpoint"`
	Source                string   `json:"source"`
}

type ChatBotReqSchema struct {
	InteractRequest

	Caller string `json:"caller"`
}

type InteractResponse struct {
	SessionID     string `json:"session_id"`
	Question      string `json:"question"`
	Prompt        string `json:"prompt"`
	Response      string `json:"response"`
	ErrorMsg      string `json:"error_msg"`
	Code          int    `json:"code"`
	Database      string `json:"database"`
	StartTime     string `json:"start_time"`
	EndTime       string `json:"end_time"`
	StartTimeMS   int64  `json:"start_time_ms"`
	EndTimeMS     int64  `json:"end_time_ms"`
	ExecuteTime   int64  `json:"execute_time"`
	PromptVersion string `json:"prompt_version"`
	TokenUsed     int    `json:"token_used"`
	JobID         string `json:"job_id"`
}

type UserInfo struct {
	UserID uint64 `json:"user_id"`
	Email  string `json:"email"`
}

type SqlEditorGetColumnParam struct {
	OrgID     string `json:"org_id"`
	ProjectID string `json:"project_id"`
	ClusterID string `json:"cluster_id"`
	DbName    string `json:"db_name"`
	TableName string `json:"table_name"`
}

type SqlEditorGetColumnResult struct {
	Code int `json:"code"`
	Data struct {
		Clusterid int    `json:"clusterid"`
		Database  string `json:"database"`
		Table     string `json:"table"`
		Columns   []struct {
			Col      string `json:"col"`
			DataType string `json:"data_type"`
			Nullable bool   `json:"nullable"`
		} `json:"columns"`
	} `json:"data"`
	Message string `json:"message"`
}

type InteractRecFilter struct {
	ID        []uint64 `json:"id" form:"id"`
	SessionID string   `json:"session_id" form:"session_id"`
	Question  string   `json:"question" form:"question"`
}

type ListInteractRecParam struct {
	Filter   InteractRecFilter
	SortBy   string `json:"sort_by" form:"sort_by"`
	Order    string `json:"order" form:"order"`
	PageNo   int    `json:"page_no" form:"page_no"`
	PageSize int    `json:"page_size" form:"page_size"`
}

type ListInteractRecResult struct {
	Total int64
	List  []*entity.BotInteractRecord
}

type ColumnDescription struct {
	Column      string `json:"column"`
	Type        string `json:"type"`
	Description string `json:"description"`
}

type Chat2QueryCompletion struct {
	RQ         string              `json:"RQ"`
	NT         []string            `json:"NT"`
	NC         map[string][]string `json:"NC"`
	Sql        string              `json:"sql"`
	RewriteSql string              `json:"rewriteSql"`
}

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

type QuestionInfo struct {
	QuestionID      int    `json:"question_id"`
	Question        string `json:"question"`
	QuestionKeyword string `json:"question_keyword"`
}

type Chat2ChartCompletion struct {
	Sql   string    `json:"sql"`
	Chart ChartInfo `json:"chart"`
}

type GuessDatabaseCompletion struct {
	DBName string `json:"dbName"`
}

// part filed

type Chat2QuestionCompletion struct {
	Desc      string         `json:"desc"`
	Grain     string         `json:"grain"`
	Questions []QuestionInfo `json:"questions"`
}

func ConvertPromptTemplate(template *entity.PromptTemplate) *api.Prompt {
	return &api.Prompt{
		AiModel:    template.AIModel,
		Deployment: template.Deployment,
		Version:    template.Version,
		Content:    template.Content,
	}
}

type GetAPIKeyParams struct {
	AccessKey string `json:"access_key" form:"access_key"`
}

type GetClusterDBURLParams struct {
	ClusterID uint64 `json:"cluster_id" form:"cluster_id"`
}

type CheckAPIKeyPermissionParams struct {
	APIKey      string `json:"api_key"`
	ResourceURI string `json:"resource_uri"`
}

type CheckUserPermissionParams struct {
	UserToken   string `json:"user_token"`
	UserID      uint64 `json:"user_id"`
	ResourceURI string `json:"resource_uri"`
}