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

SCD Strategy: {{ scd_strategy }}
Business Key: {{ business_key | join(', ') }}
{%- if scd_strategy == 'scd2' %}
Tracked Fields: {{ tracked_fields | join(', ') }}
{%- endif %}

Source: {{ source_contract }}

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

from __future__ import annotations

import hashlib
from datetime import datetime, timezone, date
from typing import Optional

import polars as pl


def run(
    good_df,
    *,
    contract=None,
    engine: str = "polars",
{%- if existing_dim_param %}
    {{ existing_dim_param }}: str = "",
{%- endif %}
    **_kwargs,
):
    """
    LakeLogic external_logic entrypoint — Dimension builder.

    Parameters
    ----------
    good_df : DataFrame
        {{ source_contract }} records that passed quality rules.
{%- if existing_dim_param %}
    {{ existing_dim_param }} : str
        Path/table for existing dimension (used for SCD merge).
{%- endif %}
    """
    df = _to_polars(good_df)

{%- if renames %}

    # ── Rename columns ────────────────────────────────────────────────────
    df = df.rename({
{%- for old, new in renames.items() %}
        "{{ old }}": "{{ new }}",
{%- endfor %}
    })
{%- endif %}

{%- if dedup_key %}

    # ── Deduplicate on business key ───────────────────────────────────────
    df = (
        df
        .sort("{{ dedup_timestamp }}", descending=True)
        .unique(subset={{ business_key | tojson }}, keep="first")
    )
{%- endif %}

{%- if derived_columns %}

    # ── Derive dimension attributes ───────────────────────────────────────
{%- for d in derived_columns %}
    df = df.with_columns({{ d.expression }}.alias("{{ d.name }}"))
{%- endfor %}
{%- endif %}

{%- if scd_strategy == 'scd1' %}

    # ── SCD Type 1 — Overwrite (latest values only) ──────────────────────
    # Select final dimension columns
    result = df.select([
{%- for col in output_columns %}
        "{{ col.name }}",
{%- endfor %}
    ])

    # Add audit columns
    result = result.with_columns(
        pl.lit(datetime.now(timezone.utc).replace(microsecond=0)).alias("dim_updated_at"),
    )

{%- elif scd_strategy == 'scd2' %}

    # ── SCD Type 2 — Full History Tracking ────────────────────────────────
    now = datetime.now(timezone.utc).replace(microsecond=0)
    end_of_time = datetime(9999, 12, 31, tzinfo=timezone.utc)

    # Hash tracked fields for change detection
    tracked = {{ tracked_fields | tojson }}

    df = df.with_columns(
        pl.concat_str([pl.col(c).cast(pl.Utf8).fill_null("__NULL__") for c in tracked], separator="|")
        .map_elements(lambda s: hashlib.md5(s.encode()).hexdigest(), return_dtype=pl.Utf8)
        .alias("_row_hash"),
    )

{%- if existing_dim_param %}

    # Load existing dimension for merge
    existing = _try_load_table({{ existing_dim_param }}, engine)

    if existing is not None:
        # Find changed / new records
        current = existing.filter(pl.col("is_current") == True)
        merged = df.join(
            current.select({{ business_key | tojson }} + ["_row_hash"]).rename({"_row_hash": "_existing_hash"}),
            on={{ business_key | tojson }},
            how="left",
        )
        changed = merged.filter(
            pl.col("_existing_hash").is_null() | (pl.col("_row_hash") != pl.col("_existing_hash"))
        ).drop("_existing_hash")

        # Expire old records for changed keys
        expired_keys = changed.select({{ business_key | tojson }}).unique()
        expire_updates = (
            existing
            .join(expired_keys, on={{ business_key | tojson }}, how="inner")
            .filter(pl.col("is_current") == True)
            .with_columns(
                pl.lit(now).alias("valid_to"),
                pl.lit(False).alias("is_current"),
            )
        )

        # Unchanged existing records stay as-is
        unchanged = existing.join(expired_keys, on={{ business_key | tojson }}, how="anti")

        # New version rows
        new_versions = changed.with_columns(
            pl.lit(now).alias("valid_from"),
            pl.lit(end_of_time).alias("valid_to"),
            pl.lit(True).alias("is_current"),
        )

        result = pl.concat([unchanged, expire_updates, new_versions], how="diagonal")
    else:
{%- endif %}
        # Initial load — all records are current
        result = df.with_columns(
            pl.lit(now).alias("valid_from"),
            pl.lit(end_of_time).alias("valid_to"),
            pl.lit(True).alias("is_current"),
        )

    # Select final columns
    result = result.select([
{%- for col in output_columns %}
        "{{ col.name }}",
{%- endfor %}
        "valid_from",
        "valid_to",
        "is_current",
        "_row_hash",
    ])

    # Add audit column
    result = result.with_columns(
        pl.lit(now).alias("dim_processed_at"),
    )

{%- endif %}

    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 _try_load_table(table_name: str, engine: str) -> Optional[pl.DataFrame]:
    """Try to load a table. Returns None if unavailable."""
    if not table_name:
        return None
    try:
        if engine == "spark":
            from pyspark.sql import SparkSession
            spark = SparkSession.builder.getOrCreate()
            return pl.from_pandas(spark.table(table_name).toPandas())
        else:
            return pl.read_delta(table_name.replace("table:", ""))
    except Exception:
        return None


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
