"""Alias allocation for the composer.

Every table referenced in a composed query gets a short, query-unique alias.
`AliasMap` owns that allocation so the rest of the composer never threads a
bare dict around: ask it for a table's alias (allocating on first use) or for a
qualified `exp.Column`, and it remembers the assignment.

Allocation strategy (stable — golden SQL depends on it): first letter, then the
first two/three letters, then the full name, then a numbered `<first><n>` on
collision. So `loan→l`, `account→a`, `district→d`.
"""

from __future__ import annotations

from sqlglot import exp


def _alloc(table: str, used: set[str]) -> str:
    """Pick an alias for `table` not already in `used`."""
    for base in (table[:1], table[:2], table[:3], table):
        if base and base not in used:
            return base
    i = 1
    while True:
        cand = f"{table[:1]}{i}"
        if cand not in used:
            return cand
        i += 1


class AliasMap:
    """Maps table name → alias, allocating on first request."""

    def __init__(self) -> None:
        self._by_table: dict[str, str] = {}

    def has(self, table: str) -> bool:
        return table in self._by_table

    def alias(self, table: str) -> str:
        """Return `table`'s alias, allocating a fresh unique one if needed."""
        if table not in self._by_table:
            self._by_table[table] = _alloc(table, set(self._by_table.values()))
        return self._by_table[table]

    def col(self, table: str, column: str) -> exp.Column:
        """A quoted, alias-qualified column reference for the final tree."""
        return exp.column(column, table=self.alias(table), quoted=True)

    def aliased_table(self, table: str) -> exp.Expression:
        """`"<table>" AS "<alias>"` for FROM / JOIN."""
        return exp.alias_(exp.to_table(table, quoted=True), self.alias(table), quoted=True)

    def tables(self) -> list[str]:
        """Sorted names of every table that has been referenced."""
        return sorted(self._by_table)