package main

import (
	"flag"
	"fmt"
	"github.com/pingcap/tidb/pkg/parser"
	"github.com/pingcap/tidb/pkg/parser/ast"
	"github.com/pingcap/tidb/pkg/parser/test_driver" // Required for parser to work
	"gopkg.in/yaml.v2"
	"io"
	"log"
	"os"
	"reflect"
	"strings"
)

type visitor struct {
	nodes []map[string]any
}

func (v *visitor) Enter(n ast.Node) (out ast.Node, skipChildren bool) {
	ndata := map[string]any{
		"_type":  reflect.TypeOf(n).Elem().Name(),
		"_phase": "enter",
	}

	switch node := n.(type) {
	case *ast.SelectStmt:
		ndata["distinct"] = node.Distinct
	case *ast.SelectField:
		ndata["offset"] = node.Offset
		ndata["as_name"] = node.AsName.String()
	case *ast.FuncCallExpr:
		ndata["schema"] = node.Schema.String()
		ndata["fn_name"] = node.FnName.String()
	case *ast.ColumnName:
		ndata["schema"] = node.Schema.String()
		ndata["table"] = node.Table.String()
		ndata["name"] = node.Name.String()
	case *ast.TableSource:
		ndata["as_name"] = node.AsName.String()
	case *ast.Join:
		using := []map[string]any{}
		for _, u := range node.Using {
			using = append(using, map[string]any{
				"schema": u.Schema.String(),
				"table":  u.Table.String(),
				"name":   u.Name.String(),
			})
		}
		ndata["using"] = using
	case *ast.TableName:
		ndata["schema"] = node.Schema.String()
		ndata["name"] = node.Name.String()
	case *ast.BinaryOperationExpr:
		ndata["op"] = node.Op.String()
	case *test_driver.ParamMarkerExpr:
		ndata["offset"] = node.Offset
		ndata["order"] = node.Order
		ndata["in_execute"] = node.InExecute
	}

	v.nodes = append(v.nodes, ndata)
	return n, false
}

func (v *visitor) Leave(n ast.Node) (out ast.Node, ok bool) {
	v.nodes = append(v.nodes, map[string]any{
		"_type":  reflect.TypeOf(n).Elem().Name(),
		"_phase": "leave",
	})
	return n, true
}

func main() {
	flag.Parse()
	var inputSQL string
	if flag.NArg() > 0 {
		inputSQL = strings.Join(flag.Args(), " ")
	} else {
		bytes, err := io.ReadAll(os.Stdin)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error reading from stdin: %v\n", err)
			os.Exit(1)
		}
		inputSQL = string(bytes)
	}
	p := parser.New()
	stmts, _, err := p.Parse(inputSQL, "", "")
	if err != nil {
		log.Fatalf("error: %v", err)
	}
	if len(stmts) == 0 {
		fmt.Println("No SQL statements found to process.")
		return
	}
	// Iterate and tag each statement's AST.
	for _, stmt := range stmts {
		v := visitor{nodes: make([]map[string]any, 0)}
		stmt.Accept(&v)
		bytes, err := yaml.Marshal(v.nodes)
		if err != nil {
			log.Fatalf("error: %v", err)
		}
		fmt.Printf(string(bytes))
	}
}