# Copyright 2025-present DatusAI, Inc.
# Licensed under the Apache License, Version 2.0.
# See http://www.apache.org/licenses/LICENSE-2.0 for details.

"""Unit tests for sql_utils module."""

from datus_db_core.constants import SQLType
from datus_db_core.sql_utils import (
    _first_statement,
    metadata_identifier,
    parse_context_switch,
    parse_dialect,
    parse_read_dialect,
    parse_sql_type,
    strip_sql_comments,
)


class TestParseReadDialect:
    def test_postgres_variants(self):
        assert parse_read_dialect("postgres") == "postgres"
        assert parse_read_dialect("postgresql") == "postgres"
        assert parse_read_dialect("redshift") == "postgres"
        assert parse_read_dialect("greenplum") == "postgres"

    def test_hive_variants(self):
        assert parse_read_dialect("spark") == "hive"
        assert parse_read_dialect("databricks") == "hive"
        assert parse_read_dialect("hive") == "hive"
        assert parse_read_dialect("starrocks") == "hive"

    def test_mssql_variants(self):
        assert parse_read_dialect("mssql") == "tsql"
        assert parse_read_dialect("sqlserver") == "tsql"

    def test_passthrough(self):
        assert parse_read_dialect("snowflake") == "snowflake"
        assert parse_read_dialect("mysql") == "mysql"

    def test_empty_and_none(self):
        assert parse_read_dialect("") == ""
        assert parse_read_dialect(None) == ""


class TestParseDialect:
    def test_postgres_variants(self):
        assert parse_dialect("postgres") == "postgres"
        assert parse_dialect("postgresql") == "postgres"

    def test_mssql_variants(self):
        assert parse_dialect("mssql") == "tsql"
        assert parse_dialect("sqlserver") == "tsql"

    def test_passthrough(self):
        assert parse_dialect("snowflake") == "snowflake"
        assert parse_dialect("mysql") == "mysql"


class TestStripSqlComments:
    def test_block_comment(self):
        assert strip_sql_comments("SELECT /* comment */ 1").strip() == "SELECT   1"

    def test_line_comment(self):
        result = strip_sql_comments("SELECT 1 -- comment\nFROM t")
        assert "comment" not in result
        assert "FROM t" in result

    def test_no_comments(self):
        assert strip_sql_comments("SELECT 1") == "SELECT 1"

    def test_multiline_block_comment(self):
        sql = "SELECT /* this\nis\nmultiline */ 1"
        result = strip_sql_comments(sql)
        assert "multiline" not in result
        assert "1" in result


class TestFirstStatement:
    def test_single_statement(self):
        assert _first_statement("SELECT 1") == "SELECT 1"

    def test_multiple_statements(self):
        assert _first_statement("SELECT 1; SELECT 2") == "SELECT 1"

    def test_semicolon_in_string(self):
        result = _first_statement("SELECT 'a;b' FROM t")
        assert result == "SELECT 'a;b' FROM t"

    def test_empty_string(self):
        assert _first_statement("") == ""

    def test_whitespace_only(self):
        assert _first_statement("   ") == ""

    def test_with_comments(self):
        result = _first_statement("/* comment */ SELECT 1; SELECT 2")
        assert result == "SELECT 1"


