// Copyright 2026 PingCAP, Inc. Licensed under Apache-2.0.

package export

import (
	"bytes"
	"strings"

	"github.com/pingcap/errors"
	"github.com/pingcap/tidb/pkg/ddl"
	"github.com/pingcap/tidb/pkg/meta/metabuild"
	"github.com/pingcap/tidb/pkg/meta/model"
	"github.com/pingcap/tidb/pkg/parser"
	"github.com/pingcap/tidb/pkg/parser/ast"
	"github.com/pingcap/tidb/pkg/parser/format"
	parsermodel "github.com/pingcap/tidb/pkg/parser/model"
	"github.com/pingcap/tidb/pkg/parser/mysql"
	_ "github.com/pingcap/tidb/pkg/planner/core" // initialize expression builders used by DDL validation
)

type projectedTableSchema struct {
	createTable     *ast.CreateTableStmt
	retainedColumns map[string]struct{}
	validationInfo  *model.TableInfo
}

type projectedTableSchemas map[tableName]*projectedTableSchema

func normalizedTableName(database, table string) tableName {
	return tableName{db: strings.ToLower(database), table: strings.ToLower(table)}
}

func buildProjectedTableSchema(
	p *parser.Parser,
	originSQL string,
	selectedColumns []string,
) (*projectedTableSchema, error) {
	stmt, err := p.ParseOneStmt(originSQL, "", "")
	if err != nil {
		return nil, errors.Annotate(err, "failed to parse CREATE TABLE for column projection")
	}
	createTable, ok := stmt.(*ast.CreateTableStmt)
	if !ok {
		return nil, errors.Errorf("expected CREATE TABLE for column projection, got %T", stmt)
	}
	partitionColumns := collectPartitionColumns(createTable)

	retainedColumns := make(map[string]struct{}, len(selectedColumns))
	for _, selectedColumn := range selectedColumns {
		retainedColumns[strings.ToLower(selectedColumn)] = struct{}{}
	}
	// A generated column can only depend on generated columns defined before it.
	for _, column := range createTable.Cols {
		for _, option := range column.Options {
			if option.Tp == ast.ColumnOptionGenerated && usesOnlyRetainedColumns(option.Expr, retainedColumns) {
				retainedColumns[column.Name.Name.L] = struct{}{}
				break
			}
		}
	}

	columns := make([]*ast.ColumnDef, 0, len(createTable.Cols))
	for _, column := range createTable.Cols {
		if _, ok := retainedColumns[column.Name.Name.L]; !ok {
			continue
		}
		options, err := filterColumnOptions(column, retainedColumns)
		if err != nil {
			return nil, err
		}
		column.Options = options
		columns = append(columns, column)
	}
	createTable.Cols = columns

	constraints := make([]*ast.Constraint, 0, len(createTable.Constraints))
	for _, constraint := range createTable.Constraints {
		if filterTableConstraint(constraint, retainedColumns) {
			constraints = append(constraints, constraint)
		}
	}
	createTable.Constraints = constraints

	if !allColumnsRetained(partitionColumns, retainedColumns) {
		return nil, errors.New("partition definition references a removed column")
	}
	if err := validateTTLColumns(createTable.Options, retainedColumns); err != nil {
		return nil, err
	}
	return &projectedTableSchema{
		createTable:     createTable,
		retainedColumns: retainedColumns,
	}, nil
}

func (s *projectedTableSchema) buildValidationInfo() (*model.TableInfo, error) {
	if s.validationInfo != nil {
		return s.validationInfo, nil
	}

	createTable := *s.createTable
	// The schema SQL is restored before this validation-only build mutates the AST.
	createTable.Constraints = make([]*ast.Constraint, 0, len(s.createTable.Constraints))
	createTable.Options = make([]*ast.TableOption, 0, len(s.createTable.Options))
	// Source metadata validation must not enforce CSE's fixed TTL job interval.
	for _, option := range s.createTable.Options {
		if option.Tp != ast.TableOptionTTLJobInterval {
			createTable.Options = append(createTable.Options, option)
		}
	}
	// Columnar indexes cannot support foreign key lookups.
	for _, constraint := range s.createTable.Constraints {
		switch constraint.Tp {
		case ast.ConstraintColumnar, ast.ConstraintFulltext, ast.ConstraintVector:
			continue
		default:
			createTable.Constraints = append(createTable.Constraints, constraint)
		}
	}
	validationInfo, err := ddl.BuildTableInfoWithStmt(
		metabuild.NewContext(),
		&createTable,
		mysql.DefaultCharset,
		"",
		nil,
	)
	if err != nil {
		return nil, err
	}
	s.validationInfo = validationInfo
	return validationInfo, nil
}

func restoreProjectedSchema(createTable *ast.CreateTableStmt) (string, error) {
	var buffer bytes.Buffer
	err := createTable.Restore(&format.RestoreCtx{
		Flags: format.DefaultRestoreFlags | format.RestoreTiDBSpecialComment,
		In:    &buffer,
	})
	if err != nil {
		return "", errors.Annotate(err, "failed to restore projected CREATE TABLE")
	}
	return buffer.String(), nil
}

func filterColumnOptions(
	column *ast.ColumnDef,
	retained map[string]struct{},
) ([]*ast.ColumnOption, error) {
	options := make([]*ast.ColumnOption, 0, len(column.Options))
	for _, option := range column.Options {
		switch option.Tp {
		case ast.ColumnOptionCheck:
			// CHECK constraints can be removed as a whole; expressions on retained columns cannot.
			if !usesOnlyRetainedColumns(option.Expr, retained) {
				continue
			}
		case ast.ColumnOptionDefaultValue, ast.ColumnOptionOnUpdate:
			if !usesOnlyRetainedColumns(option.Expr, retained) {
				return nil, errors.Errorf(
					"column `%s` expression references a removed column",
					column.Name.Name.O,
				)
			}
		}
		options = append(options, option)
	}
	return options, nil
}

