"""BLCE Wright Pipeline Builder — generate 4-object Wright pipeline SQL.

Produces VW_1 (staging), DT_2 (join), DT_3A (pre-agg), DT_3 (presentation)
SQL for each fact table in a ProposedModel.
"""
from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional

from .contracts import (
    GeneratedWrightPipeline,
    ProposedDimension,
    ProposedFact,
)

logger = logging.getLogger(__name__)


class WrightPipelineBuilder:
    """Generate Wright 4-object pipeline SQL from proposed facts."""

    def build_pipeline(
        self,
        fact: ProposedFact,
        dimensions: List[ProposedDimension],
        schema: str = "EDW",
    ) -> GeneratedWrightPipeline:
        """Build a complete 4-object pipeline for a fact table.

        Args:
            fact: Proposed fact table specification.
            dimensions: All proposed dimensions (for join SQL).
            schema: Target Snowflake schema.

        Returns:
            GeneratedWrightPipeline with VW_1, DT_2, DT_3A, DT_3 SQL.
        """
        base = fact.name.replace("FACT_", "")

        vw1 = self._generate_vw1(fact, dimensions, schema, base)
        dt2 = self._generate_dt2(fact, dimensions, schema, base)
        dt3a = self._generate_dt3a(fact, schema, base)
        dt3 = self._generate_dt3(fact, dimensions, schema, base)
        union_tpl = self._generate_union_template(fact, schema, base) if fact.needs_union else ""

        return GeneratedWrightPipeline(
            fact_name=fact.name,
            vw1_sql=vw1,
            dt2_sql=dt2,
            dt3a_sql=dt3a,
            dt3_sql=dt3,
            union_template=union_tpl,
        )

    # ------------------------------------------------------------------
    # VW_1: Staging view — SELECT with renames/casts
    # ------------------------------------------------------------------

    def _generate_vw1(
        self,
        fact: ProposedFact,
        dimensions: List[ProposedDimension],
        schema: str,
        base: str,
    ) -> str:
        """VW_1_*: staging SELECT with column renames and type casts."""
        src = fact.source_tables[0] if fact.source_tables else f"SOURCE_{base}"

        select_cols: List[str] = []

        # Grain columns
        for gc in fact.grain_columns:
            select_cols.append(f"    src.{gc}")

        # Measures
        for m in fact.measures:
            select_cols.append(f"    COALESCE(src.{m}, 0) AS {m}")

        if not select_cols:
            select_cols.append("    src.*")

        col_block = ",\n".join(select_cols)
        return (
            f"-- VW_1_{base}: Staging view\n"
            f"CREATE OR REPLACE VIEW {schema}.VW_1_{base} AS\n"
            f"SELECT\n"
            f"{col_block}\n"
            f"FROM {src} src\n"
            f";"
        )

    # ------------------------------------------------------------------
    # DT_2: Join to dimensions
    # ------------------------------------------------------------------

    def _generate_dt2(
        self,
        fact: ProposedFact,
        dimensions: List[ProposedDimension],
        schema: str,
        base: str,
    ) -> str:
        """DT_2_*: JOIN staging view to dimension tables."""
        joins: List[str] = []
        dim_cols: List[str] = []

        dim_map = {d.name.upper(): d for d in dimensions}

        for fk in fact.dimension_fks:
            fk_upper = fk.upper()
            # Find matching dimension
            for dim_name, dim in dim_map.items():
                dim_short = dim_name.replace("DIM_", "")
                if dim_short.lower() in fk_upper.lower() or fk_upper in dim_name:
                    nk = dim.natural_key or f"{dim_short}_ID"
                    sk = dim.key_column or f"{dim_name}_SK"
                    joins.append(
                        f"LEFT JOIN {schema}.{dim_name} {dim_short.lower()}\n"
                        f"    ON stg.{fk} = {dim_short.lower()}.{nk}"
                    )
                    dim_cols.append(f"    {dim_short.lower()}.{sk} AS {dim_short}_SK")
                    break

        stg_cols = []
        for gc in fact.grain_columns:
            stg_cols.append(f"    stg.{gc}")
        for m in fact.measures:
            stg_cols.append(f"    stg.{m}")

        all_cols = stg_cols + dim_cols
        if not all_cols:
            all_cols = ["    stg.*"]

        col_block = ",\n".join(all_cols)
        join_block = "\n".join(joins) if joins else "-- No dimension joins detected"

        return (
            f"-- DT_2_{base}: Dimension joins\n"
            f"CREATE OR REPLACE TABLE {schema}.DT_2_{base} AS\n"
            f"SELECT\n"
            f"{col_block}\n"
            f"FROM {schema}.VW_1_{base} stg\n"
            f"{join_block}\n"
            f";"
        )

    # ------------------------------------------------------------------
    # DT_3A: Pre-aggregation + UNION ALL
    # ------------------------------------------------------------------

    def _generate_dt3a(
        self,
        fact: ProposedFact,
        schema: str,
        base: str,
    ) -> str:
        """DT_3A_*: pre-aggregation step (optional UNION ALL)."""
        grain_cols = fact.grain_columns or ["-- grain columns TBD"]
        measure_aggs = [f"    SUM({m}) AS {m}" for m in fact.measures]

        grain_block = ",\n".join(f"    {gc}" for gc in fact.grain_columns) if fact.grain_columns else "    -- grain TBD"
        measure_block = ",\n".join(measure_aggs) if measure_aggs else "    -- measures TBD"
        group_block = ", ".join(str(i + 1) for i in range(len(fact.grain_columns)))

        sql = (
            f"-- DT_3A_{base}: Pre-aggregation\n"
            f"CREATE OR REPLACE TABLE {schema}.DT_3A_{base} AS\n"
            f"SELECT\n"
            f"{grain_block},\n"
            f"{measure_block}\n"
            f"FROM {schema}.DT_2_{base}\n"
        )

        if group_block:
            sql += f"GROUP BY {group_block}\n"

        sql += ";"
        return sql

    # ------------------------------------------------------------------
    # DT_3: Presentation layer
    # ------------------------------------------------------------------

    def _generate_dt3(
        self,
        fact: ProposedFact,
        dimensions: List[ProposedDimension],
        schema: str,
        base: str,
    ) -> str:
        """DT_3_*: presentation layer with surrogate key lookups."""
        select_cols: List[str] = []

        for gc in fact.grain_columns:
            select_cols.append(f"    agg.{gc}")
        for m in fact.measures:
            select_cols.append(f"    agg.{m}")

        # Add metadata
        select_cols.append("    CURRENT_TIMESTAMP() AS LOADED_AT")
        select_cols.append("    'PIPELINE' AS SOURCE_SYSTEM")

        col_block = ",\n".join(select_cols) if select_cols else "    agg.*"

        return (
            f"-- DT_3_{base}: Presentation / final fact\n"
            f"CREATE OR REPLACE TABLE {schema}.DT_3_{base} AS\n"
            f"SELECT\n"
            f"{col_block}\n"
            f"FROM {schema}.DT_3A_{base} agg\n"
            f";"
        )

    # ------------------------------------------------------------------
    # UNION ALL template
    # ------------------------------------------------------------------

    def _generate_union_template(
        self,
        fact: ProposedFact,
        schema: str,
        base: str,
        source_columns: Optional[Dict[str, List[str]]] = None,
    ) -> str:
        """Generate UNION ALL template for multi-source facts.

        Collects the superset of columns across all union sources and pads
        missing columns with ``NULL AS col_name`` so every SELECT has the
        same column list.

        Args:
            fact: Proposed fact with union_sources.
            schema: Target schema.
            base: Fact base name (without FACT_ prefix).
            source_columns: Optional {table_name: [col, ...]} for each source.
                            When absent, falls back to fact grain + measures.
        """
        if not fact.union_sources:
            return ""

        # Determine columns per source
        per_source: Dict[str, List[str]] = {}
        if source_columns:
            for src in fact.union_sources:
                per_source[src] = source_columns.get(
                    src, source_columns.get(src.upper(), [])
                )
        # Fallback: use grain + measures as universal column set
        if not any(per_source.values()):
            default_cols = fact.grain_columns + fact.measures
            for src in fact.union_sources:
                per_source[src] = default_cols

        # Build superset column list (preserving insertion order)
        all_cols: List[str] = []
        seen: set = set()
        for src in fact.union_sources:
            for col in per_source.get(src, []):
                col_upper = col.upper()
                if col_upper not in seen:
                    seen.add(col_upper)
                    all_cols.append(col)

        if not all_cols:
            all_cols = ["*"]

        # Build SELECT for each source with NULL padding for missing columns
        parts: List[str] = []
        for i, src in enumerate(fact.union_sources):
            src_cols_upper = {c.upper() for c in per_source.get(src, [])}
            select_items: List[str] = []
            for col in all_cols:
                if col.upper() in src_cols_upper:
                    select_items.append(col)
                else:
                    select_items.append(f"NULL AS {col}")
            select_items.append(f"'{src}' AS SOURCE_TABLE")

            col_block = ", ".join(select_items)
            union_kw = "UNION ALL\n" if i > 0 else ""
            prefix = "    " if i > 0 else ""
            parts.append(f"{union_kw}{prefix}SELECT {col_block}\n    FROM {src}")

        body = "\n".join(parts)
        return (
            f"-- UNION ALL template for {fact.name}\n"
            f"CREATE OR REPLACE VIEW {schema}.VW_UNION_{base} AS\n"
            f"{body}\n"
            f";"
        )

    def build_all(
        self,
        facts: List[ProposedFact],
        dimensions: List[ProposedDimension],
        schema: str = "EDW",
    ) -> List[GeneratedWrightPipeline]:
        """Build pipelines for all facts in a model."""
        return [self.build_pipeline(f, dimensions, schema) for f in facts]
