package analytics

import (
	"context"
	"fmt"
	"strings"

	apperrors "github.com/nexuscrm/backend/pkg/errors"
	"github.com/nexuscrm/backend/internal/features/permissions"
	"github.com/nexuscrm/shared/pkg/constants"
	"github.com/nexuscrm/shared/pkg/models"

	"github.com/pingcap/tidb/pkg/parser"
	"github.com/pingcap/tidb/pkg/parser/ast"
	"github.com/pingcap/tidb/pkg/parser/format"
	"github.com/pingcap/tidb/pkg/parser/opcode"
	"github.com/pingcap/tidb/pkg/parser/test_driver" // Using test_driver for ValueExpr
)

// PermissionChecker defines the subset of PermissionService needed by SecurityValidator
type PermissionChecker interface {
	CheckObjectPermissionWithUser(ctx context.Context, objectName string, permission string, user *models.UserSession) bool
	CheckFieldVisibilityWithUser(ctx context.Context, objectName string, fieldName string, user *models.UserSession) bool
	GetSubordinateUserIDs(ctx context.Context, roleID string) ([]string, error)
	BuildRLSResult(ctx context.Context, schema *models.ObjectMetadata, objectName string, user *models.UserSession) *permissions.RLSResult
}

// MetadataProvider defines the subset of MetadataService needed by SecurityValidator
type MetadataProvider interface {
	GetSchema(ctx context.Context, objectName string) *models.ObjectMetadata
}

// SecurityValidator handles SQL parsing and security enforcement
type SecurityValidator struct {
	permissions PermissionChecker
	metadata    MetadataProvider
}

// NewSecurityValidator creates a new SecurityValidator
func NewSecurityValidator(permissions PermissionChecker, metadata MetadataProvider) *SecurityValidator {
	return &SecurityValidator{
		permissions: permissions,
		metadata:    metadata,
	}
}

// ValidateAndRewrite parses the SQL, validates permissions, and rewrites it for RLS.
// contextVars is an optional map for named variable substitution (e.g., "start_date", "end_date").
// Security-sensitive variables like "user_id" are always injected from the user session, not contextVars.
func (v *SecurityValidator) ValidateAndRewrite(ctx context.Context, sql string, params []interface{}, user *models.UserSession, contextVars map[string]string) (string, []interface{}, error) {

	// 1. Parse SQL
	// TiDB parser is not thread-safe, so we create a new instance for each request
	p := parser.New()
	p.EnableWindowFunc(true)
	stmtNodes, _, err := p.Parse(sql, "", "")
	if err != nil {
		return "", nil, apperrors.NewValidationError("SQL", fmt.Sprintf("SQL parse error: %v", err))
	}

	if len(stmtNodes) != 1 {
		return "", nil, apperrors.NewValidationError("SQL", "only single SQL statements are allowed")
	}

	stmt := stmtNodes[0]

	// 2. Only allow SELECT or UNION statements
	var selectStmt *ast.SelectStmt
	var setOprStmt *ast.SetOprStmt
	var isUnion bool

	if s, ok := stmt.(*ast.SelectStmt); ok {
		selectStmt = s
	} else if s, ok := stmt.(*ast.SetOprStmt); ok {
		setOprStmt = s
		isUnion = true
	} else {
		return "", nil, apperrors.NewValidationError("SQL", "only SELECT or UNION statements are allowed in analytics")
	}

	// 3. Visitor for Validation AND Transformation
	// Extract primary table name to handle implicit column references
	var defaultTableName string
	if !isUnion {
		if selectStmt.From != nil && selectStmt.From.TableRefs != nil && selectStmt.From.TableRefs.Left != nil {
			if ts, ok := selectStmt.From.TableRefs.Left.(*ast.TableSource); ok {
				if tn, ok := ts.Source.(*ast.TableName); ok {
					defaultTableName = tn.Name.O
				}
			}
		}
	} else if setOprStmt.SelectList != nil && len(setOprStmt.SelectList.Selects) > 0 {
		// Attempt to extract from the first SELECT in the UNION
		if s, ok := setOprStmt.SelectList.Selects[0].(*ast.SelectStmt); ok {
			if s.From != nil && s.From.TableRefs != nil && s.From.TableRefs.Left != nil {
				if ts, ok := s.From.TableRefs.Left.(*ast.TableSource); ok {
					if tn, ok := ts.Source.(*ast.TableName); ok {
						defaultTableName = tn.Name.O
					}
				}
			}
		}
	}

	visitor := &SecurityVisitor{
		ctx:              ctx,
		user:             user,
		permissions:      v.permissions,
		metadata:         v.metadata,
		defaultTableName: defaultTableName,
		contextVars:      contextVars, // NEW: Pass context vars for substitution
		err:              nil,
	}

	// Capture the transformed node
	newNode, _ := stmt.Accept(visitor)
	if s, ok := newNode.(ast.StmtNode); ok {
		stmt = s
	}

	if visitor.err != nil {
		return "", nil, visitor.err
	}

	// Warn about unknown @variables (helps catch typos like @start_dat instead of @start_date)
	if len(visitor.unknownVars) > 0 {
		return "", nil, apperrors.NewValidationError("SQL",
			fmt.Sprintf("unknown SQL variable(s): @%s. Supported variables: @user_id, @user_email, and context_vars keys (start_date, end_date)",
				strings.Join(visitor.unknownVars, ", @")))
	}

	// 4. Transform: Inject RLS Logic if needed
	// BuildRLSResult (called inside applyRLS) is the single source of truth for RLS decisions,
	// including SuperUser/admin bypass. No outer guard needed.
	if isUnion {
		for _, s := range setOprStmt.SelectList.Selects {
			if sel, ok := s.(*ast.SelectStmt); ok {
				if err := v.applyRLS(ctx, sel, user); err != nil {
					return "", nil, err
				}
			}
		}
	} else {
		if err := v.applyRLS(ctx, selectStmt, user); err != nil {
			return "", nil, err
		}
	}

	// 5. Restore SQL
	var sb strings.Builder
	restoreCtx := format.NewRestoreCtx(format.DefaultRestoreFlags, &sb)
	if err := stmt.Restore(restoreCtx); err != nil {
		return "", nil, apperrors.NewInternalError("SQL restore error", err)
	}

	// Note: We currently inject owner_id as a literal string in AST to avoid shifting `params`.
	// This is safe because user.ID is a trusted internal UUID.

	return sb.String(), params, nil
}

