package gudusoft.gsqlparser.resolver2.namespace;

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.model.ColumnSource;
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;

    /** 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);
        }
    }

    /**
     * Recursively flatten UNION branches into a list.
     * Handles nested UNION structures like: (A UNION B) UNION C
     */
    private void flattenUnionBranches(TSelectSqlStatement stmt, List<TSelectSqlStatement> branches) {
        if (stmt == null) {
            return;
        }

        if (stmt.isCombinedQuery()) {
            // Recursively flatten left side
            flattenUnionBranches(stmt.getLeftStmt(), branches);
            // Recursively flatten right side
            flattenUnionBranches(stmt.getRightStmt(), branches);
        } else {
            // This is a leaf branch - add it
            branches.add(stmt);
        }
    }

    @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) {
            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);
            }

            // 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, colName);

                        // 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
                                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 && alias.equalsIgnoreCase(tableQualifier)) ||
                                            (name != null && name.equalsIgnoreCase(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
        for (String existingCol : columnSources.keySet()) {
            if (nameMatcher.matches(existingCol, columnName)) {
                return ColumnLevel.EXISTS;
            }
        }

        // Check in inferred columns
        if (inferredColumns != null && inferredColumns.containsKey(columnName)) {
            return ColumnLevel.EXISTS;
        }

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

        return ColumnLevel.NOT_EXISTS;
    }

    @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
        if (inferredColumns != null) {
            for (Map.Entry<String, ColumnSource> entry : inferredColumns.entrySet()) {
                if (nameMatcher.matches(entry.getKey(), columnName)) {
                    return entry.getValue();
                }
            }
        }

        // 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();
    }

    @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<>();
        }

        // Check if already exists in explicit columns
        if (columnSources != null && columnSources.containsKey(columnName)) {
            return false;
        }

        // Check if already inferred
        if (inferredColumns.containsKey(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) + ")";
    }
}