package gudusoft.gsqlparser.resolver2.namespace;

import gudusoft.gsqlparser.nodes.TObjectNameList;
import gudusoft.gsqlparser.nodes.TResultColumn;
import gudusoft.gsqlparser.nodes.TResultColumnList;
import gudusoft.gsqlparser.nodes.TTable;
import gudusoft.gsqlparser.resolver2.ColumnLevel;
import gudusoft.gsqlparser.resolver2.matcher.INameMatcher;
import gudusoft.gsqlparser.resolver2.matcher.VendorNameMatcher;
import gudusoft.gsqlparser.resolver2.model.ColumnSource;
import gudusoft.gsqlparser.sqlenv.ESQLDataObjectType;
import gudusoft.gsqlparser.stmt.TSelectSqlStatement;

import java.util.*;

/**
 * Namespace representing a UNION/INTERSECT/EXCEPT query.
 *
 * Key characteristics:
 * - Schema is defined by the FIRST branch (SQL standard)
 * - Columns must be pushed down to ALL branches
 * - hasStarColumn() returns true if ANY branch has SELECT *
 * - addInferredColumn() propagates to ALL branches
 *
 * Example:
 * FROM (
 *     SELECT * FROM table_1
 *     UNION ALL
 *     SELECT * FROM table_2
 *     UNION ALL
 *     SELECT * FROM table_3
 * ) Combined
 *
 * When outer query references "col_1", it should be pushed to ALL branches.
 */
public class UnionNamespace extends AbstractNamespace {

    private final TSelectSqlStatement unionQuery;
    private final String alias;

    /** All branches of the UNION (flattened) */
    private final List<TSelectSqlStatement> allBranches;

    /** The derived table's alias column list, {@code (... UNION ALL ...) AS r(rank, type)}: renames the
     *  columns by position (Mantis 4814). Null when absent. */
    private TObjectNameList aliasColumns;

    public void setAliasColumns(TObjectNameList aliasColumns) {
        this.aliasColumns = (aliasColumns != null && aliasColumns.size() > 0) ? aliasColumns : null;
    }

    public TObjectNameList getAliasColumns() {
        return aliasColumns;
    }

    /** Namespace for each branch */
    private final List<SubqueryNamespace> branchNamespaces;

    /** Inferred columns from star push-down */
    private Map<String, ColumnSource> inferredColumns;

    /** Track inferred column names */
    private Set<String> inferredColumnNames;

    public UnionNamespace(TSelectSqlStatement unionQuery,
                          String alias,
                          INameMatcher nameMatcher) {
        super(unionQuery, nameMatcher);
        this.unionQuery = unionQuery;
        this.alias = alias;

        // Flatten all UNION branches
        this.allBranches = new ArrayList<>();
        flattenUnionBranches(unionQuery, allBranches);

        // Create namespace for each branch
        this.branchNamespaces = new ArrayList<>();
        for (int i = 0; i < allBranches.size(); i++) {
            TSelectSqlStatement branch = allBranches.get(i);
            SubqueryNamespace branchNs = new SubqueryNamespace(branch, "branch_" + i, nameMatcher);
            branchNamespaces.add(branchNs);
        }
    }

    /**
     * Iteratively flatten UNION branches into a list.
     * Handles nested UNION structures like: (A UNION B) UNION C
     * Uses explicit stack to avoid StackOverflow on deeply nested chains.
     */
    private void flattenUnionBranches(TSelectSqlStatement stmt, List<TSelectSqlStatement> branches) {
        if (stmt == null) {
            return;
        }

        Deque<TSelectSqlStatement> stack = new ArrayDeque<>();
        stack.push(stmt);

        while (!stack.isEmpty()) {
            TSelectSqlStatement current = stack.pop();
            if (current == null) {
                continue;
            }
            if (current.isCombinedQuery()) {
                // Push right first so left is processed first (LIFO)
                stack.push(current.getRightStmt());
                stack.push(current.getLeftStmt());
            } else {
                branches.add(current);
            }
        }
    }

    @Override
    public String getDisplayName() {
        return alias != null ? alias : "<union>";
    }

    @Override
    public TTable getFinalTable() {
        // For UNION, return the first branch's final table
        // This is used for single-table resolution
        if (!branchNamespaces.isEmpty()) {
            return branchNamespaces.get(0).getFinalTable();
        }
        return null;
    }

