import hashlib
from collections import deque
from typing import Callable, Iterable, List, Set

from schema import Graph
from statement import Statement


class CorrelationGraph(Graph):
    def __init__(
        self,
        statements: List[Statement],
        guss_table_by_column: Callable[[str, List[str]], str],
    ):
        self.correlations: dict[str, Set[str]] = {}

        for statement in statements:
            for left_key, right_key in statement.correlations():
                # Process left_key
                left_is_subquery_alias = False
                if "." not in left_key:
                    table_name = guss_table_by_column(left_key, statement.table_names)
                    left_key = table_name + "." + left_key
                else:
                    # Check if table prefix is a real table or subquery alias
                    parts = left_key.split(".", 1)
                    if len(parts) == 2:
                        table_prefix, col_name = parts
                        # table_names may have schema prefix (e.g., "test.orders")
                        is_real_table = any(
                            tn == table_prefix or tn.endswith("." + table_prefix)
                            for tn in statement.table_names
                        )
                        if not is_real_table:
                            # Table prefix is a subquery alias, skip this correlation
                            left_is_subquery_alias = True

                # Process right_key
                right_is_subquery_alias = False
                if "." not in right_key:
                    table_name = guss_table_by_column(right_key, statement.table_names)
                    right_key = table_name + "." + right_key
                else:
                    # Check if table prefix is a real table or subquery alias
                    parts = right_key.split(".", 1)
                    if len(parts) == 2:
                        table_prefix, col_name = parts
                        # table_names may have schema prefix (e.g., "test.orders")
                        is_real_table = any(
                            tn == table_prefix or tn.endswith("." + table_prefix)
                            for tn in statement.table_names
                        )
                        if not is_real_table:
                            # Table prefix is a subquery alias, skip this correlation
                            right_is_subquery_alias = True

                # Skip correlations involving subquery aliases
                if left_is_subquery_alias or right_is_subquery_alias:
                    continue

                left_key = statement.schema_name + "." + left_key
                right_key = statement.schema_name + "." + right_key
                self.set_correlation_pair(left_key, right_key)

    def set_correlation_pair(self, left_key: str, right_key: str):
        if left_key not in self.correlations:
            self.correlations[left_key] = set([left_key])
        if right_key not in self.correlations:
            self.correlations[right_key] = set([right_key])
        self.correlations[left_key].add(right_key)
        self.correlations[right_key].add(left_key)

    def get_column_identity(self, key: str) -> str:
        m = hashlib.md5()
        if key not in self.correlations:
            m.update(key.encode())
            return m.hexdigest()
        for k in sorted(self.correlations[key]):
            m.update(k.encode())
        return m.hexdigest()

    def get_related_columns(self, key: str) -> Set[str]:
        if key not in self.correlations:
            return {key}
        return set(self.correlations[key])

    def iter_clusters(self) -> Iterable[Set[str]]:
        visited: Set[str] = set()
        for root in self.correlations.keys():
            if root in visited:
                continue
            cluster = set()
            queue = deque([root])
            while queue:
                current = queue.popleft()
                if current in visited:
                    continue
                visited.add(current)
                cluster.add(current)
                for neighbour in self.correlations.get(current, []):
                    if neighbour not in visited:
                        queue.append(neighbour)
            if cluster:
                yield cluster