package gudusoft.gsqlparser.resolver2.context;

import gudusoft.gsqlparser.EDbObjectType;
import gudusoft.gsqlparser.EDbVendor;
import gudusoft.gsqlparser.ESqlStatementType;
import gudusoft.gsqlparser.TCustomSqlStatement;
import gudusoft.gsqlparser.TStatementList;
import gudusoft.gsqlparser.sqlenv.TSQLEnv;
import gudusoft.gsqlparser.stmt.TSetSchemaStmt;
import gudusoft.gsqlparser.stmt.TUseDatabase;
import gudusoft.gsqlparser.stmt.TUseStmt;

/**
 * Delta 4: Database context tracker for USE/SET statements.
 *
 * Tracks USE DATABASE, USE SCHEMA, SET SCHEMA, and similar statements
 * to maintain the current database/schema context during resolution.
 *
 * This enables proper resolution of unqualified table names when the
 * context is set by prior statements in the same SQL batch.
 *
 * Example:
 * <pre>
 * USE mydb;
 * SELECT * FROM t;  -- 't' should resolve to 'mydb.dbo.t' (or 'mydb.t')
 * </pre>
 */
public class DatabaseContextTracker {

    private String currentDatabase;
    private String currentSchema;

    /**
     * Process all statements in the list to extract database/schema context.
     *
     * @param statements The list of SQL statements to process
     */
    public void processStatements(TStatementList statements) {
        if (statements == null) {
            return;
        }

        for (int i = 0; i < statements.size(); i++) {
            TCustomSqlStatement stmt = statements.get(i);
            if (stmt != null) {
                processStatement(stmt);
            }
        }
    }

    /**
     * Process a single statement and update context if it's a USE/SET statement.
     *
     * @param stmt The statement to process
     */
    public void processStatement(TCustomSqlStatement stmt) {
        if (stmt == null) {
            return;
        }

        ESqlStatementType stmtType = stmt.sqlstatementtype;

        // Handle USE DATABASE
        if (stmtType == ESqlStatementType.sstUseDatabase) {
            processUseDatabase((TUseDatabase) stmt);
        }
        // Handle general USE statement (deprecated but still used)
        else if (stmtType == ESqlStatementType.sstUse) {
            processUseStmt((TUseStmt) stmt);
        }
        // Handle USE SCHEMA (Snowflake, Databricks)
        else if (stmtType == ESqlStatementType.sstUseSchema) {
            processUseSchema(stmt);
        }
        // Handle USE CATALOG (Databricks)
        else if (stmtType == ESqlStatementType.sstUseCatalog) {
            processUseCatalog(stmt);
        }
        // Handle SET SCHEMA (Netezza, etc.)
        else if (stmtType == ESqlStatementType.sstSetSchema) {
            processSetSchema((TSetSchemaStmt) stmt);
        }
    }

    /**
     * Process USE DATABASE statement.
     */
    private void processUseDatabase(TUseDatabase useDb) {
        if (useDb == null) {
            return;
        }

        // Check if it's actually a USE SCHEMA (some vendors use TUseDatabase for both)
        if (useDb.isSchema() && useDb.getSchemaName() != null) {
            currentSchema = useDb.getSchemaName().toString();
        } else if (useDb.getDatabaseName() != null) {
            currentDatabase = useDb.getDatabaseName().toString();
            // MantisBT 4677 — a Snowflake database switch also changes the
            // schema: the docs say the current schema becomes PUBLIC, or is
            // left unset when PUBLIC does not exist. Either way the schema
            // carried over from the PREVIOUS database is no longer valid, and
            // keeping it would silently resolve later unqualified names under
            // <newdb>.<oldschema> -- a table that need not exist.
            //
            // We clear rather than assign PUBLIC: that the old schema is wrong
            // is provable from the statement alone, whereas PUBLIC existing in
            // the new database is not. A catalog-backed consumer can apply the
            // PUBLIC default on top of a cleared schema; it cannot recover from
            // a stale one. Gated to Snowflake so vendors whose USE DATABASE has
            // no such side effect are unaffected.
            if (useDb.dbvendor == EDbVendor.dbvsnowflake) {
                currentSchema = null;
            }
        }
    }

