import aiomysql
import pymysql
from dependency_injector import containers, providers

from .async_executor import AsyncExecutor
from .prepare import S3Backend, TiDBBackend
from .replayer import Replayer


def create_data_backend(cfg, connection_provider):
    backend_type = cfg["data"]["backend"]
    if backend_type == "tidb":
        return TiDBBackend(connection_provider())
    elif backend_type == "s3":
        return S3Backend(**cfg["data"]["s3"])
    else:
        raise ValueError(f"unknown data backend type: {backend_type}")


class ReplayerContainer(containers.DeclarativeContainer):
    config = providers.Configuration()

    meta = providers.Dependency()
    plan = providers.Dependency()

    tidb_connection = providers.Factory(
        pymysql.connect,
        host=config.tidb.host,
        port=config.tidb.port,
        user=config.tidb.user,
        password=config.tidb.password,
        ssl_ca=config.tidb.ssl_ca,
        read_timeout=600,
        write_timeout=600,
    )

    tidb_connection_async = providers.Factory(
        aiomysql.connect,
        host=config.tidb.host,
        port=config.tidb.port,
        user=config.tidb.user,
        password=config.tidb.password,
    )

    data_backend = providers.Singleton(
        create_data_backend,
        config,
        tidb_connection.provider,
    )

    executor = providers.Singleton(
        AsyncExecutor,
        concurrency=config.tidb.concurrency,
        async_connection_provider=tidb_connection_async.provider,
        async_queue_size=config.tidb.queue_size,
    )

    replayer = providers.Singleton(
        Replayer,
        partition_id=config.partition.id,
        partition_count=config.partition.count,
        meta=meta,
        plan=plan,
        executor=executor,
        data_backend=data_backend,
        state_config=config.state,
        prepare_batch_size=config.data.batch,
    )