// gen_clickhouse.go
package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"io/ioutil"
	"os"
	"reflect"
	"strings"
)

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Usage: gen_clickhouse <source_file>")
		os.Exit(1)
	}
	sourceFile := os.Args[1]

	// 解析源文件
	fset := token.NewFileSet()
	node, err := parser.ParseFile(fset, sourceFile, nil, parser.AllErrors)
	if err != nil {
		fmt.Printf("Failed to parse file: %v\n", err)
		os.Exit(1)
	}

	// 提取结构体定义并生成代码
	var generatedCode bytes.Buffer
	for _, decl := range node.Decls {
		if code := processDecl(decl); code != "" {
			generatedCode.WriteString(code + "\n")
		}
	}

	// 如果有生成的代码，插入到源文件
	if generatedCode.Len() > 0 {
		err = updateSourceFile(sourceFile, generatedCode.String())
		if err != nil {
			fmt.Printf("Failed to update file: %v\n", err)
			os.Exit(1)
		}
		fmt.Printf("Successfully updated %s\n", sourceFile)
	}
}

// 解析结构体定义并生成表创建和批量插入方法
func processDecl(decl ast.Decl) string {
	genDecl, ok := decl.(*ast.GenDecl)
	if !ok || genDecl.Tok != token.TYPE {
		return ""
	}

	var result strings.Builder
	for _, spec := range genDecl.Specs {
		typeSpec, ok := spec.(*ast.TypeSpec)
		if !ok {
			continue
		}

		structType, ok := typeSpec.Type.(*ast.StructType)
		if !ok {
			continue
		}

		// 提取结构体信息
		structName := typeSpec.Name.Name
		fields := extractFields(structType)

		// 生成表创建语句方法
		result.WriteString(generateCreateTableMethod(structName, fields))
		result.WriteString("\n")

		// 生成批量插入方法
		result.WriteString(generateInsertMethod(structName, fields))
		result.WriteString("\n")
	}
	return result.String()
}

// 提取结构体字段信息
func extractFields(structType *ast.StructType) []map[string]string {
	fields := []map[string]string{}

	for _, field := range structType.Fields.List {
		if field.Tag == nil || len(field.Names) == 0 {
			continue
		}

		tags := reflect.StructTag(strings.Trim(field.Tag.Value, "`"))
		columnName := tags.Get("ch_name")
		columnType := tags.Get("ch_type")
		order := tags.Get("ch_order")

		if columnName != "" && columnType != "" {
			fields = append(fields, map[string]string{
				"FieldName":  field.Names[0].Name,
				"ColumnName": columnName,
				"ColumnType": columnType,
				"order":      order,
			})
		}
	}

	return fields
}

// 生成表创建语句的方法
func generateCreateTableMethod(structName string, fields []map[string]string) string {
	var query string
	var orderBy string

	// 生成硬编码的 SQL 语句
	query = fmt.Sprintf("\nfunc CreateTable%s() string {\nquery := `CREATE TABLE IF NOT EXISTS %s (\n", CamelToUnderline(structName), CamelToUnderline(structName))
	// 遍历字段，生成列定义
	for _, field := range fields {
		columnName := field["ColumnName"]
		columnType := field["ColumnType"]
		query += fmt.Sprintf("\t%s %s,\n", columnName, columnType)
		if order := field["order"]; order != "" {
			orderBy = columnName
		}
	}

	// 去掉最后的逗号并添加 ENGINE 和 ORDER BY
	query = query[:len(query)-2] + "\n) ENGINE = MergeTree()"
	if orderBy != "" {
		query += fmt.Sprintf(" ORDER BY %s", orderBy)
	} else {
		// 默认排序字段为 id
		query += " ORDER BY id"
	}
	query += "`\n"

	query += "\n\treturn query\n}\n"

	// 返回完整的 SQL 生成代码
	return query
}

// 生成批量插入的方法
func generateInsertMethod(structName string, fields []map[string]string) string {
	var query string
	var columnNames []string
	var fieldVars []string

	// 构建 INSERT INTO 语句的列名部分
	for _, field := range fields {
		columnName := field["ColumnName"]
		columnNames = append(columnNames, columnName)
		fieldVars = append(fieldVars, strings.ToLower(structName)+"."+field["FieldName"])
	}

	// 拼接 INSERT INTO 语句
	query = fmt.Sprintf(`
func Insert%s(conn clickhouse.Conn, %ss []%s) error {
	batch, err := conn.PrepareBatch(context.Background(), "INSERT INTO %s (%s)")
	if err != nil {
		return fmt.Errorf("failed to prepare batch: %%w", err)
	}

	for _, %s := range %ss {
		if err := batch.Append(%s); err != nil {
			return fmt.Errorf("failed to append data: %%w", err)
		}
	}

	if err := batch.Send(); err != nil {
		return fmt.Errorf("failed to send batch: %%w", err)
	}
	return nil
}
`, structName, structName, structName, CamelToUnderline(structName), join(columnNames, ", "), strings.ToLower(structName), structName, join(fieldVars, ", "))

	return query
}

// join 是一个辅助函数，用于将字符串切片连接为一个由逗号分隔的字符串
func join(arr []string, separator string) string {
	return strings.Join(arr, separator)
}

// 更新源文件
func updateSourceFile(sourceFile string, generatedCode string) error {
	content, err := ioutil.ReadFile(sourceFile)
	if err != nil {
		return fmt.Errorf("failed to read source file: %w", err)
	}
	startMarker := "// Code generated by gen_clickhouse.go DO NOT EDIT.\n"
	endMarker := "// End of generated code"
	startIndex := bytes.Index(content, []byte(startMarker))
	endIndex := bytes.Index(content, []byte(endMarker))

	var updatedContent []byte
	if startIndex != -1 && endIndex != -1 {
		startIndex += len(startMarker)
		updatedContent = append(content[:startIndex-1], []byte(generatedCode)...)
		updatedContent = append(updatedContent, content[endIndex:]...)
	} else {
		updatedContent = append(content, []byte("\n"+startMarker+"\n")...)
		updatedContent = append(updatedContent, []byte(generatedCode)...)
		updatedContent = append(updatedContent, []byte(endMarker+"\n")...)
	}

	return ioutil.WriteFile(sourceFile, updatedContent, 0644)
}

// CamelToUnderline converts a camel case string to an underscore string.

func CamelToUnderline(s string) string {
	var result string
	for i, c := range s {
		if 'A' <= c && c <= 'Z' {
			if i != 0 {
				result += "_"
			}
			result += string(c + 32)
		} else {
			result += string(c)
		}
	}
	return result
}
