import argparse
import json

from dependency_injector import containers, providers
from dependency_injector.wiring import Provide, inject

from meta import MetaContainer
from plan import PlanContainer
from replayer import ReplayerContainer
from samplers import SamplerContainer
from schema import SchemaContainer
from statement import StatementContainer


class ApplicationContainer(containers.DeclarativeContainer):
    config = providers.Configuration(yaml_files=["config.yaml"])

    sampler = providers.Container(SamplerContainer).sampler

    statement_container = providers.Container(StatementContainer)

    schema_container = providers.Container(
        SchemaContainer,
        sampler=sampler,
    )

    meta = providers.Container(
        MetaContainer,
        meta_uri=config.meta.uri,
        statement=statement_container.statement,
        database=schema_container.database,
    ).meta

    plan = providers.Container(
        PlanContainer,
        plan_type=config.partition.plan,
        meta=meta,
    ).plan

    replayer = providers.Container(
        ReplayerContainer,
        config=config,
        plan=plan,
        meta=meta,
    ).replayer


@inject
def bootstrap(
    command: str,
    r=Provide[ApplicationContainer.replayer],
    p=Provide[ApplicationContainer.plan],
    c=Provide[ApplicationContainer.config],
):
    if command == "prepare":
        r.rebuild_schemas()
        r.prepare_data()
    elif command == "replay":
        r.replay()
    elif command == "groups":
        print(json.dumps(p.groups(), indent=4))
    elif command == "partitions":
        ps = p.partitions(c["partition"]["count"])
        print(
            json.dumps(
                [[f"{t.schema_name()}.{t.name}" for t in p.tables()] for p in ps],
                indent=4,
            )
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Prepare & Replay Synthesized Workload."
    )
    subparsers = parser.add_subparsers(dest="command", help="available sub commands")
    subparsers.add_parser("groups", help="print table groups")
    subparsers.add_parser("partitions", help="print partitions")
    subparsers.add_parser("prepare", help="prepare synthesized data to data backend")
    subparsers.add_parser("replay", help="replay synthesized workload")

    args = parser.parse_args()
    ApplicationContainer().wire(modules=[__name__])
    bootstrap(args.command)