from collections.abc import Sequence

from dagster_duckdb import DuckDBIOManager
from dagster_duckdb.io_manager import DbTypeHandler
from dagster_duckdb_pandas import DuckDBPandasTypeHandler
from dagster_polars import PolarsParquetIOManager

from dagster import Definitions

from .assets import deals, new_deals, works
from .jobs import check_partititions_job, deals_job, musicbrainz_lookup_job
from .schedules import deals_schedule
from .sensors import musicbrainz_lookup_sensor


class PandasDuckDBIOManager(DuckDBIOManager):
    @staticmethod
    def type_handlers() -> Sequence[DbTypeHandler]:
        return [DuckDBPandasTypeHandler()]


vinyl = Definitions(
    assets=[deals, new_deals, works],
    resources={
        "polars_parquet_io_manager": PolarsParquetIOManager(),
        "duckdb_io_manager": PandasDuckDBIOManager(database="vinyl"),
    },
    jobs=[deals_job, check_partititions_job, musicbrainz_lookup_job],
    schedules=[deals_schedule],
    sensors=[musicbrainz_lookup_sensor],
)