package gudusoft.gsqlparser.resolver2.namespace;

import gudusoft.gsqlparser.nodes.TObjectName;
import gudusoft.gsqlparser.nodes.TObjectNameList;
import gudusoft.gsqlparser.nodes.TResultColumn;
import gudusoft.gsqlparser.nodes.TResultColumnList;
import gudusoft.gsqlparser.nodes.TTable;
import gudusoft.gsqlparser.nodes.TPivotClause;
import gudusoft.gsqlparser.nodes.TPivotInClause;
import gudusoft.gsqlparser.nodes.TUnpivotInClause;
import gudusoft.gsqlparser.nodes.TUnpivotInClauseItem;
import gudusoft.gsqlparser.nodes.TExpression;
import gudusoft.gsqlparser.nodes.TExpressionList;
import gudusoft.gsqlparser.nodes.TFunctionCall;
import gudusoft.gsqlparser.nodes.TConstant;
import gudusoft.gsqlparser.EExpressionType;
import gudusoft.gsqlparser.resolver2.ColumnLevel;
import gudusoft.gsqlparser.resolver2.matcher.INameMatcher;
import gudusoft.gsqlparser.resolver2.model.ColumnSource;

import java.util.*;

/**
 * Namespace representing a PIVOT table.
 *
 * <p>A PIVOT transforms rows into columns using an aggregate function:
 * <pre>
 * SELECT ...
 * FROM source_table
 * PIVOT (aggregate_function(value_column) FOR pivot_column IN (val1, val2, ...)) AS alias
 * </pre>
 *
 * <p>The PIVOT produces:
 * <ul>
 *   <li>Columns from the IN clause (val1, val2, ...) - these are the pivoted columns</li>
 *   <li>Pass-through columns from the source table (all columns NOT used in aggregate or FOR clause)</li>
 * </ul>
 *
 * <p>For column resolution purposes:
 * <ul>
 *   <li>Pivot columns (from IN clause) resolve to the PivotNamespace</li>
 *   <li>Other columns may come from the source table (pass-through columns)</li>
 * </ul>
 */
public class PivotNamespace extends AbstractNamespace {

    /** The pivot table this namespace represents */
    private final TTable pivotTable;

    /** The source table that the PIVOT operates on */
    private final TTable sourceTable;

    /** The PIVOT clause */
    private final TPivotClause pivotClause;

    /** Alias for this pivot table (from AS clause) */
    private final String alias;

    /** Namespace for the source table (for pass-through column resolution) */
    private INamespace sourceNamespace;

    /** Columns defined by the PIVOT IN clause */
    private final Set<String> pivotColumns = new LinkedHashSet<>();

    /** Columns consumed by UNPIVOT IN (...) list (these do NOT exist as output columns) */
    private final Set<String> unpivotInColumns = new LinkedHashSet<>();

    /** Display name format for PIVOT table */
    public static final String PIVOT_TABLE_SUFFIX = "(piviot_table)";

    /** Display name format for UNPIVOT table */
    public static final String UNPIVOT_TABLE_SUFFIX = "(unpivot_table)";

    /** Whether this is an UNPIVOT clause */
    private boolean isUnpivot = false;

    public PivotNamespace(TTable pivotTable, TPivotClause pivotClause,
                          TTable sourceTable, String alias, INameMatcher nameMatcher) {
        super(pivotTable, nameMatcher);
        this.pivotTable = pivotTable;
        this.pivotClause = pivotClause;
        this.sourceTable = sourceTable;
        this.alias = alias;
        // Check if this is an UNPIVOT clause
        if (pivotClause != null) {
            this.isUnpivot = (pivotClause.getType() == TPivotClause.unpivot);
        }
    }

    /**
     * Check if this is an UNPIVOT namespace.
     */
    public boolean isUnpivot() {
        return isUnpivot;
    }

    /**
     * Set the namespace representing the source table.
     * This is used for resolving pass-through columns.
     */
    public void setSourceNamespace(INamespace sourceNamespace) {
        this.sourceNamespace = sourceNamespace;
    }

    /**
     * Get the namespace representing the source table.
     */
    public INamespace getSourceNamespace() {
        return sourceNamespace;
    }

