package query

import (
	"fmt"
	"strings"

	"github.com/pingcap/tidb/pkg/parser"
	"github.com/pingcap/tidb/pkg/parser/ast"
	_ "github.com/pingcap/tidb/pkg/parser/test_driver"
)

// ExtractTableReferences returns database.table pairs referenced in a SELECT.
func ExtractTableReferences(sql, defaultDatabase string) ([]TableRef, error) {
	sql = strings.TrimSpace(strings.TrimSuffix(strings.TrimSpace(sql), ";"))
	if sql == "" {
		return nil, fmt.Errorf("sql is required")
	}
	p := parser.New()
	stmts, _, err := p.Parse(sql, "", "")
	if err != nil {
		return nil, fmt.Errorf("invalid sql: %w", err)
	}
	if len(stmts) != 1 {
		return nil, fmt.Errorf("only a single SELECT statement is allowed")
	}
	visitor := &tableRefVisitor{defaultDB: strings.TrimSpace(defaultDatabase)}
	switch stmt := stmts[0].(type) {
	case *ast.SelectStmt:
		stmt.Accept(visitor)
	case *ast.SetOprStmt:
		stmt.Accept(visitor)
	default:
		return nil, fmt.Errorf("only SELECT statements are allowed")
	}
	return visitor.refs(), nil
}

type tableRefVisitor struct {
	defaultDB string
	seen      map[string]struct{}
	refsList  []TableRef
}

func (v *tableRefVisitor) refs() []TableRef {
	return append([]TableRef(nil), v.refsList...)
}

func (v *tableRefVisitor) add(db, table string) {
	table = strings.TrimSpace(table)
	if table == "" {
		return
	}
	db = strings.TrimSpace(db)
	if db == "" {
		db = v.defaultDB
	}
	if db == "" {
		return
	}
	key := strings.ToLower(db) + "." + strings.ToLower(table)
	if v.seen == nil {
		v.seen = map[string]struct{}{}
	}
	if _, ok := v.seen[key]; ok {
		return
	}
	v.seen[key] = struct{}{}
	v.refsList = append(v.refsList, TableRef{Database: db, Table: table})
}

func (v *tableRefVisitor) Enter(n ast.Node) (ast.Node, bool) {
	if ts, ok := n.(*ast.TableSource); ok {
		if name, ok := ts.Source.(*ast.TableName); ok {
			v.add(name.Schema.O, name.Name.O)
		}
	}
	return n, false
}

func (v *tableRefVisitor) Leave(n ast.Node) (ast.Node, bool) {
	return n, true
}

// RegistryTableRefs returns all built-in tables as references.
func RegistryTableRefs(tables []ResolvedTable) []TableRef {
	out := make([]TableRef, 0, len(tables))
	for _, t := range tables {
		out = append(out, TableRef{Database: t.Database, Table: t.Name})
	}
	return out
}