from typing import List

import sqlalchemy

import shovl.schemas as schemas
import shovl.services as services

from .service import ServiceProvider


class DatabaseProvider(ServiceProvider):
    """
    Database provider for connecting to and inspecting databases using
    SQLAlchemy.
    """

    config: schemas.DatabaseConfig

    engine: sqlalchemy.engine.Engine | None
    inspector: sqlalchemy.Inspector | None
    connection: sqlalchemy.Connection | None

    def __init__(self, config: schemas.DatabaseConfig) -> None:
        """
        Intialize the database provider with the given configuration.
        """
        super().__init__()
        self.config = config

    def _cleanup(self) -> None:
        """
        Cleanup the provider's state.
        """
        self.engine = None
        self.inspector = None
        self.connection = None

    @property
    def connected(self) -> bool:
        """
        Check if the provider is connected to a database.
        """
        return (
            self.engine is not None
            or self.inspector is not None
            or self.connection is not None
        )

    def connect(self) -> None:
        """
        Connect to a database.
        """
        if self.connected:
            self.disconnect()

        self.engine = sqlalchemy.create_engine(
            url=services.resolve_credentials(self.config.url, parse_url=True)
        )

        self.inspector = sqlalchemy.inspect(self.engine)

        self.connection = self.engine.connect()

    def disconnect(self) -> None:
        """
        Disconnect from the database.
        """
        if self.connection:
            self.connection.close()
            self.connection = None

        if self.engine:
            self.engine.dispose()
            self.engine = None

        self._cleanup()

    def inspect_schemas(self) -> List[schemas.DatabaseEntity]:
        """
        Inspect the database's schemas.
        """
        if not self.connected:
            raise Exception("Not connected to a database.")

        if not self.inspector:
            raise Exception("Database inspector is not initialized.")

        names = self.inspector.get_schema_names()

        result = [
            schemas.DatabaseEntity(
                name=name,
                type=schemas.DatabaseEntityType.schema,
            )
            for name in names
        ]

        return result

    def inspect_tables(self, schema: str) -> List[schemas.DatabaseEntity]:
        """
        Inspect a schema's tables.
        """
        if not self.connected:
            raise Exception("Not connected to a database.")

        if not self.inspector:
            raise Exception("Database inspector is not initialized.")

        names = self.inspector.get_table_names(schema=schema)

        result = [
            schemas.DatabaseEntity(
                name=name,
                type=schemas.DatabaseEntityType.table,
                dbschema=schema,
            )
            for name in names
        ]

        return result

    def inspect_views(self, schema: str) -> List[schemas.DatabaseEntity]:
        """
        Inspect a schema's views.
        """
        if not self.connected:
            raise Exception("Not connected to a database.")

        if not self.inspector:
            raise Exception("Database inspector is not initialized.")

        names = self.inspector.get_view_names(schema=schema)

        result = [
            schemas.DatabaseEntity(
                name=name,
                type=schemas.DatabaseEntityType.view,
                dbschema=schema,
            )
            for name in names
        ]

        return result

    def inspect_table(
        self, schema: str, table_name: str
    ) -> List[schemas.DatabaseEntity]:
        """
        Inspect a table.
        """
        if not self.connected:
            raise Exception("Not connected to a database.")

        if not self.inspector:
            raise Exception("Database inspector is not initialized.")

        columns = self.inspector.get_columns(
            table_name=table_name, schema=schema
        )

        result = []
        for column in columns:
            result.append(
                schemas.DatabaseEntity(
                    name=column["name"],
                    type=schemas.DatabaseEntityType.column,
                    dbschema=schema,
                    datatype=str(column["type"] or ""),
                    nullable=column["nullable"] or True,
                )
            )

        return result

    def inspect_view(
        self, schema: str, view_name: str
    ) -> List[schemas.DatabaseEntity]:
        """
        Inspect a view.
        """
        if not self.connected:
            raise Exception("Not connected to a database.")

        if not self.inspector:
            raise Exception("Database inspector is not initialized.")

        columns = self.inspector.get_columns(
            table_name=view_name, schema=schema
        )

        result = []
        for column in columns:
            result.append(
                schemas.DatabaseEntity(
                    name=column["name"],
                    type=schemas.DatabaseEntityType.column,
                    dbschema=schema,
                    datatype=str(column["type"] or ""),
                    nullable=column["nullable"] or True,
                )
            )

        return result

    def fetch_preview(
        self, schema: str, name: str
    ) -> List[schemas.DatabaseEntity]:
        """
        Fetch a preview of a table's data.
        """
        if not self.connected:
            raise Exception("Not connected to a database.")

        if not self.connection:
            raise Exception("Database connection is not initialized.")

        metadata = sqlalchemy.MetaData()
        table = sqlalchemy.Table(
            name,
            metadata,
            autoload_with=self.engine,
            schema=schema,
        )

        if self.config.max_fetch_rows > 0:
            query = table.select().limit(self.config.max_fetch_rows)
        else:
            query = table.select()

        result = self.connection.execute(query)

        rows = [
            {
                k: ("(bytes)" if isinstance(v, bytes) else v)
                for k, v in row.items()
            }
            for row in result.mappings()
        ]

        if not rows:
            return []

        data = [
            schemas.DatabaseEntity(
                name=f"row {index}",
                type=schemas.DatabaseEntityType.row,
                data=row,
            )
            for index, row in enumerate(rows)
        ]

        return data
