package gudusoft.gsqlparser.resolver2.enhancement;

import gudusoft.gsqlparser.TBaseType;
import gudusoft.gsqlparser.nodes.TCTE;
import gudusoft.gsqlparser.nodes.TObjectName;
import gudusoft.gsqlparser.nodes.TTable;
import gudusoft.gsqlparser.resolver2.ResolutionStatus;
import gudusoft.gsqlparser.resolver2.ScopeBuildResult;
import gudusoft.gsqlparser.resolver2.inference.EvidenceType;
import gudusoft.gsqlparser.resolver2.model.ResolutionEvidence;
import gudusoft.gsqlparser.resolver2.model.ResolutionResult;
import gudusoft.gsqlparser.resolver2.namespace.CTENamespace;
import gudusoft.gsqlparser.resolver2.namespace.INamespace;
import gudusoft.gsqlparser.resolver2.namespace.SubqueryNamespace;
import gudusoft.gsqlparser.resolver2.namespace.UnionNamespace;
import gudusoft.gsqlparser.resolver2.scope.IScope;
import gudusoft.gsqlparser.resolver2.scope.SelectScope;
import gudusoft.gsqlparser.stmt.TSelectSqlStatement;

import java.util.*;

/**
 * Handles explicit namespace enhancement between resolution passes.
 *
 * The enhancement process:
 * 1. During resolution, collect column references that target namespaces with star columns
 * 2. Between passes, explicitly add these columns to the namespaces
 * 3. Track which columns were added in which pass for debugging
 *
 * This makes the column inference process explicit and traceable, rather than
 * happening implicitly during resolution.
 */
public class NamespaceEnhancer {

    /** Collected column references pending enhancement */
    private final List<CollectedColumnRef> pendingRefs;

    /** History of all enhancement results */
    private final List<EnhancementResult> enhancementHistory;

    /** Current pass number */
    private int currentPass;

    /** All namespaces that have star columns (cached) */
    private Set<INamespace> starNamespaces;

    /** Map of TTable -> INamespace for subqueries in FROM clauses (for propagation) */
    private Map<TTable, INamespace> tableToNamespaceMap;

    public NamespaceEnhancer() {
        this(false);
    }

    public NamespaceEnhancer(boolean debug) {
        // debug parameter kept for API compatibility but no longer used
        this.pendingRefs = new ArrayList<>();
        this.enhancementHistory = new ArrayList<>();
        this.currentPass = 0;
        this.starNamespaces = null;
    }

    /**
     * Initialize the enhancer with scope build result.
     * Caches namespaces that have star columns and builds TTable -> INamespace mapping.
     */
    public void initialize(ScopeBuildResult scopeBuildResult) {
        this.starNamespaces = new HashSet<>();
        this.tableToNamespaceMap = new HashMap<>();

        // Collect from all SelectScopes in the statement scope map
        // This is more reliable than traversing from GlobalScope since the scope tree
        // may have different structures for different SQL constructs
        for (IScope scope : scopeBuildResult.getStatementScopeMap().values()) {
            collectStarNamespacesInternal(scope, new HashSet<>());
        }
    }

    /**
     * Recursively collect all namespaces that have star columns.
     * Uses a visited set to avoid infinite recursion.
     */
    private void collectStarNamespaces(IScope scope) {
        collectStarNamespacesInternal(scope, new HashSet<>());
    }

    private void collectStarNamespacesInternal(IScope scope, Set<IScope> visited) {
        if (scope == null || visited.contains(scope)) return;
        visited.add(scope);

        // Check if this scope has namespaces with star columns
        if (scope instanceof SelectScope) {
            SelectScope selectScope = (SelectScope) scope;
            IScope fromScope = selectScope.getFromScope();
            if (fromScope != null && !visited.contains(fromScope)) {
                visited.add(fromScope);
                // Check children (ScopeChild contains namespaces)
                for (gudusoft.gsqlparser.resolver2.model.ScopeChild child : fromScope.getChildren()) {
                    INamespace childNs = child.getNamespace();
                    if (childNs != null) {
                        collectFromNamespace(childNs, new HashSet<>());
                    }
                }
            }
        }

        // Check children of this scope (ScopeChild contains namespaces)
        for (gudusoft.gsqlparser.resolver2.model.ScopeChild child : scope.getChildren()) {
            INamespace childNs = child.getNamespace();
            if (childNs != null) {
                collectFromNamespace(childNs, new HashSet<>());
            }
        }
    }

