
from typing import Optional

import psycopg
import pytest

from qgis.core import (
    QgsAbstractDatabaseProviderConnection,
    QgsDataSourceUri,
    QgsVectorLayer,
)

from lizmap.project_checker_tools import (
    InvalidType,
    auto_generated_primary_key_field,
    invalid_type_primary_key,
    table_type,
)

# To run these tests:
# * copy credentials.py.example to credentials.py
# * edit info credentials.py about PostgreSQL

try:
    from .credentials import (
        PG_DATABASE,
        PG_HOST,
        PG_PASSWORD,
        PG_PORT,
        PG_USER,
    )

    CREDENTIALS = True
except ImportError:
    CREDENTIALS = False
    PG_HOST = ""
    PG_PORT = ""
    PG_PASSWORD = ""
    PG_USER = ""
    PG_DATABASE = ""


from .compat import TestCase


@pytest.fixture(scope="class")
def connection():
    return psycopg.connect(
        f"host={PG_HOST} user={PG_USER} password={PG_PASSWORD} port={PG_PORT} dbname={PG_DATABASE}"
    )


@pytest.fixture(scope="class")
def cursor(connection):
    cursor = connection.cursor()
    cursor.execute(f"DROP SCHEMA IF EXISTS {TestSql.schema} CASCADE;")
    cursor.execute("CREATE EXTENSION IF NOT EXISTS postgis;")
    cursor.execute(f"CREATE SCHEMA {TestSql.schema};")

    connection.commit()
    return cursor