    /**
     * Process general USE statement.
     */
    private void processUseStmt(TUseStmt useStmt) {
        if (useStmt == null || useStmt.getDbObjectName() == null) {
            return;
        }

        EDbObjectType objType = useStmt.getDbObjectType();
        String name = useStmt.getDbObjectName().toString();

        if (objType == EDbObjectType.database) {
            currentDatabase = name;
        } else if (objType == EDbObjectType.schema) {
            currentSchema = name;
        }
    }

    /**
     * Process USE SCHEMA statement (Snowflake, Databricks).
     */
    private void processUseSchema(TCustomSqlStatement stmt) {
        // USE SCHEMA typically has a schema name in the first token after USE SCHEMA
        // We need to extract it from the statement text or tokens
        if (stmt != null && stmt.toString() != null) {
            String text = stmt.toString().trim();
            // Pattern: USE SCHEMA schemaName or USE DATABASE schemaName
            String[] parts = text.split("\\s+");
            if (parts.length >= 3) {
                currentSchema = stripQuotes(parts[parts.length - 1]);
            } else if (parts.length >= 2) {
                currentSchema = stripQuotes(parts[parts.length - 1]);
            }
        }
    }

    /**
     * Process USE CATALOG statement (Databricks).
     */
    private void processUseCatalog(TCustomSqlStatement stmt) {
        // USE CATALOG catalogName or SET CATALOG catalogName
        if (stmt != null && stmt.toString() != null) {
            String text = stmt.toString().trim();
            String[] parts = text.split("\\s+");
            if (parts.length >= 2) {
                currentDatabase = stripQuotes(parts[parts.length - 1]);
            }
        }
    }

    /**
     * Process SET SCHEMA statement.
     */
    private void processSetSchema(TSetSchemaStmt setSchema) {
        if (setSchema != null && setSchema.getSchemaName() != null) {
            currentSchema = setSchema.getSchemaName().toString();
        }
    }

    /**
     * Strip quotes from an identifier.
     */
    private String stripQuotes(String name) {
        if (name == null || name.isEmpty()) {
            return name;
        }
        // Strip double quotes
        if (name.startsWith("\"") && name.endsWith("\"") && name.length() > 2) {
            return name.substring(1, name.length() - 1);
        }
        // Strip backticks
        if (name.startsWith("`") && name.endsWith("`") && name.length() > 2) {
            return name.substring(1, name.length() - 1);
        }
        // Strip brackets
        if (name.startsWith("[") && name.endsWith("]") && name.length() > 2) {
            return name.substring(1, name.length() - 1);
        }
        // Strip single quotes (for schema names in SET SCHEMA 'name')
        if (name.startsWith("'") && name.endsWith("'") && name.length() > 2) {
            return name.substring(1, name.length() - 1);
        }
        // Strip trailing semicolon
        if (name.endsWith(";")) {
            return name.substring(0, name.length() - 1);
        }
        return name;
    }

    /**
     * Apply the tracked defaults to a TSQLEnv.
     *
     * @param env The SQL environment to update
     */
    public void applyDefaults(TSQLEnv env) {
        if (env == null) {
            return;
        }

        if (currentDatabase != null && !currentDatabase.isEmpty()) {
            env.setDefaultCatalogName(currentDatabase);
        }
        if (currentSchema != null && !currentSchema.isEmpty()) {
            env.setDefaultSchemaName(currentSchema);
        }
    }

    /**
     * Check if any context has been tracked.
     *
     * @return true if either database or schema context has been set
     */
    public boolean hasContext() {
        return currentDatabase != null || currentSchema != null;
    }

    /**
     * Get the current database name.
     *
     * @return The current database name, or null if not set
     */
    public String getCurrentDatabase() {
        return currentDatabase;
    }

    /**
     * Get the current schema name.
     *
     * @return The current schema name, or null if not set
     */
    public String getCurrentSchema() {
        return currentSchema;
    }

    /**
     * Reset the context tracker.
     */
    public void reset() {
        currentDatabase = null;
        currentSchema = null;
    }

    @Override
    public String toString() {
        return "DatabaseContextTracker{" +
            "database='" + currentDatabase + '\'' +
            ", schema='" + currentSchema + '\'' +
            '}';
    }
}