    /**
     * Recursively collect namespaces with star columns from a namespace and its nested namespaces.
     */
    private void collectFromNamespace(INamespace namespace, Set<INamespace> visitedNs) {
        if (namespace == null || visitedNs.contains(namespace)) return;
        visitedNs.add(namespace);

        // Add this namespace if it has star columns
        if (namespace.hasStarColumn()) {
            starNamespaces.add(namespace);
        }

        // If this is a CTENamespace, also check its internal UnionNamespace
        if (namespace instanceof CTENamespace) {
            CTENamespace cteNs = (CTENamespace) namespace;
            UnionNamespace unionNs = cteNs.getUnionNamespace();
            if (unionNs != null) {
                collectFromNamespace(unionNs, visitedNs);
            }
        }

        // If this is a UnionNamespace, also add its branch namespaces
        if (namespace instanceof UnionNamespace) {
            UnionNamespace unionNs = (UnionNamespace) namespace;
            for (SubqueryNamespace branchNs : unionNs.getBranchNamespaces()) {
                collectFromNamespace(branchNs, visitedNs);
            }
        }
    }

    /**
     * Start a new resolution pass.
     * Clears pending references from previous pass.
     */
    public void startPass(int passNumber) {
        this.currentPass = passNumber;
        this.pendingRefs.clear();
    }

    /**
     * Collect a column reference that could not be resolved.
     * Called during resolution when a column targets a star namespace.
     *
     * @param columnName The column name
     * @param targetNamespace The namespace the column should belong to
     * @param sourceReference The AST node
     * @param confidence Confidence score
     * @param evidence Reason for this inference
     *
     * @deprecated Use {@link #collectColumnRef(String, INamespace, TObjectName, ResolutionEvidence)} instead
     */
    public void collectColumnRef(String columnName,
                                  INamespace targetNamespace,
                                  TObjectName sourceReference,
                                  double confidence,
                                  String evidence) {
        if (columnName == null || columnName.isEmpty()) {
            return;
        }

        // Only collect for star namespaces or namespaces that support dynamic inference
        if (targetNamespace == null ||
            (!targetNamespace.hasStarColumn() && !targetNamespace.supportsDynamicInference())) {
            return;
        }

        CollectedColumnRef ref = new CollectedColumnRef(
            columnName,
            targetNamespace,
            sourceReference,
            currentPass,
            confidence,
            evidence
        );

        // Avoid duplicates
        if (!pendingRefs.contains(ref)) {
            pendingRefs.add(ref);
        }
    }

    /**
     * Collect a column reference with structured evidence.
     * Called during resolution when a column targets a star namespace.
     *
     * @param columnName The column name
     * @param targetNamespace The namespace the column should belong to
     * @param sourceReference The AST node
     * @param evidenceDetail Structured evidence detail
     */
    public void collectColumnRef(String columnName,
                                  INamespace targetNamespace,
                                  TObjectName sourceReference,
                                  ResolutionEvidence evidenceDetail) {
        if (columnName == null || columnName.isEmpty()) {
            return;
        }

        // Only collect for star namespaces or namespaces that support dynamic inference
        if (targetNamespace == null ||
            (!targetNamespace.hasStarColumn() && !targetNamespace.supportsDynamicInference())) {
            return;
        }

        CollectedColumnRef ref = new CollectedColumnRef(
            columnName,
            targetNamespace,
            sourceReference,
            currentPass,
            evidenceDetail
        );

        // Avoid duplicates
        if (!pendingRefs.contains(ref)) {
            pendingRefs.add(ref);
        }
    }

    /**
     * Collect a column reference based on resolution result.
     * Automatically determines if collection is needed based on resolution status.
     */
    public void collectFromResolution(TObjectName column,
                                       ResolutionResult result,
                                       INamespace candidateNamespace) {
        if (result == null || column == null) return;

        String columnName = column.getColumnNameOnly();
        if (columnName == null || columnName.isEmpty()) return;

        // Collect if:
        // 1. Column was not found, but we have a candidate namespace with star
        // 2. Column was resolved with low confidence through star inference
        if (result.getStatus() == ResolutionStatus.NOT_FOUND) {
            if (candidateNamespace != null && candidateNamespace.hasStarColumn()) {
                ResolutionEvidence evidence = ResolutionEvidence.fromStarInference(
                    columnName,
                    candidateNamespace.getDisplayName(),
                    false, // unqualified reference
                    column
                );
                collectColumnRef(columnName, candidateNamespace, column, evidence);
            }
        } else if (result.getStatus() == ResolutionStatus.EXACT_MATCH) {
            // If resolved through star namespace, already tracked - no action needed
        }
    }