// extractAllTableSources recursively walks a Join/TableSource tree and returns all TableSources.
func (v *SecurityValidator) extractAllTableSources(node ast.ResultSetNode) []*ast.TableSource {
	if node == nil {
		return nil
	}
	switch n := node.(type) {
	case *ast.TableSource:
		return []*ast.TableSource{n}
	case *ast.Join:
		left := v.extractAllTableSources(n.Left)
		right := v.extractAllTableSources(n.Right)
		return append(left, right...)
	}
	return nil
}

// applyRLS injects RLS filtering into the WHERE clause for all tables in the FROM clause.
// Uses BuildRLSResult + RLSResultToFilterExpr as the single source of truth.
func (v *SecurityValidator) applyRLS(ctx context.Context, stmt *ast.SelectStmt, user *models.UserSession) error {
	if stmt.From == nil || stmt.From.TableRefs == nil || stmt.From.TableRefs.Left == nil {
		return nil
	}

	// Extract all table sources (handles both simple tables and JOINs)
	sources := v.extractAllTableSources(stmt.From.TableRefs)
	if len(sources) == 0 {
		return nil
	}

	for _, ts := range sources {
		tn, ok := ts.Source.(*ast.TableName)
		if !ok {
			continue // subquery or other complex source, skip
		}
		objName := tn.Name.O
		tableAlias := ts.AsName.O

		schema := v.metadata.GetSchema(ctx, objName)
		if schema == nil {
			return apperrors.NewPermissionError("read", fmt.Sprintf("unknown table '%s'", objName))
		}

		rlsResult := v.permissions.BuildRLSResult(ctx, schema, objName, user)
		// Use table alias if present (e.g., SELECT ... FROM Account a),
		// otherwise the physical table name from the SQL AST.
		qualifier := objName
		if tableAlias != "" {
			qualifier = tableAlias
		}
		filterExpr := permissions.RLSResultToFilterExpr(rlsResult, qualifier)
		if filterExpr == "" {
			continue
		}

		// Parse the filter expression string into an AST condition node
		cond, err := parseFilterExprToAST(filterExpr, tableAlias)
		if err != nil {
			return apperrors.NewInternalError("RLS filter parse error", err)
		}

		if stmt.Where == nil {
			stmt.Where = cond
		} else {
			stmt.Where = &ast.BinaryOperationExpr{
				Op: opcode.LogicAnd,
				L:  stmt.Where,
				R:  cond,
			}
		}
	}

	return nil
}

