import random
from typing import Iterator, List

from tqdm import tqdm

from meta import Meta
from plan import PartitionPlan
from samplers.sampling_phase import SamplingPhase
from samplers.state_manager import StateManager
from statement import MarkerType, Statement

from .async_executor import AsyncExecutor
from .prepare import DataBackend


class Replayer:
    def __init__(
        self,
        partition_id: int,
        partition_count: int,
        meta: Meta,
        plan: PartitionPlan,
        executor: AsyncExecutor,
        data_backend: DataBackend,
        state_config: dict,
        prepare_batch_size: int = 10000,
    ):
        partitions = plan.partitions(partition_count)
        self.partition = partitions[partition_id]
        self.meta = meta
        self.data_backend = data_backend
        self.state_manager = StateManager(state_config)
        self.prepare_batch_size = prepare_batch_size
        self.executor = executor

    def rebuild_schemas(self):
        self.data_backend.rebuild_schemas(self.partition.tables())

    def prepare_data(self):
        for table in self.partition.tables():
            column_names = [column.name for column in table.columns]

            values = []
            for i in tqdm(
                range(table.estimated_row_count),
                desc=f"preparing {table.schema_name()}.{table.name}",
                position=1,
                leave=False,
            ):
                row = table.sample_row(phase=SamplingPhase.PREPARE)
                values.append([row[cn] for cn in column_names])
                if (i + 1) % self.prepare_batch_size == 0 and len(values) > 0:
                    self.data_backend.insert_data(table, values)
                    values = []
            if len(values) > 0:
                self.data_backend.insert_data(table, values)

        # Save state after prepare is complete
        self.save_all_states()

    def sample_statement_from_batch(
        self, batch: List[Statement]
    ) -> Iterator[Statement]:
        weights = [s.exec_count for s in batch]

        while sum(weights) > 0:
            i = random.choices(range(len(batch)), weights=weights)[0]  # seed
            weights[i] -= 1
            yield batch[i]

    def replay(self):
        # Load state before replay starts
        assert self.load_all_states(), "load column states failed"

        statement_groups = {}
        for statement in self.partition.statements():
            if statement.summary_begin_time not in statement_groups:
                statement_groups[statement.summary_begin_time] = []
            statement_groups[statement.summary_begin_time].append(statement)

        for g in tqdm(
            [statement_groups[t] for t in sorted(statement_groups.keys())],
            desc="Replaying Queries",
            position=0,
        ):
            for statement in tqdm(
                self.sample_statement_from_batch(g),
                total=sum([s.exec_count for s in g]),
                desc=f"{g[0].summary_begin_time}",
                position=1,
                leave=False,
            ):
                if statement.stmt_type == "Insert":
                    assert (
                        len(statement.table_names) == 1
                    ), "insert statement should affect exact one table"
                    schema_name, table_name = statement.table_names[0].split(".")
                    db = self.meta.get_database(schema_name)
                    table = db.get_table(table_name)
                    row = table.sample_row(phase=SamplingPhase.REPLAY)
                    column_names = [column.name for column in table.columns]
                    values = [row[cn] for cn in column_names]
                    sql_template = f"""
                        insert into `{schema_name}`.`{table.name}` ({",".join(column_names)})
                        values ({",".join(['%s' for _ in range(len(column_names))])})
                    """
                    self.executor.execute(schema_name, sql_template, values)
                    continue

                ps = statement.placeholders()
                for i in range(len(ps)):
                    if "." not in ps[i][1]:
                        table_name = self.meta.find_column_owner(
                            ps[i][1], statement.table_names
                        )
                        ps[i] = (ps[i][0], ".".join([table_name, ps[i][1]]))

                db = self.meta.get_database(statement.schema_name)
                xs = []
                for p in ps:
                    (table_name, column_name) = p[1].split(".")
                    table = db.get_table(table_name)
                    if table is None:
                        break
                    c = table.lookup_column(column_name)
                    if p[0] == MarkerType.SELECT:
                        xs.append(c.sample_for_select())
                    elif p[0] == MarkerType.UPDATE:
                        xs.append(c.sample_for_update())
                    elif p[0] == MarkerType.INSERT:
                        xs.append(c.sample_for_insert(phase=SamplingPhase.REPLAY))
                    elif p[0] == MarkerType.DELETE:
                        xs.append(c.sample_for_delete())
                else:
                    self.executor.execute(
                        statement.schema_name, statement.fixed_statement(), xs
                    )
        self.executor.kill()

    def save_all_states(self) -> bool:
        """Save states for all tables in the partition"""
        success_count = 0
        total_count = len(self.partition.tables())

        for table in self.partition.tables():
            if table.save_all_column_states(self.state_manager):
                success_count += 1

        if success_count == total_count:
            return True
        else:
            # print(f"⚠️  Saved states for {success_count}/{total_count} tables")
            return False

    def load_all_states(self) -> bool:
        """Load states for all tables in the partition"""
        success_count = 0
        found_count = 0

        for table in self.partition.tables():
            if table.has_saved_states(self.state_manager):
                found_count += 1
                if table.load_all_column_states(self.state_manager):
                    success_count += 1

        if found_count == 0:
            # print("📄 No saved states found for any tables - starting fresh")
            return False
        elif success_count == found_count:
            return True
        else:
            # print(
            #     f"⚠️  Loaded states for {success_count}/{found_count} tables with saved states"
            # )
            return False