    /**
     * Perform namespace enhancement using collected column references.
     * This is called BETWEEN resolution passes.
     *
     * @return EnhancementResult with details of what was added
     */
    public EnhancementResult enhance() {
        EnhancementResult result = new EnhancementResult(currentPass);

        for (CollectedColumnRef ref : pendingRefs) {
            INamespace namespace = ref.getTargetNamespace();
            String columnName = ref.getColumnName();

            // Try to add the column using the collected evidence
            // Use evidenceDetail if available, otherwise fall back to legacy parameters
            boolean added;
            if (ref.getEvidenceDetail() != null) {
                added = namespace.addInferredColumn(
                    columnName,
                    ref.getEvidenceDetail().getWeight(),
                    ref.getEvidenceDetail().toLegacyEvidence()
                );
            } else {
                added = namespace.addInferredColumn(
                    columnName,
                    ref.getConfidence(),
                    ref.getEvidence()
                );
            }

            if (added) {
                result.recordAdded(namespace, columnName, ref);
            } else {
                result.recordSkipped(namespace, columnName, ref);
            }
        }

        // Propagate inferred columns through CTE chains
        // This handles cases like: WITH cte2 AS (SELECT * FROM cte1) where cte1 also has SELECT *
        // Columns pushed to cte2 should also be pushed to cte1
        propagateThroughCTEChains(result);

        // Store in history
        enhancementHistory.add(result);

        // Clear pending refs
        pendingRefs.clear();

        return result;
    }

    /**
     * Propagate inferred columns through CTE chains.
     *
     * When a CTE references another CTE with star columns (e.g., WITH cte2 AS (SELECT * FROM cte1)),
     * any columns inferred in cte2 should also be pushed to cte1.
     *
     * Also propagates from SubqueryNamespace to referenced CTEs:
     * When a subquery like (SELECT * FROM cte1) has inferred columns,
     * those columns should propagate to cte1.
     *
     * This method uses the actual CTENamespace instances from starNamespaces (not new instances),
     * ensuring that changes persist to the scope tree.
     */
    private void propagateThroughCTEChains(EnhancementResult result) {
        // Build a map of TCTE -> CTENamespace for the actual instances from the scope tree
        Map<TCTE, CTENamespace> cteMap = buildCTEMap();

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

        // Keep propagating until no more changes
        boolean changed;
        int iterations = 0;
        int maxIterations = 10; // Safety limit

        do {
            changed = false;
            iterations++;

            for (INamespace ns : starNamespaces) {
                if (ns instanceof CTENamespace) {
                    CTENamespace cteNs = (CTENamespace) ns;
                    changed |= propagateToReferencedCTEs(cteNs, cteMap, result);
                } else if (ns instanceof UnionNamespace) {
                    // Propagate from UnionNamespace to CTEs/subqueries referenced in its branches
                    UnionNamespace unionNs = (UnionNamespace) ns;
                    changed |= propagateFromUnionToCTEs(unionNs, cteMap, result);
                } else if (ns instanceof SubqueryNamespace) {
                    // Also propagate from SubqueryNamespace to CTEs it references
                    SubqueryNamespace subNs = (SubqueryNamespace) ns;
                    changed |= propagateFromSubqueryToCTEs(subNs, cteMap, result);
                }
            }
        } while (changed && iterations < maxIterations);
    }

    /**
     * Build a map of TCTE -> CTENamespace from the actual instances in starNamespaces.
     */
    private Map<TCTE, CTENamespace> buildCTEMap() {
        Map<TCTE, CTENamespace> map = new HashMap<>();
        for (INamespace ns : starNamespaces) {
            if (ns instanceof CTENamespace) {
                CTENamespace cteNs = (CTENamespace) ns;
                if (cteNs.getCTE() != null) {
                    map.put(cteNs.getCTE(), cteNs);
                }
            }
        }
        return map;
    }