    /**
     * Get the pivot table.
     */
    public TTable getPivotTable() {
        return pivotTable;
    }

    /**
     * Get the underlying table that PIVOT operates on.
     * This is the table whose columns are being pivoted.
     */
    public TTable getUnderlyingSourceTable() {
        return sourceTable;
    }

    /**
     * {@inheritDoc}
     * For PivotNamespace, returns the PIVOT table (the result of the PIVOT operation).
     * This is the immediate source table for columns resolved through this PIVOT.
     */
    @Override
    public TTable getSourceTable() {
        return pivotTable;
    }

    /**
     * Get the PIVOT clause.
     */
    public TPivotClause getPivotClause() {
        return pivotClause;
    }

    /**
     * Get the set of pivot column names (from IN clause).
     */
    public Set<String> getPivotColumns() {
        return Collections.unmodifiableSet(pivotColumns);
    }

    /**
     * Check if a column is a pivot column (from IN clause).
     */
    public boolean isPivotColumn(String columnName) {
        if (columnName == null) return false;
        for (String pivotCol : pivotColumns) {
            if (nameMatcher.matches(pivotCol, columnName)) {
                return true;
            }
        }
        return false;
    }

    @Override
    public String getDisplayName() {
        String suffix = isUnpivot ? UNPIVOT_TABLE_SUFFIX : PIVOT_TABLE_SUFFIX;
        if (alias != null && !alias.isEmpty()) {
            return alias + suffix;
        }
        return (isUnpivot ? "unpivot_alias" : "pivot_alias") + suffix;
    }

    @Override
    public TTable getFinalTable() {
        // Return the pivot table itself so that pivot columns are properly attributed
        // to the pivot table (e.g., "(pivot-table:p(piviot_table)).[1]")
        // For lineage tracing to source columns, use getAllFinalTables()
        return pivotTable;
    }

    @Override
    public List<TTable> getAllFinalTables() {
        // Return source table for lineage tracing
        if (sourceTable != null) {
            List<TTable> tables = new ArrayList<>();
            tables.add(sourceTable);
            return tables;
        }
        return Collections.emptyList();
    }

