from typing import List

from meta import Meta

from .partition_plan import PartitionPlan


class GroupPartitionPlan(PartitionPlan):
    def __init__(self, meta: Meta):
        self.meta = meta

    def groups(self) -> List[List[str]]:
        correlations = {}
        statements = self.meta.statements()
        for s in statements:
            for tn in s.table_names:
                if tn not in correlations:
                    correlations[tn] = set()
                correlations[tn] |= set(s.table_names) - set([tn])
        groups = []
        for database in self.meta.databases():
            for table in database.tables:
                table_id = f"{database.name}.{table.name}"
                if table_id not in correlations:
                    groups.append([table_id])
                    continue
                for g in groups:
                    for t in g:
                        if t in correlations[table_id]:
                            g.append(table_id)
                            break
                    else:
                        continue
                    break
                else:
                    groups.append([table_id])
        return groups