"""Tests for andlake.trino."""

from __future__ import annotations

import sys
from unittest.mock import MagicMock

import pytest

from andlake.config import reset_config


@pytest.fixture(autouse=True)
def _clean_config():
    reset_config()
    yield
    reset_config()


@pytest.fixture()
def mock_trino(monkeypatch: pytest.MonkeyPatch):
    """Install a mock ``trino`` package into sys.modules."""
    mock = MagicMock()
    monkeypatch.setitem(sys.modules, "trino", mock)
    monkeypatch.setitem(sys.modules, "trino.dbapi", mock.dbapi)
    monkeypatch.setitem(sys.modules, "trino.auth", mock.auth)
    return mock


@pytest.fixture()
def mock_sqlalchemy(monkeypatch: pytest.MonkeyPatch):
    """Install a mock ``sqlalchemy`` into sys.modules."""
    mock = MagicMock()
    monkeypatch.setitem(sys.modules, "sqlalchemy", mock)
    return mock


class TestGetTrinoConnection:
    """get_trino_connection() creates a properly configured connection."""

    def test_default_no_auth(
        self, monkeypatch: pytest.MonkeyPatch, mock_trino: MagicMock
    ):
        """Without tenant_id/access_token, connects without auth."""
        # No ANDLAKE_TENANT_ID or ANDLAKE_ACCESS_TOKEN set
        monkeypatch.delenv("ANDLAKE_TENANT_ID", raising=False)
        monkeypatch.delenv("ANDLAKE_ACCESS_TOKEN", raising=False)

        mock_conn = MagicMock()
        mock_trino.dbapi.connect.return_value = mock_conn

        from andlake.trino import get_trino_connection

        conn = get_trino_connection()

        call_kwargs = mock_trino.dbapi.connect.call_args.kwargs
        assert call_kwargs["host"] == "notebook-service"
        assert call_kwargs["port"] == 8082
        assert call_kwargs["catalog"] == "lake"
        assert call_kwargs["http_scheme"] == "http"
        assert "auth" not in call_kwargs
        assert "user" not in call_kwargs
        assert conn is mock_conn

    def test_with_auth(
        self, monkeypatch: pytest.MonkeyPatch, mock_trino: MagicMock
    ):
        """With tenant_id and access_token, uses JWT auth."""
        monkeypatch.setenv("ANDLAKE_TENANT_ID", "tenant-xyz")
        monkeypatch.setenv("ANDLAKE_ACCESS_TOKEN", "jwt-test-token")

        mock_auth_instance = MagicMock()
        mock_trino.auth.JWTAuthentication.return_value = mock_auth_instance

        from andlake.trino import get_trino_connection

        get_trino_connection()

        call_kwargs = mock_trino.dbapi.connect.call_args.kwargs
        assert call_kwargs["user"] == "tenant-xyz"
        assert call_kwargs["auth"] is mock_auth_instance
        mock_trino.auth.JWTAuthentication.assert_called_once_with("jwt-test-token")

    def test_custom_catalog_and_schema(
        self, monkeypatch: pytest.MonkeyPatch, mock_trino: MagicMock
    ):
        from andlake.trino import get_trino_connection

        get_trino_connection(catalog="analytics", schema="gold")

        call_kwargs = mock_trino.dbapi.connect.call_args.kwargs
        assert call_kwargs["catalog"] == "analytics"
        assert call_kwargs["schema"] == "gold"

    def test_custom_gateway(
        self, monkeypatch: pytest.MonkeyPatch, mock_trino: MagicMock
    ):
        monkeypatch.setenv("TRINO_HOST", "custom-gw")
        monkeypatch.setenv("TRINO_PORT", "9999")

        from andlake.trino import get_trino_connection

        get_trino_connection()

        call_kwargs = mock_trino.dbapi.connect.call_args.kwargs
        assert call_kwargs["host"] == "custom-gw"
        assert call_kwargs["port"] == 9999


class TestGetTrinoEngine:
    """get_trino_engine() builds a valid SQLAlchemy connection URL."""

    def test_default_engine_url_no_auth(
        self,
        monkeypatch: pytest.MonkeyPatch,
        mock_trino: MagicMock,
        mock_sqlalchemy: MagicMock,
    ):
        """Without tenant_id, URL has no user@ prefix."""
        monkeypatch.delenv("ANDLAKE_TENANT_ID", raising=False)
        monkeypatch.delenv("ANDLAKE_ACCESS_TOKEN", raising=False)

        from andlake.trino import get_trino_engine

        get_trino_engine()

        url = mock_sqlalchemy.create_engine.call_args.args[0]
        assert url == "trino://notebook-service:8082/lake"

        connect_args = mock_sqlalchemy.create_engine.call_args.kwargs["connect_args"]
        assert "auth" not in connect_args

    def test_engine_with_auth(
        self,
        monkeypatch: pytest.MonkeyPatch,
        mock_trino: MagicMock,
        mock_sqlalchemy: MagicMock,
    ):
        monkeypatch.setenv("ANDLAKE_TENANT_ID", "tenant-xyz")
        monkeypatch.setenv("ANDLAKE_ACCESS_TOKEN", "jwt-test-token")

        from andlake.trino import get_trino_engine

        get_trino_engine()

        url = mock_sqlalchemy.create_engine.call_args.args[0]
        assert url == "trino://tenant-xyz@notebook-service:8082/lake"

    def test_engine_with_schema(
        self,
        monkeypatch: pytest.MonkeyPatch,
        mock_trino: MagicMock,
        mock_sqlalchemy: MagicMock,
    ):
        from andlake.trino import get_trino_engine

        get_trino_engine(catalog="analytics", schema="gold")

        url = mock_sqlalchemy.create_engine.call_args.args[0]
        assert url.endswith("/analytics/gold")
