package engine

import (
	"fmt"
	"strings"

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

func trimSQLIdent(s string) string {
	return strings.Trim(s, `"'`+"`")
}

func cteMapKey(name ast.CIStr) string {
	return strings.ToLower(trimSQLIdent(name.O))
}

func cloneCTEScope(m map[string]struct{}) map[string]struct{} {
	if len(m) == 0 {
		return map[string]struct{}{}
	}
	o := make(map[string]struct{}, len(m)+4)
	for k := range m {
		o[k] = struct{}{}
	}
	return o
}

func processWithClause(w *ast.WithClause, cteOuter map[string]struct{}, refs *[]TableReference, seen map[string]bool, fn func(visibleMain map[string]struct{})) {
	if w == nil {
		fn(cteOuter)
		return
	}
	visibleMain := cloneCTEScope(cteOuter)
	for _, cte := range w.CTEs {
		visibleMain[cteMapKey(cte.Name)] = struct{}{}
	}
	for i, cte := range w.CTEs {
		visBody := cloneCTEScope(cteOuter)
		for j := 0; j < i; j++ {
			visBody[cteMapKey(w.CTEs[j].Name)] = struct{}{}
		}
		if cte.IsRecursive {
			visBody[cteMapKey(cte.Name)] = struct{}{}
		}
		if cte.Query != nil && cte.Query.Query != nil {
			walkResultSet(cte.Query.Query, visBody, refs, seen)
		}
	}
	fn(visibleMain)
}

func addTableName(tn *ast.TableName, cteVisible map[string]struct{}, refs *[]TableReference, seen map[string]bool) {
	if tn == nil || tn.IsAlias {
		return
	}
	schema := trimSQLIdent(tn.Schema.String())
	name := trimSQLIdent(tn.Name.String())
	if schema == "" && name != "" {
		if _, ok := cteVisible[strings.ToLower(name)]; ok {
			return
		}
	}
	key := schema + "." + name
	if seen[key] {
		return
	}
	seen[key] = true
	*refs = append(*refs, TableReference{Database: schema, Table: name})
}

func walkJoin(j *ast.Join, cteVisible map[string]struct{}, refs *[]TableReference, seen map[string]bool) {
	if j == nil {
		return
	}
	walkResultSet(j.Left, cteVisible, refs, seen)
	if j.Right != nil {
		walkResultSet(j.Right, cteVisible, refs, seen)
	}
	if j.On != nil && j.On.Expr != nil {
		walkExpr(j.On.Expr, cteVisible, refs, seen)
	}
}

func walkResultSet(rs ast.ResultSetNode, cteVisible map[string]struct{}, refs *[]TableReference, seen map[string]bool) {
	if rs == nil {
		return
	}
	switch r := rs.(type) {
	case *ast.SelectStmt:
		extractFromSelectStmt(r, cteVisible, refs, seen)
	case *ast.SetOprStmt:
		extractFromSetOprStmt(r, cteVisible, refs, seen)
	case *ast.TableSource:
		switch src := r.Source.(type) {
		case *ast.TableName:
			addTableName(src, cteVisible, refs, seen)
		case *ast.SelectStmt:
			extractFromSelectStmt(src, cteVisible, refs, seen)
		case *ast.SetOprStmt:
			extractFromSetOprStmt(src, cteVisible, refs, seen)
		case *ast.SubqueryExpr:
			if src.Query != nil {
				walkResultSet(src.Query, cteVisible, refs, seen)
			}
		}
	case *ast.Join:
		walkJoin(r, cteVisible, refs, seen)
	}
}

type tableRefExprWalker struct {
	cteVisible map[string]struct{}
	refs       *[]TableReference
	seen       map[string]bool
}

func (w *tableRefExprWalker) Enter(n ast.Node) (ast.Node, bool) {
	switch x := n.(type) {
	case *ast.SelectStmt:
		extractFromSelectStmt(x, w.cteVisible, w.refs, w.seen)
		return n, true
	case *ast.SetOprStmt:
		extractFromSetOprStmt(x, w.cteVisible, w.refs, w.seen)
		return n, true
	case *ast.SubqueryExpr:
		if x.Query != nil {
			walkResultSet(x.Query, w.cteVisible, w.refs, w.seen)
		}
		return n, true
	case *ast.TableName:
		addTableName(x, w.cteVisible, w.refs, w.seen)
		return n, true
	}
	return n, false
}

func (w *tableRefExprWalker) Leave(n ast.Node) (ast.Node, bool) {
	return n, true
}

func walkExpr(e ast.ExprNode, cteVisible map[string]struct{}, refs *[]TableReference, seen map[string]bool) {
	if e == nil {
		return
	}
	w := &tableRefExprWalker{cteVisible: cteVisible, refs: refs, seen: seen}
	e.Accept(w)
}

func walkOrderByAndLimit(ob *ast.OrderByClause, lim *ast.Limit, cteVisible map[string]struct{}, refs *[]TableReference, seen map[string]bool) {
	if ob != nil {
		for _, item := range ob.Items {
			walkExpr(item.Expr, cteVisible, refs, seen)
		}
	}
	if lim != nil {
		walkExpr(lim.Count, cteVisible, refs, seen)
		walkExpr(lim.Offset, cteVisible, refs, seen)
	}
}

func extractFromSelectStmt(s *ast.SelectStmt, cteOuter map[string]struct{}, refs *[]TableReference, seen map[string]bool) {
	if s == nil {
		return
	}
	processWithClause(s.With, cteOuter, refs, seen, func(visibleMain map[string]struct{}) {
		if s.From != nil && s.From.TableRefs != nil {
			walkJoin(s.From.TableRefs, visibleMain, refs, seen)
		}
		walkExpr(s.Where, visibleMain, refs, seen)
		if s.Having != nil {
			walkExpr(s.Having.Expr, visibleMain, refs, seen)
		}
		if s.Fields != nil {
			for _, f := range s.Fields.Fields {
				if f.Expr != nil {
					walkExpr(f.Expr, visibleMain, refs, seen)
				}
			}
		}
		if s.GroupBy != nil {
			for _, item := range s.GroupBy.Items {
				walkExpr(item.Expr, visibleMain, refs, seen)
			}
		}
		walkOrderByAndLimit(s.OrderBy, s.Limit, visibleMain, refs, seen)
	})
}

func extractFromSetOprSelectList(list *ast.SetOprSelectList, cteOuter map[string]struct{}, refs *[]TableReference, seen map[string]bool) {
	if list == nil {
		return
	}
	processWithClause(list.With, cteOuter, refs, seen, func(visibleMain map[string]struct{}) {
		for _, node := range list.Selects {
			switch n := node.(type) {
			case *ast.SelectStmt:
				extractFromSelectStmt(n, visibleMain, refs, seen)
			case *ast.SetOprSelectList:
				extractFromSetOprSelectList(n, visibleMain, refs, seen)
			}
		}
		walkOrderByAndLimit(list.OrderBy, list.Limit, visibleMain, refs, seen)
	})
}

func extractFromSetOprStmt(s *ast.SetOprStmt, cteOuter map[string]struct{}, refs *[]TableReference, seen map[string]bool) {
	if s == nil {
		return
	}
	var unionVisible map[string]struct{}
	processWithClause(s.With, cteOuter, refs, seen, func(visibleMain map[string]struct{}) {
		unionVisible = visibleMain
		if s.SelectList != nil {
			extractFromSetOprSelectList(s.SelectList, visibleMain, refs, seen)
		}
	})
	walkOrderByAndLimit(s.OrderBy, s.Limit, unionVisible, refs, seen)
}

func extractTableReferencesAST(sql string) ([]TableReference, error) {
	p := parser.New()
	stmt, err := p.ParseOneStmt(sql, "", "")
	if err != nil {
		return nil, err
	}
	var refs []TableReference
	seen := make(map[string]bool)
	outer := map[string]struct{}{}
	switch s := stmt.(type) {
	case *ast.SelectStmt:
		extractFromSelectStmt(s, outer, &refs, seen)
	case *ast.SetOprStmt:
		extractFromSetOprStmt(s, outer, &refs, seen)
	default:
		return nil, fmt.Errorf("unsupported statement type %T", stmt)
	}
	return refs, nil
}