class TestParseSqlType:
    def test_select(self):
        assert parse_sql_type("SELECT * FROM t", "snowflake") == SQLType.SELECT

    def test_select_with_whitespace(self):
        assert parse_sql_type("  SELECT 1  ", "snowflake") == SQLType.SELECT

    def test_insert(self):
        assert parse_sql_type("INSERT INTO t VALUES (1)", "mysql") == SQLType.INSERT

    def test_update(self):
        assert parse_sql_type("UPDATE t SET col=1", "mysql") == SQLType.UPDATE

    def test_delete(self):
        assert parse_sql_type("DELETE FROM t WHERE id=1", "mysql") == SQLType.DELETE

    def test_merge(self):
        assert (
            parse_sql_type("MERGE INTO t USING s ON t.id=s.id WHEN MATCHED THEN UPDATE SET t.v=s.v", "snowflake")
            == SQLType.MERGE
        )

    def test_create_table(self):
        assert parse_sql_type("CREATE TABLE t (id INT)", "mysql") == SQLType.DDL

    def test_alter_table(self):
        assert parse_sql_type("ALTER TABLE t ADD COLUMN name VARCHAR(50)", "mysql") == SQLType.DDL

    def test_drop_table(self):
        assert parse_sql_type("DROP TABLE t", "mysql") == SQLType.DDL

    def test_truncate(self):
        assert parse_sql_type("TRUNCATE TABLE t", "mysql") == SQLType.DDL

    def test_show_tables(self):
        assert parse_sql_type("SHOW TABLES", "mysql") == SQLType.METADATA_SHOW

    def test_describe(self):
        assert parse_sql_type("DESCRIBE t", "mysql") == SQLType.METADATA_SHOW

    def test_explain(self):
        assert parse_sql_type("EXPLAIN SELECT 1", "snowflake") == SQLType.EXPLAIN

    def test_use_database(self):
        assert parse_sql_type("USE my_db", "mysql") == SQLType.CONTENT_SET

    def test_set(self):
        assert parse_sql_type("SET search_path TO my_schema", "postgres") == SQLType.CONTENT_SET

    def test_with_cte_select(self):
        assert parse_sql_type("WITH cte AS (SELECT 1) SELECT * FROM cte", "snowflake") == SQLType.SELECT

    def test_empty_string(self):
        assert parse_sql_type("", "mysql") == SQLType.UNKNOWN

    def test_none_input(self):
        assert parse_sql_type(None, "mysql") == SQLType.UNKNOWN

    def test_whitespace_only(self):
        assert parse_sql_type("   ", "mysql") == SQLType.UNKNOWN

    def test_grant(self):
        assert parse_sql_type("GRANT SELECT ON t TO user1", "mysql") == SQLType.DDL

    def test_begin_transaction(self):
        assert parse_sql_type("BEGIN", "mysql") == SQLType.CONTENT_SET

    def test_commit(self):
        assert parse_sql_type("COMMIT", "mysql") == SQLType.CONTENT_SET

    def test_rollback(self):
        assert parse_sql_type("ROLLBACK", "mysql") == SQLType.CONTENT_SET

    def test_values(self):
        assert parse_sql_type("VALUES (1, 2, 3)", "snowflake") == SQLType.SELECT

    def test_replace(self):
        assert parse_sql_type("REPLACE INTO t VALUES (1)", "mysql") == SQLType.INSERT

    def test_vacuum(self):
        # VACUUM is parsed as Command by sqlglot, routed to CONTENT_SET
        # It only falls back to DDL via keyword map if sqlglot parsing fails
        result = parse_sql_type("VACUUM", "snowflake")
        assert result in (SQLType.DDL, SQLType.CONTENT_SET)

    def test_analyze(self):
        assert parse_sql_type("ANALYZE t", "postgres") == SQLType.DDL

    def test_pragma(self):
        assert parse_sql_type("PRAGMA table_info(t)", "sqlite") == SQLType.METADATA_SHOW