    @Override
    protected void doValidate() {
        columnSources = new LinkedHashMap<>();
        pivotColumns.clear();
        unpivotInColumns.clear();

        if (pivotClause == null) {
            return;
        }

        if (isUnpivot) {
            // UNPIVOT case: Add generated columns (value column and FOR column)
            // UNPIVOT (yearly_total FOR order_mode IN (store AS 'direct', internet AS 'online'))
            // - yearly_total is the value column (new column containing values)
            // - order_mode is the FOR column (new column containing labels)

            // Add value columns (e.g., yearly_total)
            // Note: For single value column UNPIVOT (like Oracle), use getValueColumn() (deprecated singular)
            // For multi-value column UNPIVOT (like SQL Server), use getValueColumnList()
            TObjectNameList valueColumns = pivotClause.getValueColumnList();
            if (valueColumns != null && valueColumns.size() > 0) {
                for (int i = 0; i < valueColumns.size(); i++) {
                    TObjectName valueCol = valueColumns.getObjectName(i);
                    if (valueCol != null) {
                        String columnName = valueCol.getColumnNameOnly();
                        if (columnName != null && !columnName.isEmpty() && !columnSources.containsKey(columnName)) {
                            addPivotColumn(columnName, "unpivot_value_column");
                        }
                    }
                }
            } else {
                // Fallback to deprecated singular method for Oracle compatibility
                @SuppressWarnings("deprecation")
                TObjectName valueCol = pivotClause.getValueColumn();
                if (valueCol != null) {
                    String columnName = valueCol.getColumnNameOnly();
                    if (columnName != null && !columnName.isEmpty() && !columnSources.containsKey(columnName)) {
                        addPivotColumn(columnName, "unpivot_value_column");
                    }
                }
            }

            // Add FOR columns (e.g., order_mode)
            // Note: For single FOR column UNPIVOT (like Oracle), use getPivotColumn() (deprecated singular)
            // For multi-FOR column UNPIVOT, use getPivotColumnList()
            TObjectNameList pivotColumnList = pivotClause.getPivotColumnList();
            if (pivotColumnList != null && pivotColumnList.size() > 0) {
                for (int i = 0; i < pivotColumnList.size(); i++) {
                    TObjectName forCol = pivotColumnList.getObjectName(i);
                    if (forCol != null) {
                        String columnName = forCol.getColumnNameOnly();
                        if (columnName != null && !columnName.isEmpty() && !columnSources.containsKey(columnName)) {
                            addPivotColumn(columnName, "unpivot_for_column");
                        }
                    }
                }
            } else {
                // Fallback to deprecated singular method for Oracle compatibility
                @SuppressWarnings("deprecation")
                TObjectName forCol = pivotClause.getPivotColumn();
                if (forCol != null) {
                    String columnName = forCol.getColumnNameOnly();
                    if (columnName != null && !columnName.isEmpty() && !columnSources.containsKey(columnName)) {
                        addPivotColumn(columnName, "unpivot_for_column");
                    }
                }
            }

            // UNPIVOT case: Collect consumed IN(...) source columns so they are NOT treated as pass-through
            // Example: UNPIVOT (col3 FOR col4 IN (p.col2, p.col3))
            // - p.col2, p.col3 are consumed source columns and should not be visible via unpivot alias
            TUnpivotInClause unpivotInClause = pivotClause.getUnpivotInClause();
            if (unpivotInClause != null && unpivotInClause.getItems() != null) {
                for (int i = 0; i < unpivotInClause.getItems().size(); i++) {
                    TUnpivotInClauseItem item = unpivotInClause.getItems().getElement(i);
                    if (item == null) continue;

                    if (item.getColumn() != null) {
                        String name = item.getColumn().getColumnNameOnly();
                        if (name != null && !name.isEmpty()) {
                            unpivotInColumns.add(name);
                        }
                    }
                    if (item.getColumnList() != null) {
                        TObjectNameList cols = item.getColumnList();
                        for (int j = 0; j < cols.size(); j++) {
                            TObjectName col = cols.getObjectName(j);
                            if (col != null) {
                                String name = col.getColumnNameOnly();
                                if (name != null && !name.isEmpty()) {
                                    unpivotInColumns.add(name);
                                }
                            }
                        }
                    }
                }
            }
        } else {
            // PIVOT case: Get output columns
            // If alias clause has a column list (e.g., AS p (col1, col2, col3)), use those
            // as they REPLACE the original pivot column names. Otherwise use IN clause columns.
            boolean hasAliasColumnList = pivotClause.getAliasClause() != null &&
                pivotClause.getAliasClause().getColumns() != null &&
                pivotClause.getAliasClause().getColumns().size() > 0;

            if (hasAliasColumnList) {
                // Alias column list replaces the default pivot column names
                // e.g., PIVOT(...) AS p (empid_renamed, Q1, Q2, Q3, Q4)
                // The alias columns provide the output column names
                for (TObjectName column : pivotClause.getAliasClause().getColumns()) {
                    String columnName = stripDelimiters(column.toString());
                    if (columnName != null && !columnName.isEmpty() && !columnSources.containsKey(columnName)) {
                        addPivotColumn(columnName, "pivot_alias_clause");
                    }
                }
            } else {
                // No alias column list - use IN clause columns as pivot column names
                TPivotInClause inClause = pivotClause.getPivotInClause();
                if (inClause != null) {
                    // Case 1: IN clause has items (e.g., IN ([1], [2], [3]) or IN ([Sammich], [Pickle]))
                    if (inClause.getItems() != null) {
                        TResultColumnList items = inClause.getItems();
                        for (int i = 0; i < items.size(); i++) {
                            TResultColumn resultColumn = items.getResultColumn(i);
                            String columnName = extractPivotColumnName(resultColumn);
                            if (columnName != null && !columnName.isEmpty() && !columnSources.containsKey(columnName)) {
                                addPivotColumn(columnName, "pivot_in_clause");
                            }
                        }
                    }

                    // Case 2: IN clause has a subquery (e.g., IN (SELECT DISTINCT col FROM table))
                    if (inClause.getSubQuery() != null) {
                        TResultColumnList subqueryColumns = inClause.getSubQuery().getResultColumnList();
                        if (subqueryColumns != null) {
                            for (int i = 0; i < subqueryColumns.size(); i++) {
                                TResultColumn resultColumn = subqueryColumns.getResultColumn(i);
                                String columnName = stripDelimiters(resultColumn.getDisplayName());
                                if (columnName != null && !columnName.isEmpty() && !columnSources.containsKey(columnName)) {
                                    addPivotColumn(columnName, "pivot_in_subquery");
                                }
                            }
                        }
                    }

                    // Case 3: IN clause has value list (e.g., IN ((val1, val2), (val3, val4)))
                    if (inClause.getValueList() != null) {
                        for (TExpressionList exprList : inClause.getValueList()) {
                            if (exprList != null) {
                                for (int i = 0; i < exprList.size(); i++) {
                                    TExpression expr = exprList.getExpression(i);
                                    String columnName = extractColumnNameFromExpression(expr);
                                    if (columnName != null && !columnName.isEmpty() && !columnSources.containsKey(columnName)) {
                                        addPivotColumn(columnName, "pivot_in_valuelist");
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    /**
     * Extract the pivot column name from a result column in the IN clause.
     * Handles both column references and constant values.
     */
    private String extractPivotColumnName(TResultColumn resultColumn) {
        if (resultColumn == null) return null;

        TExpression expr = resultColumn.getExpr();
        if (expr != null) {
            return extractColumnNameFromExpression(expr);
        }

        // Fallback to display name
        return stripDelimiters(resultColumn.getDisplayName());
    }

    /**
     * Extract column name from an expression (handles constants and column references).
     */
    private String extractColumnNameFromExpression(TExpression expr) {
        if (expr == null) return null;

        EExpressionType exprType = expr.getExpressionType();

        // Column reference (e.g., [Sammich], [Apple], "HOUSTON" in BigQuery)
        // For BigQuery, double-quoted identifiers like "HOUSTON" are parsed as simple_object_name_t
        // and we need to strip the quotes for display purposes.
        if (exprType == EExpressionType.simple_object_name_t) {
            TObjectName objName = expr.getObjectOperand();
            if (objName != null) {
                // Strip delimiters (quotes) for display - per identifier_normalization.md
                return stripDelimiters(objName.toString());
            }
        }

        // Constant value (e.g., [1], [2], 'value', "HOUSTON")
        if (exprType == EExpressionType.simple_constant_t) {
            TConstant constant = expr.getConstantOperand();
            if (constant != null && constant.getValueToken() != null) {
                String value = constant.getValueToken().toString();
                // Strip string delimiters (quotes)
                return stripDelimiters(value);
            }
        }

        // Fallback: use expression string with delimiters stripped
        return stripDelimiters(expr.toString());
    }

    /**
     * Add a pivot column to the namespace.
     */
    private void addPivotColumn(String columnName, String evidence) {
        pivotColumns.add(columnName);
        ColumnSource source = new ColumnSource(
            this,
            columnName,
            null,
            1.0,  // Definite - from PIVOT IN clause
            evidence
        );
        columnSources.put(columnName, source);
    }

    /**
     * Strip SQL string delimiters (double quotes and single quotes) from a name.
     * Note: SQL Server brackets and backticks are preserved as they're often
     * needed for special names like [1], [Date], etc.
     */
    private String stripDelimiters(String name) {
        if (name == null || name.isEmpty()) {
            return name;
        }
        // Strip double quotes (BigQuery style identifier)
        if (name.startsWith("\"") && name.endsWith("\"") && name.length() > 2) {
            return name.substring(1, name.length() - 1);
        }
        // Strip single quotes (string literal)
        if (name.startsWith("'") && name.endsWith("'") && name.length() > 2) {
            return name.substring(1, name.length() - 1);
        }
        // SQL Server brackets [] are preserved as they're part of the column identity
        return name;
    }

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

        // First check if it's a pivot column
        if (isPivotColumn(columnName)) {
            return ColumnLevel.EXISTS;
        }

        // UNPIVOT: IN(...) columns are consumed and do NOT exist in the virtual table output
        if (isUnpivot && isUnpivotInColumn(columnName)) {
            return ColumnLevel.NOT_EXISTS;
        }

        // Check in columnSources
        for (String existingCol : columnSources.keySet()) {
            if (nameMatcher.matches(existingCol, columnName)) {
                return ColumnLevel.EXISTS;
            }
        }

        // For PIVOT tables without explicit column list, any column MAYBE exists
        // (pass-through columns from source table)
        return ColumnLevel.MAYBE;
    }

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

        // First try to find in pivot columns (from IN clause)
        for (Map.Entry<String, ColumnSource> entry : columnSources.entrySet()) {
            if (nameMatcher.matches(entry.getKey(), columnName)) {
                return entry.getValue();
            }
        }

        // UNPIVOT: IN(...) columns are consumed and must NOT resolve as pass-through columns
        if (isUnpivot && isUnpivotInColumn(columnName)) {
            return null;
        }

        // Delta 2: Delegate to source namespace for pass-through columns
        // Pass-through columns are all source columns EXCEPT:
        // - The FOR column (pivot_column)
        // - The aggregate input column (value_column)
        if (sourceNamespace != null) {
            // Check if this is a pass-through column (not FOR or aggregate column)
            if (!isForColumn(columnName) && !isAggregateInputColumn(columnName)) {
                ColumnSource sourceColumn = sourceNamespace.resolveColumn(columnName);
                if (sourceColumn != null) {
                    // Both PIVOT and UNPIVOT pass-through columns keep their source table attribution
                    // (e.g., SELECT year FROM ... UNPIVOT ... -> year attributed to source table)
                    // This matches the expected behavior where passthrough columns should be
                    // attributed to their original source table, not the PIVOT/UNPIVOT virtual table.
                    return sourceColumn;
                }
            }
        }

        return null;
    }

    private boolean isUnpivotInColumn(String columnName) {
        if (columnName == null || columnName.isEmpty()) return false;
        for (String consumed : unpivotInColumns) {
            if (nameMatcher.matches(consumed, columnName)) {
                return true;
            }
        }
        return false;
    }

    /**
     * Check if the column is the FOR column (pivot_column) in the PIVOT clause.
     * The FOR column is consumed by PIVOT and should not be passed through.
     */
    private boolean isForColumn(String columnName) {
        if (pivotClause == null || columnName == null) {
            return false;
        }

        // Try pivotColumnList first (modern API, used by BigQuery and others)
        TObjectNameList pivotColumns = pivotClause.getPivotColumnList();
        if (pivotColumns != null && pivotColumns.size() > 0) {
            for (int i = 0; i < pivotColumns.size(); i++) {
                TObjectName forColumn = pivotColumns.getObjectName(i);
                if (forColumn != null) {
                    String forColName = forColumn.getColumnNameOnly();
                    if (forColName != null && nameMatcher.matches(forColName, columnName)) {
                        return true;
                    }
                }
            }
        }

        // Fallback to deprecated singular method
        @SuppressWarnings("deprecation")
        TObjectName forColumn = pivotClause.getPivotColumn();
        if (forColumn != null) {
            String forColName = forColumn.getColumnNameOnly();
            if (forColName != null && nameMatcher.matches(forColName, columnName)) {
                return true;
            }
        }
        return false;
    }

    /**
     * Check if the column is the aggregate input column (value_column) in the PIVOT clause.
     * The aggregate input column is consumed by PIVOT and should not be passed through.
     */
    private boolean isAggregateInputColumn(String columnName) {
        if (pivotClause == null || columnName == null) {
            return false;
        }

        // Try to extract columns from aggregate function (SQL Server / BigQuery style)
        // e.g., PIVOT(AVG(departure_delay) FOR airline IN (...))
        TFunctionCall aggFunc = pivotClause.getAggregation_function();
        if (aggFunc != null && aggFunc.getArgs() != null) {
            for (int i = 0; i < aggFunc.getArgs().size(); i++) {
                TExpression argExpr = aggFunc.getArgs().getExpression(i);
                if (argExpr != null && argExpr.getObjectOperand() != null) {
                    TObjectName argCol = argExpr.getObjectOperand();
                    String argColName = argCol.getColumnNameOnly();
                    if (argColName != null && nameMatcher.matches(argColName, columnName)) {
                        return true;
                    }
                }
            }
        }

        // Try to extract columns from aggregate function list (Oracle style)
        TResultColumnList aggFuncList = pivotClause.getAggregation_function_list();
        if (aggFuncList != null) {
            for (int i = 0; i < aggFuncList.size(); i++) {
                TResultColumn rc = aggFuncList.getResultColumn(i);
                if (rc != null && rc.getExpr() != null) {
                    TExpression expr = rc.getExpr();
                    if (expr.getFunctionCall() != null && expr.getFunctionCall().getArgs() != null) {
                        for (int j = 0; j < expr.getFunctionCall().getArgs().size(); j++) {
                            TExpression argExpr = expr.getFunctionCall().getArgs().getExpression(j);
                            if (argExpr != null && argExpr.getObjectOperand() != null) {
                                TObjectName argCol = argExpr.getObjectOperand();
                                String argColName = argCol.getColumnNameOnly();
                                if (argColName != null && nameMatcher.matches(argColName, columnName)) {
                                    return true;
                                }
                            }
                        }
                    }
                }
            }
        }

        // Fallback: use valueColumnList (for UNPIVOT or direct value column specification)
        TObjectNameList valueColumns = pivotClause.getValueColumnList();
        if (valueColumns != null && valueColumns.size() > 0) {
            for (int i = 0; i < valueColumns.size(); i++) {
                TObjectName valueCol = valueColumns.getObjectName(i);
                if (valueCol != null) {
                    String valueColName = valueCol.getColumnNameOnly();
                    if (valueColName != null && nameMatcher.matches(valueColName, columnName)) {
                        return true;
                    }
                }
            }
        }

        // Fallback to deprecated singular method
        @SuppressWarnings("deprecation")
        TObjectName valueCol = pivotClause.getValueColumn();
        if (valueCol != null) {
            String valueColName = valueCol.getColumnNameOnly();
            if (valueColName != null && nameMatcher.matches(valueColName, columnName)) {
                return true;
            }
        }
        return false;
    }

    /**
     * Get all column sources for this PIVOT namespace.
     * This includes both pivot columns (from IN clause) and pass-through columns from the source.
     *
     * <p>For PIVOT, the output columns are:
     * <ul>
     *   <li>Pivot columns: generated from IN clause values (e.g., 'AA', 'KH', 'DL', '9E')</li>
     *   <li>Pass-through columns: source columns NOT consumed by FOR or aggregate</li>
     * </ul>
     *
     * <p>This method is critical for star column expansion: when outer query uses SELECT *
     * from a PIVOT table, we need to return ALL columns the pivot table exposes.
     */
    @Override
    public Map<String, ColumnSource> getAllColumnSources() {
        ensureValidated();

        if (columnSources == null) {
            return Collections.emptyMap();
        }

        // Start with pivot columns (from IN clause)
        Map<String, ColumnSource> allColumns = new LinkedHashMap<>(columnSources);

        // Add pass-through columns from source namespace
        // Pass-through = all source columns EXCEPT FOR column and aggregate input column
        if (sourceNamespace != null) {
            Map<String, ColumnSource> sourceColumns = sourceNamespace.getAllColumnSources();
            for (Map.Entry<String, ColumnSource> entry : sourceColumns.entrySet()) {
                String colName = entry.getKey();

                // Skip if already a pivot column (from IN clause)
                if (allColumns.containsKey(colName)) {
                    continue;
                }

                // Skip FOR column (consumed by PIVOT)
                if (isForColumn(colName)) {
                    continue;
                }

                // Skip aggregate input column (consumed by PIVOT)
                if (isAggregateInputColumn(colName)) {
                    continue;
                }

                // This is a pass-through column - add it with pivot table attribution
                ColumnSource sourceCol = entry.getValue();
                ColumnSource passthroughCol = new ColumnSource(
                    this,  // PivotNamespace
                    colName,
                    sourceCol.getDefinitionNode(),
                    sourceCol.getConfidence(),
                    "pivot_passthrough",
                    pivotTable,  // attribute to pivot table
                    null
                );
                allColumns.put(colName, passthroughCol);
            }
        }
        return Collections.unmodifiableMap(allColumns);
    }

    @Override
    public String toString() {
        return "PivotNamespace(" + getDisplayName() + ", pivotCols=" + pivotColumns.size() + ")";
    }
}