package chatgpttemplate

import (
	"fmt"
	"strings"

	"github.com/tidbcloud/aipower/internal/bean/dto"

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

const DefaultNoSqlGeneratedMsg = "No SQL Generated"
const DefaultRequestOpenAIError = "Request OpenAI Error"
const DefaultInvalidParamError = "Invalid Params"
const DefaultGetDatabaseColumnsError = "Get Database Columns Error"
const DefaultNoDatabaseGuessedMsg = "No Database Guessed"
const DefaultPromptTooLongMsg = "Prompt Too Long"
const DefaultDatabaseSchemaEmptyError = "Database Schema Empty"

func formatColumnName(cols []string) []string {
	retList := make([]string, len(cols))
	for index, col := range cols {
		if !strings.Contains(col, "`") {
			retList[index] = "`" + col + "`"
		} else {
			retList[index] = col
		}
	}
	return retList
}

type Prompt struct {
	Content     string
	Version     string
	SchemaEmpty bool
}

func FormatPromptForGuessDatabase(question string, metas []*dataapi.TableMeta) Prompt {
	template := `
Act as a data scientist, you are tasked with generating a JSON file that only include the name of database which contains most relevant table(s) to answer a question from the following database schemas: %s.
The question is: "%s".
If you are unable to identify the most appropriate schema, just provide empty string instead.
Do not provide additional information or explanations.
The format for the JSON file should be:{"dbName": string}.
let's think step by step and generate an answer.json based on the above requirements. Answer:`

	formatTemplate := fmt.Sprintf(template, FormatDatabaseSchemaString(metas), question)
	return Prompt{
		Content: formatTemplate,
		Version: "data2database-20230322",
	}
}

func FormatDatabaseSchemaString(metas []*dataapi.TableMeta) string {
	databaseMap := map[string][]*dataapi.TableMeta{}
	for _, meta := range metas {
		if _, ok := databaseMap[meta.Database]; !ok {
			databaseMap[meta.Database] = []*dataapi.TableMeta{}
		}
		databaseMap[meta.Database] = append(databaseMap[meta.Database], meta)
	}
	// if only contains one database
	isContainsOneDB := len(databaseMap) == 1
	databaseLines := []string{}
	for database, metas := range databaseMap {
		databaseLine := "dbName = `" + database + "`, tables = "
		tableLines := []string{}
		for _, meta := range metas {
			tableLine := ""
			if isContainsOneDB {
				tableLine = "{dbName:`" + database + "`, table: `" + meta.Table + "`, columns: [" + strings.Join(formatColumnName(meta.Columns), ", ") + "]}"
			} else {
				tableLine = "{table: `" + meta.Table + "`, columns: [" + strings.Join(formatColumnName(meta.Columns), ", ") + "]}"
			}
			tableLines = append(tableLines, tableLine)
		}
		databaseLine = databaseLine + "[" + strings.Join(tableLines, ", ") + "]"
		databaseLines = append(databaseLines, databaseLine)
	}
	return strings.Join(databaseLines, "\n")
}

func FormatTableSchemaString(req *dto.InteractRequest, metas []*dataapi.TableMeta) string {
	tableLines := []string{}
	tables := map[string]bool{}
	for _, i := range req.Tables {
		tables[i] = true
	}

	for _, meta := range metas {
		// if req.Tables is not null, we should filter those tables not in the given list
		if len(tables) > 0 && !tables[meta.Table] {
			continue
		}

		meta.Columns = formatColumnName(meta.Columns)
		line := "Table = " + "`" + meta.Table + "`" + ", columns = [" + strings.Join(meta.Columns, ", ") + "]"
		tableLines = append(tableLines, line)
	}
	return strings.Join(tableLines, "\n")
}

func FormatCustomizedPromptForChat2Question(req *dto.InteractRequest, metas []*dataapi.TableMeta) Prompt {
	content := fmt.Sprintf(req.PromptTemplate, FormatTableSchemaString(req, metas), req.QuestionNums)
	return Prompt{
		Content: content,
		Version: req.PromptTemplateVersion,
	}
}

func FormatPromptForData2ColumnMatching(sampleData string) Prompt {
	template := `
You are tasked with generating a JSON file that includes the data type and description for every column based on a given set of sample data: %s.
The format for the JSON file should be:
[{
    "column": string,
    "type": string, // The type must be supported by MySQL
    "description": string
}...]
Act as a senior data scientist, let's think step by step and generate an answer.json based on the above requirements. Answer:`

	formatTemplate := fmt.Sprintf(template, sampleData)
	return Prompt{
		Content: formatTemplate,
		Version: "data2columnMatching-20230306",
	}
}

func FormatCustomizedPromptForData2ColumnMatching(req *dto.InteractRequest) Prompt {
	content := fmt.Sprintf(req.PromptTemplate, req.SampleData)
	return Prompt{
		Content: content,
		Version: req.PromptTemplateVersion,
	}
}

func FormatPromptForGenDatabaseQuestions(req *dto.InteractRequest, metas []*dataapi.TableMeta, questionNum int) Prompt {

	template := `
You are tasked with generating interesting questions to explore the data based on the following MySQL database schema: %s. 
The questions you generate must meet the following requirements: 
- The generated questions should not contain specific table names. 
- The question should cover different aspects of the data (e.g. descriptive statistics, relationships between variables, etc.). 
- The question should be open-ended, meaning that there are multiple ways to answer it and it encourages further exploration.

Act as a senior data scientist, Let's think step by step and use this template to generate an answer.json: {
step. Can you write a brief description of the table in one sentence? 
{"desc": {to_fill}}

step. What is the grain of the table in one sentence? 
{"grain": {to_fill}}

step. Can you generate %d diverse questions from different perspectives about this table in the following schema: [{
"question_id": number, // the id of question 
"question": string, // the content of question, must be single line question, remove the order number 
"question_keyword":string,// use only one eye-catching keyword to summarize the question, it must be one word!
}]. You MUST make sure the questions can be transformed into a workable SQL statement! 
{"questions": {to_fill}}
}. Answer:`

	formatTemplate := fmt.Sprintf(template, FormatTableSchemaString(req, metas), questionNum)

	return Prompt{
		Content: formatTemplate,
		Version: "chat2question-20230228",
	}
}

func FormatCustomizedPromptForChat2Chart(req *dto.InteractRequest, metas []*dataapi.TableMeta) Prompt {
	content := fmt.Sprintf(req.PromptTemplate, FormatTableSchemaString(req, metas), req.Question)
	return Prompt{
		Content: content,
		Version: req.PromptTemplateVersion,
	}
}

func FormatPromptForChat2Chart(req *dto.InteractRequest, metas []*dataapi.TableMeta) Prompt {
	template := `
You are tasked with generating a SQL statement to answer a specific question based on the given database schema: %s. 
And select a chart type from the following options that would be most suitable for visualizing the data generated by your SQL statement:
PieChart {label: Column, value: Column} 
LineChart {x: Column, y: Column | Column[]} 
BarChart {x: Column, y: Column | Column[]} 
Table {columns: Column[]} 

The SQL statement you generate should meet the following requirements: %s 

The output answer should meet the following schema: 
{ 
 "sql": string, // SQL statement !!! 
 "chart": {"chartName": string, "title": string, "options": ChartOptions} // must generate chart!!!
}

Assume you are an extremely experienced MySQL expert, let's think step by step and generate an answer based on the above requirements from the following question: "%s". Answer:

After generating the answer, rewrite the SQL string inside it according to the following rules: %s`

	var (
		databaseSchema string
		version        string
	)
	if req.DatabaseSchema != "" {
		databaseSchema = req.DatabaseSchema
		version = "chat2chart-20230306"
	} else {
		databaseSchema = FormatTableSchemaString(req, metas)
		version = "chat2chart-20230221"
	}
	content := fmt.Sprintf(
		template,
		databaseSchema,
		formatSQLGenerationInstructionString(),
		req.Question,
		formatSQLRewriteInstructionString())

	return Prompt{
		Content: content,
		Version: version,
	}
}

func FormatCustomizedPromptForSQLEditorDatabase(req *dto.InteractRequest, metas []*dataapi.TableMeta) Prompt {
	content := fmt.Sprintf(req.PromptTemplate, FormatTableSchemaString(req, metas), req.Question)
	return Prompt{
		Content: content,
		Version: req.PromptTemplateVersion,
	}
}

func FormatPromptForSqlEditorDatabase(req *dto.InteractRequest, metas []*dataapi.TableMeta) Prompt {
	return FormatPromptForChat2Query(req, metas)
}

// Context: Set context for Chat-GPT.
// Instruction: Requirements to generate answer.
// Task: Describe the task clearly.
// ReWrite: Post-processing, format the result...

func FormatPromptForChat2Query(req *dto.InteractRequest, metas []*dataapi.TableMeta) Prompt {
	template := `
Act as a MySQL expert, you are tasked with writing a SQL statement to answer question based on the following database schema: %s. 
The question is: "%s". Let's think step by step and use the following answer template to generate an json file: {
1. Output the revised question, including all necessary keywords. Marked as RQ. {"RQ": string}.
2. List all necessary tables involved in answering the RQ. Mark as NT. {"NT": string[]}
3. List all necessary columns of NT involved in answering the RQ, do not introduce non-existent columns. Marked as NC. {"NC": {table: string, Columns: string[]}}
4. Using best practices of MySQL to write a single SQL statement to answer the RQ. {"sql": string}
5. Rewriting the SQL statement based on the following rules carefully. If no SQL statement can be generated, output the message 'No SQL Generated'. {"rewriteSql": string}
'''
- Recheck the column names listed in NC, replace the corresponding column names in the SQL statement.
- Use the table.column format in the SQL statement when multiple tables mentioned in the NT list and all table aliases must be correctly referenced.
- Avoid the Use of SELECT Asterisk (SELECT *).
- Ensure that columns are referenced properly and only exist in the appropriate table.
- Enclose all table names and column names with backtick.
'''
}`

	content := fmt.Sprintf(
		template,
		FormatTableSchemaString(req, metas),
		req.Question)

	return Prompt{
		Content: content,
		Version: "chat2sql-20230313",
	}
}

func formatSQLGenerationInstructionString() string {
	instructionLines := []string{
		"- Output only SQL.",
		"- Use only one table if possible.",
		"- If multiple tables are needed, use the table.column format to avoid ambiguous columns.",
		"- Keep the table and column names as they are. Do not change them.",
		"- Do not add additional columns to the tables.",
		"- Ensure that columns are referenced properly and only exist in the appropriate table.",
		"- If the column name has spaces, enclose it with backticks (`) in SQL statement.",
		"- If no SQL statement can be generated, output the message 'No SQL Generated.",
	}
	return strings.Join(instructionLines, "\n")
}

func formatSQLRewriteInstructionString() string {
	instructionLines := []string{
		"- Wrap all columns with backticks (`).",
		"- Properly format the SQL statement with punctuation and end it with a semicolon.",
	}
	return strings.Join(instructionLines, "\n")
}