    @Override
    public List<TTable> getAllFinalTables() {
        // Return tables from ALL branches
        List<TTable> allTables = new ArrayList<>();
        for (SubqueryNamespace branchNs : branchNamespaces) {
            branchNs.validate();
            List<TTable> branchTables = branchNs.getAllFinalTables();
            allTables.addAll(branchTables);
        }
        return allTables;
    }

    @Override
    protected void doValidate() {
        // Extract columns from first branch's SELECT list (SQL standard)
        columnSources = new LinkedHashMap<>();

        if (allBranches.isEmpty()) {
            return;
        }

        // Validate all branch namespaces
        for (SubqueryNamespace branchNs : branchNamespaces) {
            branchNs.validate();
        }

        // Get columns from first branch (defines schema)
        TSelectSqlStatement firstBranch = allBranches.get(0);
        TResultColumnList selectList = firstBranch.getResultColumnList();
        if (selectList == null) {
            // (VALUES (1) UNION ALL VALUES (2)) AS r(x): the first branch is a VALUES row; its alias column list
            // names the cells (Mantis 4814, Codex review)
            if (aliasColumns != null && firstBranch.getValueClause() != null && firstBranch.getValueClause().getRows() != null
                    && !firstBranch.getValueClause().getRows().isEmpty()) {
                TResultColumnList firstRow = firstBranch.getValueClause().getRows().get(0);
                for (int i = 0; i < aliasColumns.size() && i < firstRow.size(); i++) {
                    if (aliasColumns.getObjectName(i) == null) continue;
                    String colName = aliasColumns.getObjectName(i).getColumnNameOnly();
                    columnSources.put(colName, new ColumnSource(this, colName, firstRow.getResultColumn(i), 1.0,
                            "union_alias_column_list"));
                }
            }
            return;
        }

        for (int i = 0; i < selectList.size(); i++) {
            TResultColumn resultCol = selectList.getResultColumn(i);
            String colName = getColumnName(resultCol);
            if (colName == null) {
                colName = "col_" + (i + 1);
            }
            // the first branch's own name: the branches are matched against it, not against the exposed name
            String firstBranchName = colName;
            // The alias column list renames the set operation's columns by position (Mantis 4814):
            // FROM (SELECT 0 AS a UNION ALL SELECT 1) AS r(rank) exposes rank.
            if (aliasColumns != null && i < aliasColumns.size() && aliasColumns.getObjectName(i) != null
                    && !SubqueryNamespace.isStarResultColumn(resultCol)) {
                colName = aliasColumns.getObjectName(i).getColumnNameOnly();
            }

            // For each column position, collect tables from ALL branches.
            // UNION/MINUS columns are matched by POSITION, not by name.
            // For each branch, check if the result column at that position is a simple column reference.
            // If so, include all tables from that branch as potential sources.
            List<TTable> columnTables = new ArrayList<>();

            for (int branchIdx = 0; branchIdx < allBranches.size(); branchIdx++) {
                TSelectSqlStatement branch = allBranches.get(branchIdx);
                TResultColumnList branchSelectList = branch.getResultColumnList();
                if (branchSelectList == null || i >= branchSelectList.size()) {
                    continue;
                }

                TResultColumn branchResultCol = branchSelectList.getResultColumn(i);
                if (branchResultCol == null || branchResultCol.getExpr() == null) {
                    continue;
                }

                // Check if this result column is a simple column reference (not NULL or expression)
                gudusoft.gsqlparser.nodes.TExpression expr = branchResultCol.getExpr();
                if (expr.getExpressionType() == gudusoft.gsqlparser.EExpressionType.simple_object_name_t) {
                    gudusoft.gsqlparser.nodes.TObjectName objName = expr.getObjectOperand();
                    if (objName != null) {
                        // Get the column name at this position in the branch
                        String branchColName = objName.getColumnNameOnly();

                        // For UNION data lineage, only include tables where the column name
                        // in this branch matches the column name from the first branch.
                        // This prevents incorrect associations like CDS_APP.bankcode when
                        // CDS_APP branch actually has c_mandant at that position.
                        boolean columnNameMatches = nameMatcher.matches(branchColName, firstBranchName);

                        // First try to get the source table from the column reference itself
                        TTable sourceTable = objName.getSourceTable();
                        if (sourceTable != null && !columnTables.contains(sourceTable) && columnNameMatches) {
                            // Phase 1 resolved this column - add if column name matches
                            columnTables.add(sourceTable);
                        } else if (sourceTable == null && columnNameMatches) {
                            // If sourceTable is not set (Phase 1 resolution hasn't happened),
                            // check if there's a qualified reference (e.g., t.col)
                            String tableQualifier = objName.getTableString();
                            if (tableQualifier != null && !tableQualifier.isEmpty() && branch.tables != null) {
                                // Qualified column - try to find the table by alias or name.
                                // Cleanup follow-up to S15: route through the per-dialect
                                // {@link gudusoft.gsqlparser.resolver2.matcher.INameMatcher},
                                // and when the matcher is a {@link VendorNameMatcher}, route
                                // through {@link ESQLDataObjectType#dotTable} explicitly. The
                                // 2-arg {@code matches(...)} default binds to {@code dotColumn}
                                // inside {@code VendorNameMatcher}, which produces wrong
                                // results for BigQuery / MySQL on table-qualifier compares
                                // (BigQuery: tables SENSITIVE / columns INSENSITIVE; MySQL:
                                // tables depend on lower_case_table_names, columns always
                                // INSENSITIVE). Same routing pattern as ListBasedScope and
                                // CTEScope below.
                                for (int ti = 0; ti < branch.tables.size(); ti++) {
                                    TTable t = branch.tables.getTable(ti);
                                    if (t != null) {
                                        String alias = t.getAliasName();
                                        String name = t.getTableName() != null ? t.getTableName().toString() : null;
                                        if ((alias != null && tableNameMatches(alias, tableQualifier)) ||
                                            (name != null && tableNameMatches(name, tableQualifier))) {
                                            if (!columnTables.contains(t)) {
                                                columnTables.add(t);
                                            }
                                            break;
                                        }
                                    }
                                }
                            } else if (branch.tables != null) {
                                // Unqualified column reference - add only the first non-subquery/join table
                                for (int ti = 0; ti < branch.tables.size(); ti++) {
                                    TTable t = branch.tables.getTable(ti);
                                    if (t != null && !columnTables.contains(t)) {
                                        // Skip subqueries and joins - they're not final tables
                                        if (t.getTableType() != gudusoft.gsqlparser.ETableSource.subquery &&
                                            t.getTableType() != gudusoft.gsqlparser.ETableSource.join) {
                                            columnTables.add(t);
                                            break;  // Only add the first table for unqualified columns
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
                // For non-column expressions (NULL, functions, etc.), don't add any tables
            }

            // Create column source with candidateTables from branches that have the column
            // Always pass the list (even if empty) so getAllFinalTables() knows we explicitly
            // determined the candidate tables rather than needing to delegate to namespace
            ColumnSource source = new ColumnSource(
                this,
                colName,
                resultCol,
                1.0,
                "union_column",
                null,  // overrideTable
                columnTables  // Pass empty list when no tables match
            );

            columnSources.put(colName, source);
        }
    }

    /**
     * Extract column name from TResultColumn.
     */
    private String getColumnName(TResultColumn resultCol) {
        // Check for alias
        if (resultCol.getAliasClause() != null &&
            resultCol.getAliasClause().getAliasName() != null) {
            return resultCol.getAliasClause().getAliasName().toString();
        }

        // Check for simple column reference
        if (resultCol.getExpr() != null) {
            gudusoft.gsqlparser.nodes.TExpression expr = resultCol.getExpr();
            if (expr.getExpressionType() == gudusoft.gsqlparser.EExpressionType.simple_object_name_t) {
                gudusoft.gsqlparser.nodes.TObjectName objName = expr.getObjectOperand();
                if (objName != null) {
                    return objName.getColumnNameOnly();
                }
            }
        }

        return null;
    }

    @Override
    public ColumnLevel hasColumn(String columnName) {
        ensureValidated();

        // Check in explicit columns from first branch (raw-keyed map; matcher-aware lookup
        // routes through INameMatcher / IdentifierService so per-dialect rules apply).
        if (containsColumnByMatcher(columnSources, columnName)) {
            return ColumnLevel.EXISTS;
        }

        // Check in inferred columns. The map is raw-keyed (= ColumnSource.
        // exposedName) — see addInferredColumn for why round-2 reverted from
        // normalized keys. The matcher-aware helper handles per-dialect case
        // rules including SQL Server COLLATION_BASED.
        if (containsColumnByMatcher(inferredColumns, columnName)) {
            return ColumnLevel.EXISTS;
        }

        // If any branch has SELECT *, unknown columns MAYBE exist
        if (hasStarColumn()) {
            return ColumnLevel.MAYBE;
        }

        return ColumnLevel.NOT_EXISTS;
    }

    /**
     * Binding-diagnostic view of the set-operation output schema.
     *
     * <p>SQL exposes set-operation columns using the first branch's output
     * names. Branch-local names from later operands are not visible to an
     * outer query. When any branch contains {@code SELECT *}, the output
     * shape is not authoritative for missing-output diagnostics.</p>
     */
    public ColumnLevel hasAuthoritativeOutputColumn(String columnName) {
        ensureValidated();

        if (columnName == null || columnName.isEmpty()) {
            return ColumnLevel.MAYBE;
        }

        if (containsColumnByMatcher(columnSources, columnName)) {
            return ColumnLevel.EXISTS;
        }

        if (containsColumnByMatcher(inferredColumns, columnName)) {
            return ColumnLevel.EXISTS;
        }

        if (hasStarColumn()) {
            return ColumnLevel.MAYBE;
        }

        if (columnSources != null && !columnSources.isEmpty()) {
            return ColumnLevel.NOT_EXISTS;
        }

        return ColumnLevel.MAYBE;
    }

    /**
     * Cleanup follow-up to S15: route table-qualifier compares through
     * {@link ESQLDataObjectType#dotTable} when the matcher is vendor-aware.
     * Without this, the 2-arg {@link INameMatcher#matches(String, String)}
     * defaults to {@code dotColumn} semantics inside {@link VendorNameMatcher},
     * producing wrong results on BigQuery / MySQL.
     */
    private boolean tableNameMatches(String storedName, String tableQualifier) {
        if (nameMatcher instanceof VendorNameMatcher) {
            return ((VendorNameMatcher) nameMatcher).matches(storedName, tableQualifier, ESQLDataObjectType.dotTable);
        }
        return nameMatcher.matches(storedName, tableQualifier);
    }

    @Override
    public ColumnSource resolveColumn(String columnName) {
        ensureValidated();

        // First check explicit columns from first branch
        ColumnSource source = super.resolveColumn(columnName);
        if (source != null) {
            return source;
        }

        // Then check inferred columns. Slice S1 + codex round 2: the map is
        // raw-keyed (= ColumnSource.exposedName), so the exact-match probe
        // is O(1) for the common case where the same identifier is queried
        // again. For case-only-different references, the matcher loop walks
        // values via getExposedName() so quote state is preserved.
        if (inferredColumns != null) {
            ColumnSource exact = inferredColumns.get(columnName);
            if (exact != null) {
                return exact;
            }
            for (ColumnSource entry : inferredColumns.values()) {
                String exposed = entry != null ? entry.getExposedName() : null;
                if (exposed != null && nameMatcher.matches(exposed, columnName)) {
                    return entry;
                }
            }
        }

        // If has star column, auto-infer this column
        if (hasStarColumn()) {
            boolean added = addInferredColumn(columnName, 0.8, "auto_inferred_from_outer_reference");
            if (added && inferredColumns != null) {
                ColumnSource inferredSource = inferredColumns.get(columnName);
                if (inferredSource != null) {
                    return inferredSource;
                }
            }
        }

        return null;
    }

    @Override
    public TSelectSqlStatement getSelectStatement() {
        return unionQuery;
    }

    @Override
    public boolean hasStarColumn() {
        // Returns true if ANY branch has SELECT *
        for (SubqueryNamespace branchNs : branchNamespaces) {
            if (branchNs.hasStarColumn()) {
                return true;
            }
        }
        return false;
    }

    @Override
    public boolean supportsDynamicInference() {
        return hasStarColumn();
    }

    /**
     * Slice S4 (plan §5.5): a set-operation's output schema is defined by the
     * first branch (SQL standard, see class javadoc). The schema is
     * authoritative once validated AND the first branch projects named
     * columns; if any branch contains an unresolved {@code SELECT *} we treat
     * the output as METADATA_UNAVAILABLE because position-based column
     * matching across branches has not yet completed. S11 will refine this
     * once arity-mismatch handling lands.
     */
    @Override
    public MetadataState getMetadataState() {
        ensureValidated();
        if (columnSources == null || columnSources.isEmpty() || hasStarColumn()) {
            return MetadataState.METADATA_UNAVAILABLE;
        }
        return MetadataState.FOUND;
    }

    @Override
    public boolean addInferredColumn(String columnName, double confidence, String evidence) {
        if (columnName == null || columnName.isEmpty()) {
            return false;
        }

        // Initialize maps if needed
        if (inferredColumns == null) {
            inferredColumns = new LinkedHashMap<>();
        }
        if (inferredColumnNames == null) {
            inferredColumnNames = new HashSet<>();
        }

        // Slice S1: dedupe checks must respect per-vendor identifier rules so
        // that case-only-different inputs do not leave the map with two entries
        // (one of which is then non-deterministically picked by the matcher
        // loops in hasColumn / resolveColumn).
        //
        // Codex round 2 caveat: the dedupe gate must use {@link
        // #containsColumnByMatcher} (which walks ColumnSource.exposedName for
        // value compares) and the storage key must be the original-cased
        // identifier — NOT the normalized form. Two matcher-distinct
        // identifiers can normalize to the same key (e.g. Postgres {@code
        // "mycol"} and unquoted {@code MYCOL} both normalize to {@code
        // mycol}), so a normalized-keyed map cannot represent both. Raw
        // (exposedName) keys avoid this collision.
        if (containsColumnByMatcher(columnSources, columnName)) {
            return false;
        }
        if (containsColumnByMatcher(inferredColumns, columnName)) {
            return false;
        }

        // Collect final tables from ALL branches that support dynamic inference (have SELECT *)
        // For data lineage, we need to track that columns could come from any UNION branch
        // The formatter will output all candidates for UNION columns when isCandidatesFromUnion is true
        java.util.List<TTable> candidateTables = new java.util.ArrayList<>();
        for (SubqueryNamespace branchNs : branchNamespaces) {
            // Only collect from branches that could have inferred columns
            if (!branchNs.supportsDynamicInference()) {
                continue;
            }
            branchNs.validate();
            java.util.List<TTable> branchTables = branchNs.getAllFinalTables();
            for (TTable table : branchTables) {
                if (table != null && !candidateTables.contains(table)) {
                    candidateTables.add(table);
                }
            }
        }

        // Create inferred column source for this union namespace WITH candidate tables
        ColumnSource source = new ColumnSource(
            this,
            columnName,
            null,
            confidence,
            evidence,
            null,  // overrideTable is null for UNION columns
            candidateTables.isEmpty() ? null : candidateTables
        );

        inferredColumns.put(columnName, source);
        inferredColumnNames.add(columnName);

        if (gudusoft.gsqlparser.TBaseType.DUMP_RESOLVER_LOG_TO_CONSOLE) {
            System.out.println("[UnionNamespace] Added '" + columnName + "' to " + alias +
                ", propagating to " + branchNamespaces.size() + " branches");
        }

        // CRITICAL: Propagate to branches that support dynamic inference (have SELECT *)
        // Only propagate to branches with star columns - branches with explicit column lists
        // either have the column explicitly or don't have it at all.
        for (SubqueryNamespace branchNs : branchNamespaces) {
            // Only propagate to branches that can accept inferred columns
            if (!branchNs.supportsDynamicInference()) {
                if (gudusoft.gsqlparser.TBaseType.DUMP_RESOLVER_LOG_TO_CONSOLE) {
                    System.out.println("[UnionNamespace] Skipping branch " + branchNs.getDisplayName() +
                        " (no star column)");
                }
                continue;
            }
            if (gudusoft.gsqlparser.TBaseType.DUMP_RESOLVER_LOG_TO_CONSOLE) {
                System.out.println("[UnionNamespace] Propagating '" + columnName + "' to branch " + branchNs.getDisplayName());
            }
            branchNs.addInferredColumn(columnName, confidence, evidence + "_union_propagate");
        }

        return true;
    }

    @Override
    public Set<String> getInferredColumns() {
        if (inferredColumnNames == null) {
            return Collections.emptySet();
        }
        return Collections.unmodifiableSet(inferredColumnNames);
    }

    /**
     * Get all branch namespaces.
     * Useful for external code that needs to iterate over branches.
     */
    public List<SubqueryNamespace> getBranchNamespaces() {
        return Collections.unmodifiableList(branchNamespaces);
    }

    /**
     * Get all branch SELECT statements.
     */
    public List<TSelectSqlStatement> getAllBranches() {
        return Collections.unmodifiableList(allBranches);
    }

    /**
     * Get the number of UNION branches.
     */
    public int getBranchCount() {
        return allBranches.size();
    }

    @Override
    public String toString() {
        int totalColumns = (columnSources != null ? columnSources.size() : 0) +
                          (inferredColumns != null ? inferredColumns.size() : 0);
        return "UnionNamespace(" + getDisplayName() +
               ", branches=" + allBranches.size() +
               ", columns=" + totalColumns +
               ", inferred=" + (inferredColumns != null ? inferredColumns.size() : 0) + ")";
    }
}