"""Cost attribution helpers for the gharchive sample pipeline.

Tags warehouse queries with team / job metadata so per-team spend can be
reported from the warehouse's query-history metadata view.

Tag schema (convention):
    team=<name>;job=<name>;run_id=<uuid>

Examples:
    "team=analytics;job=daily_report;run_id=a1b2c3"
    "team=growth;job=top_actors;run_id=d4e5f6"

Warehouse support:
    Snowflake  — uses QUERY_HISTORY.QUERY_TAG; cost in credits.
    BigQuery   — uses INFORMATION_SCHEMA.JOBS.labels; cost via billed_bytes.
    DuckDB     — no native cost metadata; adapter records wall-clock seconds
                 in a local ``_query_costs`` table (see warehouse/adapter.py).
"""

from __future__ import annotations

import re
from typing import Any


def make_tag(*, team: str, job: str, run_id: str = "") -> str:
    """Return a canonical query tag string.

    Args:
        team:   Team name, e.g. ``"analytics"``.
        job:    Job name, e.g. ``"daily_report"``.
        run_id: Optional unique run identifier (UUID or similar).

    Returns:
        Tag string in the form ``team=<team>;job=<job>;run_id=<run_id>``.
    """
    parts = [f"team={team}", f"job={job}"]
    if run_id:
        parts.append(f"run_id={run_id}")
    return ";".join(parts)


def parse_tag(tag: str) -> dict[str, str]:
    """Parse a query-tag string into a dict of key=value pairs.

    Args:
        tag: Tag string, e.g. ``"team=analytics;job=daily_report"``.

    Returns:
        Dict, e.g. ``{"team": "analytics", "job": "daily_report"}``.
    """
    result: dict[str, str] = {}
    for part in tag.split(";"):
        part = part.strip()
        if "=" in part:
            k, _, v = part.partition("=")
            result[k.strip()] = v.strip()
    return result


def extract_team(tag: str) -> str:
    """Return the team name from a query tag, or ``"untagged"`` if absent."""
    return parse_tag(tag).get("team", "untagged")


# ---------------------------------------------------------------------------
# DuckDB fallback: aggregate cost from the _query_costs table recorded by
# warehouse/adapter.py when running in local (non-cloud) mode.
# ---------------------------------------------------------------------------

_DUCKDB_SUMMARY_SQL = """
SELECT
    regexp_extract(tag, 'team=([^;]+)', 1) AS team,
    COUNT(*)                               AS query_count,
    SUM(elapsed_s)                         AS total_seconds,
    NULL                                   AS total_credits,  -- N/A on DuckDB
    NULL                                   AS estimated_usd   -- N/A on DuckDB
FROM _query_costs
WHERE tag LIKE 'team=%'
GROUP BY 1
ORDER BY total_seconds DESC
"""


def duckdb_cost_summary(con: Any) -> Any:
    """Return per-team cost data from a DuckDB connection.

    The ``_query_costs`` table is created and populated by
    ``warehouse.adapter.query()`` automatically when running against DuckDB.

    Cost attribution requires a warehouse with per-query cost metadata.
    On DuckDB we fall back to wall-clock seconds because DuckDB has no
    query-cost view equivalent to Snowflake's QUERY_HISTORY.

    Args:
        con: An open ``duckdb.DuckDBPyConnection``.

    Returns:
        A ``duckdb.DuckDBPyRelation`` (call ``.df()`` for a pandas DataFrame).
    """
    return con.execute(_DUCKDB_SUMMARY_SQL)


# ---------------------------------------------------------------------------
# Snowflake variant (SQL only — executed in the Snowflake UI or via the
# Snowflake Python connector; not run as part of the local lab).
# See governance/cost_attribution.sql for the full, annotated version.
# ---------------------------------------------------------------------------

SNOWFLAKE_COST_SQL = """
-- Per-team query cost attribution via QUERY_HISTORY.
-- Works on Snowflake; DuckDB needs the fallback above.

WITH tagged AS (
  SELECT
    query_id,
    user_name,
    warehouse_name,
    query_tag,
    total_elapsed_time / 1000.0 AS elapsed_seconds,
    credits_used_cloud_services AS credits
  FROM SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY
  WHERE start_time >= DATEADD(DAY, -7, CURRENT_TIMESTAMP())
    AND query_tag LIKE 'team=%'
)
SELECT
  REGEXP_SUBSTR(query_tag, 'team=([^;]+)', 1, 1, 'e', 1) AS team,
  COUNT(*) AS query_count,
  SUM(elapsed_seconds) AS total_seconds,
  SUM(credits) AS total_credits,
  SUM(credits) * 3.00 AS estimated_usd  -- $3/credit illustrative
FROM tagged
GROUP BY 1
ORDER BY total_credits DESC
"""


def allocate_spend(rows: list[dict[str, Any]]) -> dict[str, float]:
    """Aggregate cost rows into a per-team spend map.

    Args:
        rows: List of dicts, each with a ``"team"`` key and a numeric
              ``"total_seconds"`` or ``"total_credits"`` key.

    Returns:
        ``{"analytics": 42.1, "growth": 17.3, ...}``
    """
    totals: dict[str, float] = {}
    for row in rows:
        team = row.get("team") or "untagged"
        cost = float(row.get("total_credits") or row.get("total_seconds") or 0)
        totals[team] = totals.get(team, 0.0) + cost
    return totals


if __name__ == "__main__":
    # Quick sanity-check (no side effects).
    tag = make_tag(team="analytics", job="daily_report", run_id="abc123")
    assert extract_team(tag) == "analytics", "make_tag / extract_team mismatch"
    parsed = parse_tag(tag)
    assert parsed["job"] == "daily_report", "parse_tag mismatch"
    print("cost_attribution.py: all assertions passed.")