import sqlite3
from functools import lru_cache
from typing import Callable, List, Tuple

import pandas as pd

from schema import Database
from statement import Statement

from .graph import CorrelationGraph


class Meta:
    def __init__(
        self,
        s3_uri: str,
        statement_factory: Callable[..., Statement],
        database_factory: Callable[..., Database],
    ):
        self.s3_uri = s3_uri
        sqlite_connection = sqlite3.Connection(s3_uri)
        df = pd.read_sql("SELECT * FROM columns", con=sqlite_connection)

        @lru_cache(maxsize=1024)
        def _find_column_owner(column_name: str, candidates: Tuple[str]) -> str:
            f = df[
                (df.column_name == column_name)
                & ((df.table_schema + "." + df.table_name).isin(candidates))
            ]
            if len(f) == 0:
                return None
            return f.table_name.iloc[0]

        self.find_column_owner = lambda c, cs: _find_column_owner(c, tuple(cs))
        self._statements = self._load_statements(sqlite_connection, statement_factory)
        self.graph = CorrelationGraph(self.statements(), self.find_column_owner)
        self._databases = self._load_databases(sqlite_connection, database_factory)

    def statements(self) -> List[Statement]:
        return self._statements

    def databases(self) -> Database:
        return self._databases

    def _load_databases(self, sqlite_connection, database_factory) -> List[Database]:
        databases = []
        columns_df = pd.read_sql("SELECT * FROM columns", con=sqlite_connection)
        partitions_df = pd.read_sql("SELECT * FROM partitions", con=sqlite_connection)
        indexes_df = pd.read_sql("SELECT * FROM tidb_indexes", con=sqlite_connection)
        stats_buckets_df = pd.read_sql(
            "SELECT * FROM stats_buckets", con=sqlite_connection
        )
        stats_top_n_df = pd.read_sql("SELECT * FROM stats_top_n", con=sqlite_connection)
        for schema_name in columns_df.table_schema.unique():
            databases.append(
                database_factory(
                    name=schema_name,
                    dataframes={
                        "columns_df": columns_df[
                            columns_df.table_schema == schema_name
                        ],
                        "partitions_df": partitions_df[
                            partitions_df.table_schema == schema_name
                        ],
                        "indexes_df": indexes_df[
                            indexes_df.table_schema == schema_name
                        ],
                        "stats_buckets_df": stats_buckets_df[
                            stats_buckets_df.table_schema == schema_name
                        ],
                        "stats_top_n_df": stats_top_n_df[
                            stats_top_n_df.table_schema == schema_name
                        ],
                    },
                    graph=self.graph,
                )
            )
        return databases

    def _load_statements(self, sqlite_connection, statement_factory) -> List[Statement]:
        statements = []
        df = pd.read_sql(
            "SELECT * FROM cluster_statements_summary_history order by summary_begin_time",
            con=sqlite_connection,
        )
        df = df[(df.stmt_type.isin(["Select", "Update", "Insert", "Delete"]))]
        for statement in df.to_dict(orient="records"):
            statements.append(
                statement_factory(
                    digest_text=statement["digest_text"],
                    stmt_type=statement["stmt_type"],
                    schema_name=statement["schema_name"],
                    table_names=(
                        statement["table_names"].split(",")
                        if statement["table_names"]
                        else []
                    ),
                    exec_count=statement["exec_count"],
                    summary_begin_time=statement["summary_begin_time"],
                )
            )
        return statements

    def get_database(self, name: str) -> Database:
        for db in self.databases():
            if db.name == name:
                return db
        return None