package chatgpttemplate

import (
	"fmt"
	"strconv"
	"strings"

	"message-service/commons/utils"

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

const DefaultNoSqlGeneratedMsg = "No SQL Generated"
const DefaultNoDatabaseGuessedMsg = "No Database Guessed"
const DefaultPromptTooLongMsg = "Prompt Too Long"

const backtick = "`"

var mysqlReserveWords = []string{"rank"}

func needWrapBicktick(col string) bool {
	if strings.Contains(col, "`") {
		return false
	}
	if strings.Contains(col, " ") {
		return true
	}

	if utils.CheckInStringList(strings.ToLower(col), mysqlReserveWords) {
		return true
	}

	return false
}

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
}

func FormatPromptForGuessDatabase(prompt string, 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)
	}
	tableLines := []string{}
	databaseIndex := 1
	for database, metas := range databaseMap {
		tableIndex := 1
		tableLines = append(tableLines, ". database"+strconv.Itoa(databaseIndex)+"_name: "+database+"\n")
		for _, meta := range metas {
			tableLine := ". database" + strconv.Itoa(databaseIndex) + "_table" + strconv.Itoa(tableIndex) + "_name: " + meta.Table + "\n"
			// colLine := ""
			colLine := ". database" + strconv.Itoa(databaseIndex) + "_table" + strconv.Itoa(tableIndex) + "_columns: [" + strings.Join(formatColumnName(meta.Columns), ", ") + "]\n"

			itemLine := tableLine + colLine
			tableLines = append(tableLines, itemLine)
			tableIndex++
		}
		databaseIndex++
	}
	template := `#Database Selector
###Task:
Select one database based on following input parameters: 

` + strings.Join(tableLines, "\n") + `

Output the database name only.
If the column name has spaces, wrap it with ` + backtick + `
###Question:
Let's think step by step, select one database by interpreting this question: "` + prompt + `", output the database_name ONLY.`
	return template
}

func FormatPromptForGenQuestions(meta *dataapi.TableMeta, questionNum int) string {
	return FormatPromptForGenDatabaseQuestions(meta.Database, []*dataapi.TableMeta{meta}, questionNum)
}

func FormatTableInstructionString(database string, metas []*dataapi.TableMeta) string {
	tableLines := []string{}
	for _, meta := range metas {
		meta.Columns = formatColumnName(meta.Columns)
		line := "Table " + database + "." + meta.GetTable() + ", columns = [" + strings.Join(meta.Columns, ", ") + "]"
		tableLines = append(tableLines, line)
	}
	return strings.Join(tableLines, "\n")
}

func FormatPromptForGenDatabaseQuestions(database string, metas []*dataapi.TableMeta, questionNum int) string {
	tableLines := []string{}
	for _, meta := range metas {
		meta.Columns = formatColumnName(meta.Columns)
		line := "Table " + database + "." + meta.GetTable() + ", columns = [" + strings.Join(meta.Columns, ", ") + "]"
		tableLines = append(tableLines, line)
	}
	template := `#Question Generator for Data Exploring Purpose

` + strings.Join(tableLines, "\n") + `

###Task:
Generate ` + strconv.Itoa(questionNum) + ` interesting questions to explore the data based on the following MySQL database schema:
`
	return template
}
func FormatPromptForChat2Chart(prompt string, database string, metas []*dataapi.TableMeta) Prompt {
	template := `# MySQL SQL
# Generate an SQL statement based on the following database schema:
---
` + FormatTableInstructionString(database, metas) + `
---
# SQLOptions
If no SQL can be generated, display the message "No SQL Generated".
If the column name has spaces, enclose it in backticks (` + "`" + `).
If the column name is a reserved word in MySQL, enclose it in backticks (` + "`" + `).
---

# ChartOptions
type Column = string; // must be the column name in the SQL result!!!
PieChart {label: Column; value: Column;}
LineChart {x: Column; y: Column | Column[];}
BarChart {x: Column; y: Column | Column[];}
MapChart {country_code: Column; value: Column;}
NumberCard {label: Column | null; value: Column;}
RepoCard {repo_name: Column;}
PersonalCard {user_login: Column;}
Table {columns: Column[];}

Answer {
  sql: string; // single line sql
  chart: {chartName: string; title: string; options: ChartOptions;}; // must generate chart!!!
}
---
Let's think step by step, use best practice of writing SQL, generate a answer.json file to answer the question: "` + prompt + `".
---
answer.json // must be json!!!
---`
	return Prompt{
		Content: template,
		Version: "20230129",
	}
}

func FormatCustomizedPromptForSQLEditorDatabase(template, version, question, database string, metas []*dataapi.TableMeta) Prompt {
	content := fmt.Sprintf(template, FormatTableInstructionString(database, metas), question)
	return Prompt{
		Content: content,
		Version: version,
	}
}

func FormatPromptForSqlEditorDatabase(question string, database string, metas []*dataapi.TableMeta) Prompt {
	return FormatPromptForChat2Query(question, database, metas)
}

func FormatPromptForChat2Query(question string, database string, metas []*dataapi.TableMeta) Prompt {
	template := `---# MySQL SQL Generator---
---# Generate an SQL statement based on the following database schema:
%s
---
---
- Output only SQL.
- The generated SQL must be properly formatted and end with a semicolon (;).
- If no SQL can be generated, display the message "No SQL Generated".
- Use one table if can match the question.
- Make sure to avoid ambiguous column references and non-exitent columns by using table aliases and double-checking column names before running the query.---

# Question:
# Assume you are a MySQL expert, let's think step by step, use best practices of writing SQL, generate one SQL statement from this question: "%s"

--- After generate the sql, rewrite the sql.
- Wrap backticks with all columns.
- DO NOT enclose backtick around database.table.
---`

	content := fmt.Sprintf(template, FormatTableInstructionString(database, metas), question)
	return Prompt{
		Content: content,
		Version: "20230201",
	}
}