    /**
     * Propagate inferred columns from a CTENamespace to CTEs and subqueries it references.
     *
     * This handles:
     * 1. CTE references (WITH cte AS (SELECT * FROM other_cte))
     * 2. Subquery references (WITH cte AS (SELECT * FROM (SELECT * FROM t)))
     * 3. UNION subqueries (WITH cte AS (SELECT * FROM (SELECT ... UNION ALL SELECT ...)))
     *
     * @return true if any columns were added
     */
    private boolean propagateToReferencedCTEs(CTENamespace source,
                                               Map<TCTE, CTENamespace> cteMap,
                                               EnhancementResult result) {
        boolean changed = false;

        Set<String> inferredCols = source.getInferredColumns();
        if (inferredCols.isEmpty()) {
            return false;
        }

        TSelectSqlStatement stmt = source.getSelectStatement();
        if (stmt == null || stmt.tables == null) {
            return false;
        }

        for (int i = 0; i < stmt.tables.size(); i++) {
            TTable table = stmt.tables.getTable(i);
            if (table == null) continue;

            // Case 1: CTE reference - look up the actual CTENamespace instance
            if (table.isCTEName() && table.getCTE() != null) {
                CTENamespace targetNs = cteMap.get(table.getCTE());
                if (targetNs != null && targetNs.hasStarColumn()) {
                    changed |= propagateColumnsToNamespace(source, targetNs, inferredCols, result,
                        "CTE chain propagation");
                }
            }
            // Case 2: Subquery table - find the namespace in starNamespaces
            else if (table.getSubquery() != null) {
                TSelectSqlStatement subquery = table.getSubquery();
                INamespace targetNs = findNamespaceForStatement(subquery);
                if (targetNs != null && targetNs.hasStarColumn()) {
                    changed |= propagateColumnsToNamespace(source, targetNs, inferredCols, result,
                        "CTE to subquery propagation");
                }
            }
        }

        return changed;
    }

    /**
     * Find the namespace in starNamespaces that corresponds to a given TSelectSqlStatement.
     * This is used to look up the actual namespace instance for a subquery.
     *
     * @param stmt The TSelectSqlStatement to find
     * @return The corresponding namespace, or null if not found
     */
    private INamespace findNamespaceForStatement(TSelectSqlStatement stmt) {
        if (stmt == null) return null;

        for (INamespace ns : starNamespaces) {
            if (ns instanceof SubqueryNamespace) {
                SubqueryNamespace subNs = (SubqueryNamespace) ns;
                if (subNs.getSelectStatement() == stmt) {
                    return subNs;
                }
            } else if (ns instanceof UnionNamespace) {
                UnionNamespace unionNs = (UnionNamespace) ns;
                if (unionNs.getSelectStatement() == stmt) {
                    return unionNs;
                }
            }
        }
        return null;
    }

    /**
     * Propagate a set of columns from source namespace to target namespace.
     *
     * @param source The source namespace with inferred columns
     * @param target The target namespace to propagate to
     * @param columns The column names to propagate
     * @param result EnhancementResult to record changes
     * @param evidencePrefix Prefix for evidence string
     * @return true if any columns were added
     */
    private boolean propagateColumnsToNamespace(INamespace source,
                                                 INamespace target,
                                                 Set<String> columns,
                                                 EnhancementResult result,
                                                 String evidencePrefix) {
        boolean changed = false;
        for (String colName : columns) {
            ResolutionEvidence evidence = new ResolutionEvidence(
                EvidenceType.INFERRED_FROM_USAGE,
                0.8,
                evidencePrefix + " from " + source.getDisplayName() + " to " + target.getDisplayName(),
                (gudusoft.gsqlparser.nodes.TParseTreeNode) null
            );
            boolean added = target.addInferredColumn(colName, evidence.getWeight(), evidence.toLegacyEvidence());
            if (added) {
                changed = true;
                result.recordAdded(target, colName, null);
            }
        }
        return changed;
    }

    /**
     * Propagate inferred columns from a SubqueryNamespace to CTEs and subqueries it references.
     *
     * When a subquery has SELECT * FROM cte and the subquery has inferred columns,
     * those columns should propagate to the CTE so the CTE's star columns can be expanded.
     *
     * @param source The SubqueryNamespace with inferred columns
     * @param cteMap Map of TCTE to CTENamespace
     * @param result EnhancementResult to record changes
     * @return true if any columns were added
     */
    private boolean propagateFromSubqueryToCTEs(SubqueryNamespace source,
                                                 Map<TCTE, CTENamespace> cteMap,
                                                 EnhancementResult result) {
        boolean changed = false;

        Set<String> inferredCols = source.getInferredColumns();
        if (inferredCols.isEmpty()) {
            return false;
        }

        TSelectSqlStatement stmt = source.getSelectStatement();
        if (stmt == null || stmt.tables == null) {
            return false;
        }

        for (int i = 0; i < stmt.tables.size(); i++) {
            TTable table = stmt.tables.getTable(i);
            if (table == null) continue;

            // Case 1: CTE reference - look up the actual CTENamespace instance
            if (table.isCTEName() && table.getCTE() != null) {
                CTENamespace targetNs = cteMap.get(table.getCTE());
                if (targetNs != null && targetNs.hasStarColumn()) {
                    changed |= propagateColumnsToNamespace(source, targetNs, inferredCols, result,
                        "Subquery to CTE propagation");
                }
            }
            // Case 2: Subquery table - find the namespace in starNamespaces
            else if (table.getSubquery() != null) {
                TSelectSqlStatement subquery = table.getSubquery();
                INamespace targetNs = findNamespaceForStatement(subquery);
                if (targetNs != null && targetNs.hasStarColumn()) {
                    changed |= propagateColumnsToNamespace(source, targetNs, inferredCols, result,
                        "Subquery to subquery propagation");
                }
            }
        }

        return changed;
    }

