from typing import Iterator, List

from meta import Meta
from schema import Table
from statement import Statement


class Partition:
    def __init__(self, meta: Meta, table_ids: Iterator[str]):
        """
        table_ids is a list of `{db_name}.{tb_name}`
        """
        self.meta = meta
        self.table_ids = set(table_ids)

    def add_table(self, table_id: str) -> bool:
        self.table_ids.add(table_id)

    def add_tables(self, table_ids: Iterator[str]) -> bool:
        for t in table_ids:
            self.add_table(t)

    def has_table(self, table_id: str) -> bool:
        return table_id in self.table_ids

    def tables(self) -> List[Table]:
        xs = []
        for database in self.meta.databases():
            for table in database.tables:
                if self.has_table(f"{database.name}.{table.name}"):
                    xs.append(table)
        return xs

    def statements(self) -> List[Statement]:
        xs = []
        for statement in self.meta.statements():
            if len(statement.table_names) == 0:
                continue
            if self.has_table(statement.table_names[0]):
                xs.append(statement)
        return xs