import argparse
import sqlite3
from typing import List

import pandas as pd
import pymysql
from dependency_injector import containers, providers
from dependency_injector.wiring import Provide, inject
from sqlalchemy import create_engine
from tqdm import tqdm

from meta import MetaContainer


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

    tidb_connection = providers.Resource(
        pymysql.connect,
        host=config.tidb.host,
        port=config.tidb.port,
        user=config.tidb.user,
        password=config.tidb.password,
        ssl_ca=config.tidb.ssl_ca,
    )

    sqlite_connection = providers.Singleton(
        sqlite3.Connection,
        database=config.meta.uri,
    )

    meta = providers.Container(
        MetaContainer,
        meta_uri=config.meta.uri,
    ).meta


class Capture:
    @inject
    def analyze(
        self,
        tidb_connection_provider=Provide[CaptureContainer.tidb_connection.provider],
    ):
        tidb_connection = tidb_connection_provider()
        with tidb_connection.cursor() as cur:
            cur.execute(
                "SELECT DISTINCT table_schema, table_name FROM `simcluster`.`columns`"
            )
            for table_schema, table_name in tqdm(cur.fetchall(), desc="Analyze"):
                cur.execute(f"ANALYZE TABLE {table_schema}.{table_name} ALL COLUMNS")

    @inject
    def capture(
        self,
        tables: List[str],
        config=Provide[CaptureContainer.config],
        sqlite_connection=Provide[CaptureContainer.sqlite_connection],
    ):
        tidb_connection = create_engine(
            f"mysql+pymysql://{config['tidb']['user']}@{config['tidb']['host']}:{config['tidb']['port']}/simcluster"
        )
        for table in tqdm(tables, desc="Capture"):
            df = pd.read_sql(f"SELECT * FROM {table}", con=tidb_connection)
            df.to_sql(table, sqlite_connection, if_exists="replace", index=False)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Capture metadata for workload simulation."
    )
    subparsers = parser.add_subparsers(dest="command", help="available sub commands")
    subparsers.add_parser(
        "statistics", help="capture statistical information"
    ).add_argument(
        "--skip-analyze", action="store_true", help="skip analyze table before collect"
    )
    subparsers.add_parser("statements", help="capture statements history")

    args = parser.parse_args()

    CaptureContainer().wire(modules=[__name__])
    c = Capture()
    if args.command == "statistics":
        if not args.skip_analyze:
            c.analyze()
        c.capture(
            ["columns", "partitions", "tidb_indexes", "stats_buckets", "stats_top_n"]
        )
    elif args.command == "statements":
        c.capture(["cluster_statements_summary_history"])