    /**
     * Propagate inferred columns from a UnionNamespace to CTEs and subqueries referenced in its branches.
     *
     * UNION branches can reference CTEs or subqueries. When columns are inferred in the UNION,
     * they should propagate to all branch sources.
     *
     * @param source The UnionNamespace with inferred columns
     * @param cteMap Map of TCTE to CTENamespace
     * @param result EnhancementResult to record changes
     * @return true if any columns were added
     */
    private boolean propagateFromUnionToCTEs(UnionNamespace source,
                                              Map<TCTE, CTENamespace> cteMap,
                                              EnhancementResult result) {
        boolean changed = false;

        Set<String> inferredCols = source.getInferredColumns();
        if (inferredCols.isEmpty()) {
            return false;
        }

        // Propagate through each branch of the UNION
        for (SubqueryNamespace branchNs : source.getBranchNamespaces()) {
            TSelectSqlStatement branchStmt = branchNs.getSelectStatement();
            if (branchStmt == null || branchStmt.tables == null) continue;

            for (int i = 0; i < branchStmt.tables.size(); i++) {
                TTable table = branchStmt.tables.getTable(i);
                if (table == null) continue;

                // Case 1: CTE reference
                if (table.isCTEName() && table.getCTE() != null) {
                    CTENamespace targetNs = cteMap.get(table.getCTE());
                    if (targetNs != null && targetNs.hasStarColumn()) {
                        changed |= propagateColumnsToNamespace(source, targetNs, inferredCols, result,
                            "UNION branch to CTE propagation");
                    }
                }
                // Case 2: Subquery table
                else if (table.getSubquery() != null) {
                    TSelectSqlStatement subquery = table.getSubquery();
                    INamespace targetNs = findNamespaceForStatement(subquery);
                    if (targetNs != null && targetNs.hasStarColumn()) {
                        changed |= propagateColumnsToNamespace(source, targetNs, inferredCols, result,
                            "UNION branch to subquery propagation");
                    }
                }
            }
        }

        return changed;
    }

    /**
     * Get all enhancement results history
     */
    public List<EnhancementResult> getEnhancementHistory() {
        return Collections.unmodifiableList(enhancementHistory);
    }

    /**
     * Get total columns added across all passes
     */
    public int getTotalColumnsAdded() {
        return enhancementHistory.stream()
            .mapToInt(EnhancementResult::getTotalAdded)
            .sum();
    }

    /**
     * Get number of pending references (not yet enhanced)
     */
    public int getPendingCount() {
        return pendingRefs.size();
    }

    /**
     * Check if a namespace is a star namespace
     */
    public boolean isStarNamespace(INamespace namespace) {
        return starNamespaces != null && starNamespaces.contains(namespace);
    }

    /**
     * Get all star namespaces
     */
    public Set<INamespace> getStarNamespaces() {
        return starNamespaces != null
            ? Collections.unmodifiableSet(starNamespaces)
            : Collections.emptySet();
    }

    /**
     * Generate a summary report of all enhancements
     */
    public String generateReport() {
        StringBuilder sb = new StringBuilder();
        sb.append("=== Namespace Enhancement Summary ===\n");
        sb.append("Total passes: ").append(enhancementHistory.size()).append("\n");
        sb.append("Total columns added: ").append(getTotalColumnsAdded()).append("\n");
        sb.append("Star namespaces: ").append(starNamespaces != null ? starNamespaces.size() : 0).append("\n\n");

        for (EnhancementResult result : enhancementHistory) {
            sb.append(result.toDetailedReport());
        }

        return sb.toString();
    }
}