"""
Gold – {{ title }}
{{ '=' * (title | length + 7) }}
{{ description }}

Source: {{ source_contract }}
Dimensions: {{ dimensions | join(', ') }}
Measures:
{%- for m in measures %}
  - {{ m.alias }}: {{ m.func }}({{ m.field }})
{%- endfor %}

Called by lakelogic via external_logic in {{ contract_filename }}.
Generated by Zeus AI from template: gold_aggregate_processor.
"""

from __future__ import annotations

from datetime import datetime, timezone
from typing import Optional

import polars as pl


def run(
    good_df,
    *,
    contract=None,
    engine: str = "polars",
    **_kwargs,
):
    """
    LakeLogic external_logic entrypoint.

    Parameters
    ----------
    good_df : DataFrame
        {{ source_contract }} records that passed quality rules.
    """
    df = _to_polars(good_df)

{%- if filters %}

    # ── Apply filters ─────────────────────────────────────────────────────
{%- for f in filters %}
    df = df.filter({{ f.expression }})
{%- endfor %}
{%- endif %}

{%- if pre_derived %}

    # ── Pre-aggregation derived columns ───────────────────────────────────
{%- for d in pre_derived %}
    df = df.with_columns({{ d.expression }}.alias("{{ d.name }}"))
{%- endfor %}
{%- endif %}

    # ── Aggregate ─────────────────────────────────────────────────────────
    result = df.group_by({{ dimensions | tojson }}).agg(
{%- for m in measures %}
        pl.{{ m.func }}("{{ m.field }}").alias("{{ m.alias }}"),
{%- endfor %}
    )

{%- if post_derived %}

    # ── Post-aggregation derived columns ──────────────────────────────────
{%- for d in post_derived %}
    result = result.with_columns({{ d.expression }}.alias("{{ d.name }}"))
{%- endfor %}
{%- endif %}

{%- if sort_by %}

    # ── Sort ──────────────────────────────────────────────────────────────
    result = result.sort({{ sort_by | tojson }}, descending={{ sort_descending | default(false) | tojson }})
{%- endif %}

{%- if top_n %}

    # ── Top N ─────────────────────────────────────────────────────────────
    result = result.head({{ top_n }})
{%- endif %}

    # ── Audit column ──────────────────────────────────────────────────────
    result = result.with_columns(
        pl.lit(datetime.now(timezone.utc).replace(microsecond=0)).alias("gold_processed_at"),
    )

    return _maybe_to_spark(result, engine)


# ── Helpers ───────────────────────────────────────────────────────────────────


def _to_polars(df) -> pl.DataFrame:
    """Convert any DataFrame to Polars."""
    if isinstance(df, pl.DataFrame):
        return df
    if hasattr(df, "toPandas"):
        return pl.from_pandas(df.toPandas())
    return pl.from_pandas(df)


def _maybe_to_spark(df: pl.DataFrame, engine: str):
    """Convert back to Spark if needed."""
    if engine == "spark":
        try:
            from pyspark.sql import SparkSession
            spark = SparkSession.builder.getOrCreate()
            return spark.createDataFrame(df.to_pandas())
        except Exception:
            pass
    return df