func filterTableConstraint(
	constraint *ast.Constraint,
	retained map[string]struct{},
) bool {
	for _, key := range constraint.Keys {
		if key.Column != nil {
			if _, ok := retained[key.Column.Name.L]; !ok {
				return false
			}
		}
		if !usesOnlyRetainedColumns(key.Expr, retained) {
			return false
		}
	}
	if !usesOnlyRetainedColumns(constraint.Expr, retained) {
		return false
	}
	return true
}

func validateForeignKeyParents(
	childDB string,
	child *projectedTableSchema,
	schemas projectedTableSchemas,
) error {
	for _, column := range child.createTable.Cols {
		for _, option := range column.Options {
			if option.Tp != ast.ColumnOptionReference {
				continue
			}
			// MySQL 9.7 supports inline foreign keys as column-level REFERENCES options.
			if err := validateForeignKeyParent(option.Refer, childDB, schemas); err != nil {
				return err
			}
		}
	}
	for _, constraint := range child.createTable.Constraints {
		if err := validateForeignKeyParent(constraint.Refer, childDB, schemas); err != nil {
			return err
		}
	}
	return nil
}

func validateForeignKeyParent(reference *ast.ReferenceDef, childDB string, schemas projectedTableSchemas) error {
	if reference == nil || reference.Table == nil {
		return nil
	}

	parentTable := reference.Table.Name.O
	parentDB := reference.Table.Schema.O
	if parentDB == "" {
		parentDB = childDB
	}

	parent, ok := schemas[normalizedTableName(parentDB, parentTable)]
	if !ok {
		// The referenced table is outside this dump and is not rewritten here.
		return nil
	}
	parentColumns := make([]parsermodel.CIStr, 0, len(reference.IndexPartSpecifications))
	// Foreign key references from SHOW CREATE TABLE contain only column index parts.
	for _, key := range reference.IndexPartSpecifications {
		if _, ok := parent.retainedColumns[key.Column.Name.L]; !ok {
			return errors.Errorf(
				"foreign key references removed column `%s`.`%s`.`%s`",
				escapeString(parentDB),
				escapeString(parentTable),
				escapeString(key.Column.Name.O),
			)
		}
		parentColumns = append(parentColumns, key.Column.Name)
	}
	parentInfo, err := parent.buildValidationInfo()
	if err != nil {
		return err
	}
	if !hasParentIndex(parentInfo, parentColumns) {
		return errors.Errorf(
			"foreign key referenced columns are not indexed in table `%s`.`%s`",
			escapeString(parentDB),
			escapeString(parentTable),
		)
	}
	return nil
}

func hasParentIndex(parentInfo *model.TableInfo, parentColumns []parsermodel.CIStr) bool {
	if len(parentColumns) == 1 {
		column := model.FindColumnInfo(parentInfo.Columns, parentColumns[0].L)
		if column != nil && parentInfo.PKIsHandle && mysql.HasPriKeyFlag(column.GetFlag()) {
			return true
		}
	}
	return model.FindIndexByColumns(parentInfo, parentInfo.Indices, parentColumns...) != nil
}

func collectPartitionColumns(createTable *ast.CreateTableStmt) map[string]struct{} {
	columns := make(map[string]struct{})
	if createTable.Partition == nil {
		return columns
	}

	methods := []*ast.PartitionMethod{&createTable.Partition.PartitionMethod}
	if createTable.Partition.Sub != nil {
		methods = append(methods, createTable.Partition.Sub)
	}
	for _, method := range methods {
		for _, column := range method.ColumnNames {
			columns[column.Name.L] = struct{}{}
		}
		if method.Expr != nil {
			for _, column := range ddl.FindColumnNamesInExpr(method.Expr) {
				columns[column.Name.L] = struct{}{}
			}
		}
		if method.Tp != parsermodel.PartitionTypeKey || len(method.ColumnNames) != 0 {
			continue
		}

		for _, column := range createTable.Cols {
			for _, option := range column.Options {
				if option.Tp == ast.ColumnOptionPrimaryKey {
					columns[column.Name.Name.L] = struct{}{}
					break
				}
			}
		}
		for _, constraint := range createTable.Constraints {
			if constraint.Tp != ast.ConstraintPrimaryKey {
				continue
			}
			for _, key := range constraint.Keys {
				if key.Column != nil {
					columns[key.Column.Name.L] = struct{}{}
				}
			}
		}
	}
	return columns
}

func validateTTLColumns(options []*ast.TableOption, retained map[string]struct{}) error {
	for _, option := range options {
		if option.Tp == ast.TableOptionTTL && option.ColumnName != nil {
			if _, ok := retained[option.ColumnName.Name.L]; !ok {
				return errors.Errorf("TTL definition references removed column `%s`", option.ColumnName.Name.O)
			}
		}
	}
	return nil
}

func usesOnlyRetainedColumns(expr ast.ExprNode, retainedColumns map[string]struct{}) bool {
	if expr == nil {
		return true
	}
	for _, column := range ddl.FindColumnNamesInExpr(expr) {
		if _, ok := retainedColumns[column.Name.L]; !ok {
			return false
		}
	}
	return true
}

func allColumnsRetained(columns, retained map[string]struct{}) bool {
	for column := range columns {
		if _, ok := retained[column]; !ok {
			return false
		}
	}
	return true
}