@pytest.mark.skipif(
    not all((CREDENTIALS, PG_HOST, PG_PORT, PG_PASSWORD, PG_USER, PG_DATABASE)),
    reason="No PG configuration defined",
)
class TestSql(TestCase):
    schema = "lizmap_plugin_tests"

    @classmethod
    def uri(cls) -> QgsDataSourceUri:
        """URI of the database."""
        uri = QgsDataSourceUri()
        uri.setConnection(
            PG_HOST,
            PG_PORT,
            PG_DATABASE,
            PG_USER,
            PG_PASSWORD,
            QgsDataSourceUri.SslMode.SslPrefer,
        )
        return uri

    def _create_views(
        self,
        table_name: str,
        cursor: psycopg.Cursor,
        connection: psycopg.Connection):
        """Helper to create views."""
        sql = (
            f"CREATE VIEW {self.schema}.view_{table_name} AS SELECT * FROM {self.schema}.{table_name};"
            f"CREATE MATERIALIZED VIEW {self.schema}.view_m_{table_name} AS SELECT * FROM {self.schema}.{table_name};"
        )
        cursor.execute(sql)
        connection.commit()

    def _vector_layer(
        self,
        table_name: str,
        prefix: str = "",
        pk: Optional[str] = None,
    ) -> QgsVectorLayer:
        """Helper to have the QgsVectorLayer."""
        uri = QgsDataSourceUri(self.uri())
        uri.setSchema(self.schema)
        uri.setTable(prefix + "_" + table_name if prefix else table_name)

        if pk:
            uri.setKeyColumn(pk)
        return QgsVectorLayer(uri.uri(False), "test", "postgres")

    def test_invalid_auto_generated_pk(
        self,
        cursor: psycopg.Cursor,
        connection: psycopg.Connection,
    ):
        """Test invalid autogenerated PK."""
        table_name = "test_tid"
        field_name = "id_int"
        # The table is missing a primary key
        sql = (
            f"CREATE TABLE IF NOT EXISTS {self.schema}.{table_name}"
            f"("
            f"  {field_name} INT,"
            f"  label TEXT"
            f");"  # fmt: ignore
        )
        cursor.execute(sql)
        connection.commit()

        self._create_views(table_name, cursor, connection)

        layer = self._vector_layer(table_name)

        # Only testing QGIS
        self.assertListEqual([], layer.primaryKeyAttributes())
        new_uri = QgsDataSourceUri(layer.source())
        self.assertEqual("tid", new_uri.keyColumn())

        # Lizmap check
        result, field = auto_generated_primary_key_field(layer)
        self.assertTrue(result)
        self.assertEqual("tid", field)

        # View
        del layer
        layer_view = self._vector_layer(table_name, "view")

        # Only testing QGIS
        self.assertListEqual([], layer_view.primaryKeyAttributes())
        new_uri = QgsDataSourceUri(layer_view.source())
        self.assertEqual("", new_uri.keyColumn())

        # Lizmap check
        result, field = auto_generated_primary_key_field(layer_view)
        self.assertTrue(result)
        self.assertEqual("", field)

        # Materialized view
        del layer_view
        layer_materialized_view = self._vector_layer(table_name, "view_m")

        # Only testing QGIS
        self.assertListEqual([], layer_materialized_view.primaryKeyAttributes())
        new_uri = QgsDataSourceUri(layer_materialized_view.source())
        self.assertEqual("", new_uri.keyColumn())

        # Lizmap check
        result, field = auto_generated_primary_key_field(layer_materialized_view)
        self.assertTrue(result)
        self.assertEqual("", field)

    def test_invalid_varchar_pk(self, cursor, connection):
        """Test invalid varchar primary key."""
        table_name = "test_varchar"
        field_name = "municipality_code"
        sql = (
            f"CREATE TABLE IF NOT EXISTS {self.schema}.{table_name}"
            f"("
            f"  {field_name} VARCHAR(10) PRIMARY KEY,"
            f"  label TEXT"
            f");"
        )
        cursor.execute(sql)
        connection.commit()

        self._create_views(table_name, cursor, connection)

        layer = self._vector_layer(table_name)

        # Only testing QGIS
        self.assertListEqual([0], layer.primaryKeyAttributes())
        self.assertEqual(field_name, layer.fields().at(layer.primaryKeyAttributes()[0]).name())
        new_uri = QgsDataSourceUri(layer.source())
        self.assertEqual(field_name, new_uri.keyColumn())

        # Lizmap check
        self.assertFalse(invalid_type_primary_key(layer, InvalidType.Int8))
        self.assertTrue(invalid_type_primary_key(layer, InvalidType.Varchar))

        # View
        del layer
        layer_view = self._vector_layer(table_name, "view")

        # Only testing QGIS
        self.assertListEqual([], layer_view.primaryKeyAttributes())
        new_uri = QgsDataSourceUri(layer_view.source())
        self.assertEqual("", new_uri.keyColumn())

        # Lizmap check
        self.assertFalse(invalid_type_primary_key(layer_view, InvalidType.Int8))
        self.assertFalse(invalid_type_primary_key(layer_view, InvalidType.Varchar))

        # Materialized view
        del layer_view
        layer_materialized_view = self._vector_layer(table_name, "view_m")

        # Only testing QGIS
        self.assertListEqual([], layer_materialized_view.primaryKeyAttributes())
        new_uri = QgsDataSourceUri(layer_materialized_view.source())
        self.assertEqual("", new_uri.keyColumn())

        # Lizmap check
        self.assertFalse(invalid_type_primary_key(layer_materialized_view, InvalidType.Int8))
        self.assertFalse(invalid_type_primary_key(layer_materialized_view, InvalidType.Varchar))

    def test_invalid_int8_pk(self, cursor, connection):
        """Test invalid int8 (bigint) primary key."""
        table_name = "test_bigint"
        field_name = "id"
        sql = (
            f"CREATE TABLE IF NOT EXISTS {self.schema}.{table_name}"
            f"("
            f"  {field_name} BIGINT PRIMARY KEY,"
            f"  label TEXT"
            f");"
        )
        self.cursor.execute(sql)
        self.connection.commit()

        self._create_views(table_name, cursor, connection)

        layer = self._vector_layer(table_name)

        # Only testing QGIS
        self.assertListEqual([0], layer.primaryKeyAttributes())
        self.assertEqual(field_name, layer.fields().at(layer.primaryKeyAttributes()[0]).name())
        new_uri = QgsDataSourceUri(layer.source())
        self.assertEqual(field_name, new_uri.keyColumn())

        # Lizmap check
        self.assertTrue(invalid_type_primary_key(layer, InvalidType.Int8))
        self.assertFalse(invalid_type_primary_key(layer, InvalidType.Varchar))

        # View
        del layer
        layer_view = self._vector_layer(table_name, "view")

        # Only testing QGIS
        self.assertListEqual([], layer_view.primaryKeyAttributes())
        new_uri = QgsDataSourceUri(layer_view.source())
        self.assertEqual("", new_uri.keyColumn())

        # Lizmap check
        self.assertFalse(invalid_type_primary_key(layer_view, InvalidType.Int8))
        self.assertFalse(invalid_type_primary_key(layer_view, InvalidType.Varchar))

        # Materialized view
        del layer_view
        layer_materialized_view = self._vector_layer(table_name, "view_m")

        # Only testing QGIS
        self.assertListEqual([], layer_materialized_view.primaryKeyAttributes())
        new_uri = QgsDataSourceUri(layer_materialized_view.source())
        self.assertEqual("", new_uri.keyColumn())

        # Lizmap check
        self.assertFalse(invalid_type_primary_key(layer_materialized_view, InvalidType.Int8))
        self.assertFalse(invalid_type_primary_key(layer_materialized_view, InvalidType.Varchar))

    def test_invalid_composite_pk(self, cursor, connection):
        """Test invalid composite primary keys."""
        table_name = "test_composite"
        field_name_1 = "id_1"
        field_name_2 = "id_2"
        sql = (
            f"CREATE TABLE IF NOT EXISTS {self.schema}.{table_name}"
            f"("
            f"  {field_name_1} INT,"
            f"  {field_name_2} INT,"
            f"  key_2 INT,"
            f"  label TEXT,"
            f"  PRIMARY KEY ({field_name_1}, {field_name_2})"
            f");"
        )
        cursor.execute(sql)
        connection.commit()

        self._create_views(table_name, cursor, connection)

        layer = self._vector_layer(table_name)

        # Only testing QGIS
        self.assertListEqual([0, 1], layer.primaryKeyAttributes())
        self.assertEqual(field_name_1, layer.fields().at(layer.primaryKeyAttributes()[0]).name())
        self.assertEqual(field_name_2, layer.fields().at(layer.primaryKeyAttributes()[1]).name())
        new_uri = QgsDataSourceUri(layer.source())
        self.assertEqual(f"{field_name_1},{field_name_2}", new_uri.keyColumn())

        # Lizmap check
        # TODO check what we expect
        self.assertFalse(invalid_type_primary_key(layer, InvalidType.Int8))
        self.assertFalse(invalid_type_primary_key(layer, InvalidType.Varchar))

        result, field = auto_generated_primary_key_field(layer)
        self.assertFalse(result)
        self.assertIsNone(field)

        # View
        del layer
        layer_view = self._vector_layer(table_name, "view")

        # Only testing QGIS
        self.assertListEqual([], layer_view.primaryKeyAttributes())
        new_uri = QgsDataSourceUri(layer_view.source())
        self.assertEqual("", new_uri.keyColumn())

        # Lizmap check
        self.assertFalse(invalid_type_primary_key(layer_view, InvalidType.Int8))
        self.assertFalse(invalid_type_primary_key(layer_view, InvalidType.Varchar))
        result, field = auto_generated_primary_key_field(layer_view)
        self.assertTrue(result)
        self.assertEqual("", field)

        # Materialized view
        del layer_view
        layer_materialized_view = self._vector_layer(table_name, "view_m")

        # Only testing QGIS
        self.assertListEqual([], layer_materialized_view.primaryKeyAttributes())
        new_uri = QgsDataSourceUri(layer_materialized_view.source())
        self.assertEqual("", new_uri.keyColumn())

        # Lizmap check
        self.assertFalse(invalid_type_primary_key(layer_materialized_view, InvalidType.Int8))
        self.assertFalse(invalid_type_primary_key(layer_materialized_view, InvalidType.Varchar))
        result, field = auto_generated_primary_key_field(layer_materialized_view)
        self.assertTrue(result)
        self.assertEqual("", field)

    def test_valid_serial_pk(self, cursor, connection):
        """Test valid serial/integer primary key."""
        table_name = "test_int"
        field_name = "id"
        sql = (
            f"CREATE TABLE IF NOT EXISTS {self.schema}.{table_name}"
            f"("
            f"  {field_name} SERIAL PRIMARY KEY,"
            f"  label TEXT"
            f");"
        )
        cursor.execute(sql)
        connection.commit()

        self._create_views(table_name, cursor, connection)

        layer = self._vector_layer(table_name)

        # Only testing QGIS
        self.assertListEqual([0], layer.primaryKeyAttributes())
        self.assertEqual(field_name, layer.fields().at(layer.primaryKeyAttributes()[0]).name())
        new_uri = QgsDataSourceUri(layer.source())
        self.assertEqual(field_name, new_uri.keyColumn())

        # Lizmap check
        self.assertFalse(invalid_type_primary_key(layer, InvalidType.Int8))
        self.assertFalse(invalid_type_primary_key(layer, InvalidType.Varchar))

        # View
        # We need to give the PK
        del layer
        layer_view = self._vector_layer(table_name, "view", pk=field_name)

        # Only testing QGIS
        self.assertListEqual([0], layer_view.primaryKeyAttributes())
        self.assertEqual(field_name, layer_view.fields().at(layer_view.primaryKeyAttributes()[0]).name())
        new_uri = QgsDataSourceUri(layer_view.source())
        self.assertEqual(field_name, new_uri.keyColumn())

        # Lizmap check
        self.assertFalse(invalid_type_primary_key(layer_view, InvalidType.Int8))
        self.assertFalse(invalid_type_primary_key(layer_view, InvalidType.Varchar))
        result, field = auto_generated_primary_key_field(layer_view)
        self.assertFalse(result)
        self.assertIsNone(field)

        # Materialized view
        del layer_view
        # We need to give the PK
        layer_materialized_view = self._vector_layer(table_name, "view_m", pk=field_name)

        # Only testing QGIS
        self.assertListEqual([0], layer_materialized_view.primaryKeyAttributes())
        self.assertEqual(
            field_name,
            layer_materialized_view.fields().at(layer_materialized_view.primaryKeyAttributes()[0]).name(),
        )
        new_uri = QgsDataSourceUri(layer_materialized_view.source())
        self.assertEqual(field_name, new_uri.keyColumn())

        # Lizmap check
        self.assertFalse(invalid_type_primary_key(layer_materialized_view, InvalidType.Int8))
        self.assertFalse(invalid_type_primary_key(layer_materialized_view, InvalidType.Varchar))
        result, field = auto_generated_primary_key_field(layer_materialized_view)
        self.assertFalse(result)
        self.assertIsNone(field)

    def test_valid_smallint_pk(self, cursor, connection):
        """Test valid small int (int2) primary key."""
        table_name = "test_int"
        field_name = "id"
        sql = (
            f"CREATE TABLE IF NOT EXISTS {self.schema}.{table_name}"
            f"("
            f"  {field_name} SMALLINT PRIMARY KEY,"
            f"  label TEXT"
            f");"
        )
        cursor.execute(sql)
        connection.commit()

        self._create_views(table_name, cursor, connection)

        layer = self._vector_layer(table_name)

        # Only testing QGIS
        self.assertListEqual([0], layer.primaryKeyAttributes())
        self.assertEqual(field_name, layer.fields().at(layer.primaryKeyAttributes()[0]).name())
        new_uri = QgsDataSourceUri(layer.source())
        self.assertEqual(field_name, new_uri.keyColumn())

        # Lizmap check
        # TODO check
        self.assertFalse(invalid_type_primary_key(layer, InvalidType.Int8))
        self.assertFalse(invalid_type_primary_key(layer, InvalidType.Varchar))
        result, field = auto_generated_primary_key_field(layer)
        self.assertFalse(result)
        self.assertIsNone(field)

        # View
        # We need to give the PK
        del layer
        layer_view = self._vector_layer(table_name, "view", pk=field_name)

        # Only testing QGIS
        self.assertListEqual([0], layer_view.primaryKeyAttributes())
        self.assertEqual(field_name, layer_view.fields().at(layer_view.primaryKeyAttributes()[0]).name())
        new_uri = QgsDataSourceUri(layer_view.source())
        self.assertEqual(field_name, new_uri.keyColumn())

        # Lizmap check
        self.assertFalse(invalid_type_primary_key(layer_view, InvalidType.Int8))
        self.assertFalse(invalid_type_primary_key(layer_view, InvalidType.Varchar))
        result, field = auto_generated_primary_key_field(layer_view)
        self.assertFalse(result)
        self.assertIsNone(field)

        # Materialized view
        del layer_view
        # We need to give the PK
        layer_materialized_view = self._vector_layer(table_name, "view_m", pk=field_name)

        # Only testing QGIS
        self.assertListEqual([0], layer_materialized_view.primaryKeyAttributes())
        self.assertEqual(
            field_name,
            layer_materialized_view.fields().at(layer_materialized_view.primaryKeyAttributes()[0]).name(),
        )
        new_uri = QgsDataSourceUri(layer_materialized_view.source())
        self.assertEqual(field_name, new_uri.keyColumn())

        # Lizmap check
        self.assertFalse(invalid_type_primary_key(layer_materialized_view, InvalidType.Int8))
        self.assertFalse(invalid_type_primary_key(layer_materialized_view, InvalidType.Varchar))
        result, field = auto_generated_primary_key_field(layer_materialized_view)
        self.assertFalse(result)
        self.assertIsNone(field)

    def test_detect_view(self, cursor, connection):
        """Test to check we can detect views."""
        table_name = "test"
        sql = (
            f"CREATE TABLE IF NOT EXISTS {self.schema}.{table_name}"
            f"("
            f"  id SERIAL PRIMARY KEY,"
            f"  label TEXT"
            f");"  # fmt: ignore
        )
        cursor.execute(sql)
        connection.commit()

        self._create_views(table_name)

        self.assertEqual(
            QgsAbstractDatabaseProviderConnection.TableFlag.Vector, table_type(self._vector_layer(table_name))
        )
        self.assertEqual(
            QgsAbstractDatabaseProviderConnection.TableFlag.View,
            table_type(self._vector_layer(table_name, "view")),
        )
        self.assertEqual(
            QgsAbstractDatabaseProviderConnection.TableFlag.MaterializedView,
            table_type(self._vector_layer(table_name, "view_m")),
        )