// parseFilterExprToAST parses a SQL WHERE fragment string into an AST ExprNode.
// tableAlias is used to qualify the OwnerId column when non-empty.
func parseFilterExprToAST(filterExpr string, tableAlias string) (ast.ExprNode, error) {
	p := parser.New()
	sql := fmt.Sprintf("SELECT 1 WHERE %s", filterExpr)
	stmts, _, err := p.Parse(sql, "", "")
	if err != nil {
		return nil, fmt.Errorf("parseFilterExprToAST: %w", err)
	}
	if len(stmts) != 1 {
		return nil, fmt.Errorf("parseFilterExprToAST: expected 1 statement, got %d", len(stmts))
	}
	sel, ok := stmts[0].(*ast.SelectStmt)
	if !ok || sel.Where == nil {
		return nil, fmt.Errorf("parseFilterExprToAST: no WHERE clause parsed")
	}

	// If a table alias is provided, qualify all unqualified column references
	// so criteria-based sharing rule predicates bind to the correct table in joins
	if tableAlias != "" {
		qualifyUnqualifiedColumns(sel.Where, tableAlias)
	}

	return sel.Where, nil
}

// qualifyUnqualifiedColumns walks an AST expression and qualifies ALL bare column
// references with the given table alias. This is critical for joined queries where
// criteria-based sharing rule fragments (e.g., `Name = 'Acme'`) would otherwise
// be ambiguous or bind to the wrong table. Columns that already have a table
// qualifier (e.g., subquery references like rs.RecordId) are left untouched.
func qualifyUnqualifiedColumns(node ast.Node, tableAlias string) {
	if node == nil {
		return
	}
	node.Accept(&columnQualifier{tableAlias: tableAlias})
}

type columnQualifier struct {
	tableAlias string
}

func (q *columnQualifier) Enter(in ast.Node) (ast.Node, bool) {
	if col, ok := in.(*ast.ColumnNameExpr); ok {
		if col.Name.Table.O == "" {
			col.Name.Table = ast.NewCIStr(q.tableAlias)
		}
	}
	return in, false
}

func (q *columnQualifier) Leave(in ast.Node) (ast.Node, bool) {
	return in, true
}

type SecurityVisitor struct {
	ctx              context.Context
	user             *models.UserSession
	permissions      PermissionChecker
	metadata         MetadataProvider
	defaultTableName string
	contextVars      map[string]string // Named variables like start_date, end_date
	unknownVars      []string          // Track unknown @variables for warning
	err              error
}

func (v *SecurityVisitor) Enter(in ast.Node) (ast.Node, bool) {
	if v.err != nil {
		return in, true
	}

	// Validate Table Permissions
	if t, ok := in.(*ast.TableName); ok {
		objName := t.Name.O
		if objName != "" {
			if !v.permissions.CheckObjectPermissionWithUser(v.ctx, objName, constants.PermRead, v.user) {
				v.err = apperrors.NewPermissionError("read", fmt.Sprintf("table '%s'", objName))
				return in, true
			}
		}
	}

	// Validate Column Permissions
	if c, ok := in.(*ast.ColumnName); ok {
		tableName := c.Table.O
		if tableName == "" {
			tableName = v.defaultTableName
		}

		colName := c.Name.O

		// If Table is known, check specific field visibility
		if tableName != "" && colName != "*" {
			if !v.permissions.CheckFieldVisibilityWithUser(v.ctx, tableName, colName, v.user) {
				v.err = apperrors.NewPermissionError("read", fmt.Sprintf("field '%s' on table '%s'", colName, tableName))
				return in, true
			}
		}
	}
	// Substitute System Variables (@user_id, @user_email) and Context Variables (@start_date, @end_date)
	if v.err == nil {
		if valExpr, ok := in.(*ast.VariableExpr); ok {
			varName := strings.ToLower(valExpr.Name)

			// SECURITY: user_id and user_email are ALWAYS sourced from the authenticated session.
			// They cannot be overridden by contextVars to prevent privilege escalation.
			if varName == "user_id" {
				newVal := &test_driver.ValueExpr{}
				newVal.SetString(v.user.ID)
				return newVal, true
			}
			if varName == "user_email" && v.user.Email != nil {
				newVal := &test_driver.ValueExpr{}
				newVal.SetString(*v.user.Email)
				return newVal, true
			}

			// Check contextVars for other named variables (e.g., start_date, end_date)
			if v.contextVars != nil {
				if val, exists := v.contextVars[varName]; exists {
					newVal := &test_driver.ValueExpr{}
					newVal.SetString(val)
					return newVal, true
				}
			}

			// Track unknown variables for warning (helps catch typos like @start_dat)
			v.unknownVars = append(v.unknownVars, varName)
		}
	}

	return in, false
}

func (v *SecurityVisitor) Leave(in ast.Node) (ast.Node, bool) {
	return in, true
}