class TestParseContextSwitch:
    def test_use_database_mysql(self):
        result = parse_context_switch("USE my_db", "mysql")
        assert result is not None
        assert result["command"] == "USE"
        assert result["database_name"] == "my_db"
        assert result["target"] == "database"

    def test_use_schema_snowflake(self):
        result = parse_context_switch("USE db.schema1", "snowflake")
        assert result is not None
        assert result["database_name"] == "db"
        assert result["schema_name"] == "schema1"
        assert result["target"] == "schema"

    def test_use_database_snowflake_single(self):
        result = parse_context_switch("USE my_db", "snowflake")
        assert result is not None
        assert result["database_name"] == "my_db"
        assert result["target"] == "database"

    def test_use_catalog_keyword(self):
        result = parse_context_switch("USE CATALOG my_catalog", "trino")
        assert result is not None
        assert result["catalog_name"] == "my_catalog"
        assert result["target"] == "catalog"

    def test_use_database_keyword(self):
        result = parse_context_switch("USE DATABASE my_db", "snowflake")
        assert result is not None
        assert result["database_name"] == "my_db"
        assert result["target"] == "database"

    def test_use_schema_keyword(self):
        result = parse_context_switch("USE SCHEMA my_schema", "snowflake")
        assert result is not None
        assert result["schema_name"] == "my_schema"
        assert result["target"] == "schema"

    def test_set_catalog(self):
        result = parse_context_switch("SET CATALOG my_catalog", "trino")
        assert result is not None
        assert result["catalog_name"] == "my_catalog"
        assert result["target"] == "catalog"

    def test_set_database(self):
        result = parse_context_switch("SET DATABASE my_db", "snowflake")
        assert result is not None
        assert result["database_name"] == "my_db"
        assert result["target"] == "database"

    def test_set_schema(self):
        result = parse_context_switch("SET SCHEMA my_schema", "postgres")
        assert result is not None
        assert result["schema_name"] == "my_schema"
        assert result["target"] == "schema"

    def test_set_schema_with_equals(self):
        result = parse_context_switch("SET SCHEMA = my_schema", "postgres")
        assert result is not None
        assert result["schema_name"] == "my_schema"

    def test_set_schema_with_to(self):
        result = parse_context_switch("SET SCHEMA TO my_schema", "postgres")
        assert result is not None
        assert result["schema_name"] == "my_schema"

    def test_non_context_command_returns_none(self):
        assert parse_context_switch("SELECT 1", "mysql") is None

    def test_empty_input(self):
        assert parse_context_switch("", "mysql") is None

    def test_none_input(self):
        assert parse_context_switch(None, "mysql") is None

    def test_use_duckdb_single_identifier(self):
        result = parse_context_switch("USE my_schema", "duckdb")
        assert result is not None
        assert result["schema_name"] == "my_schema"
        assert result["fuzzy"] is True

    def test_use_duckdb_two_part(self):
        result = parse_context_switch("USE my_db.my_schema", "duckdb")
        assert result is not None
        assert result["database_name"] == "my_db"
        assert result["schema_name"] == "my_schema"
        assert result["fuzzy"] is False

    def test_use_starrocks(self):
        result = parse_context_switch("USE my_db", "starrocks")
        assert result is not None
        assert result["database_name"] == "my_db"
        assert result["target"] == "database"

    def test_raw_preserved(self):
        result = parse_context_switch("USE my_db", "mysql")
        assert result["raw"] == "USE my_db"

    def test_set_session_schema(self):
        result = parse_context_switch("SET SESSION SCHEMA my_schema", "postgres")
        assert result is not None
        assert result["schema_name"] == "my_schema"


class TestMetadataIdentifier:
    def test_sqlite_with_database(self):
        result = metadata_identifier(database_name="main", table_name="users", dialect="sqlite")
        assert result == "main.users"

    def test_sqlite_without_database(self):
        result = metadata_identifier(table_name="users", dialect="sqlite")
        assert result == "users"

    def test_duckdb(self):
        result = metadata_identifier(database_name="db", schema_name="public", table_name="t", dialect="duckdb")
        assert result == "db.public.t"

    def test_with_registry_capabilities(self):
        from datus_db_core.registry import ConnectorRegistry

        saved_caps = ConnectorRegistry._capabilities.copy()
        try:
            ConnectorRegistry._capabilities["testdialect"] = {"database", "schema"}
            result = metadata_identifier(
                database_name="mydb", schema_name="mysch", table_name="mytable", dialect="testdialect"
            )
            assert result == "mydb.mysch.mytable"
        finally:
            ConnectorRegistry._capabilities = saved_caps

    def test_table_only(self):
        result = metadata_identifier(table_name="users", dialect="snowflake")
        assert result == "users"
