"""Metric-fragment qualification.

Metric expressions in `metrics.yaml` are written *unaliased* — e.g.
`AVG(CASE WHEN status = 'B' THEN 1.0 WHEN status = 'A' THEN 0.0 END)`, `A13`,
`SUM(amount)`. Before such a fragment can sit in a multi-table query it has to
be bound to the metric's own table alias, so the composer can introduce other
tables via joins without ambiguity.

We parse the fragment with sqlglot and inject the alias as the table qualifier
on every bare column reference, returning the resulting expression node (not a
string) so it splices directly into the query tree the composer assembles.
"""

from __future__ import annotations

import sqlglot
from sqlglot import exp


def qualify_metric_expr(sql_fragment: str, table_alias: str) -> exp.Expression:
    """Parse `sql_fragment` and qualify its bare column refs to `table_alias`.

    Used for both the metric's SELECT expression and its optional WHERE filter
    (both are written as bare expressions over the metric's table)."""
    tree = sqlglot.parse_one(sql_fragment, dialect="postgres")
    for col in tree.find_all(exp.Column):
        if col.table:
            continue
        col.set("table", exp.Identifier(this=table_alias, quoted=True))
    return tree