"""Ibis code generator for fyrnheim entities.

Generates Python modules containing plain Ibis functions (no Hamilton)
that transform entity data through source, prep, dimension, and snapshot layers.
"""

from __future__ import annotations

from pathlib import Path

from fyrnheim.core.entity import Entity
from fyrnheim.core.source import (
    AggregationSource,
    DerivedSource,
    EventAggregationSource,
    TableSource,
    UnionSource,
)
from fyrnheim.core.source_mapping import SourceMapping

IBIS_TYPE_MAP: dict[str, str] = {
    "timestamp": "timestamp",
    "timestamptz": "timestamp",
    "boolean": "boolean",
    "integer": "int64",
    "bigint": "int64",
    "int64": "int64",
    "float": "float64",
    "double": "float64",
    "float64": "float64",
    "decimal": "decimal",
    "numeric": "decimal",
    "string": "string",
    "text": "string",
    "date": "date",
}


class IbisCodeGenerator:
    """Generate Ibis Python modules from entity definitions."""

    def __init__(
        self,
        entity: Entity,
        source_mapping: SourceMapping | None = None,
    ) -> None:
        self.entity = entity
        self.entity_name = entity.name
        self.source_mapping = source_mapping

    def _generate_imports(self) -> str:
        """Generate module header with imports (no Hamilton)."""
        return f'''"""{self.entity_name} entity transformations.

Auto-generated by fyrnheim.generate() -- DO NOT EDIT.
Source entity: {self.entity_name}
"""

import os

import ibis
'''

    def _bind_expression(self, expr: str) -> str:
        """Bind column references to table variable 't'.

        If the expression already references t. or starts with ibis. or (,
        pass it through unchanged. Otherwise prefix with t.
        """
        if "t." in expr:
            return expr
        if expr.startswith("ibis.") or expr.startswith("("):
            return expr
        return f"t.{expr}"

    def _generate_source_functions(self) -> str:
        """Generate source reader functions, dispatching by source type."""
        source = self.entity.source
        if source is None:
            return ""

        if isinstance(source, UnionSource):
            return self._generate_union_source_functions(source)
        elif isinstance(source, (TableSource, EventAggregationSource)):
            return self._generate_single_source_function(source)
        elif isinstance(source, AggregationSource):
            return self._generate_aggregation_source_function(source)
        elif isinstance(source, DerivedSource):
            if source.identity_graph_config is not None:
                return self._generate_derived_source_function(source)
            return ""
        else:
            return ""

    def _build_rename_suffix(self) -> str:
        """Build .rename() suffix from source_mapping.field_mappings, or empty string."""
        if not self.source_mapping or not self.source_mapping.field_mappings:
            return ""
        return f".rename({self.source_mapping.field_mappings!r})"

    def _generate_single_source_function(
        self, source: TableSource | EventAggregationSource
    ) -> str:
        """Generate a single source function with backend switching."""
        name = self.entity_name
        rename = self._build_rename_suffix()

        # Build DuckDB path
        if not source.duckdb_path:
            raise ValueError(
                f"TableSource '{source.table}' requires duckdb_path for code generation"
            )
        duckdb_path = source.duckdb_path

        func = f'''
def source_{name}(conn: ibis.BaseBackend, backend: str) -> ibis.Table:
    """Read {name} from configured source."""
    if backend == "duckdb":
        try:
            raw = conn.table("source_{name}")
        except Exception:
            parquet_path = os.path.expanduser("{duckdb_path}")
            raw = conn.read_parquet(parquet_path)
        return raw{rename}
    elif backend == "bigquery":
        return conn.table("{source.table}", database=("{source.project}", "{source.dataset}")){rename}
    else:
        raise ValueError(f"Unsupported backend: {{backend}}")
'''
        return func

    def _generate_aggregation_source_function(self, source: AggregationSource) -> str:
        """Generate source function for aggregation from another entity."""
        name = self.entity_name
        func = f'''
def source_{name}(source_{source.source_entity}: ibis.Table) -> ibis.Table:
    """Aggregate {name} from {source.source_entity}."""
    t = source_{source.source_entity}
'''
        if source.filter_expression:
            func += f"    t = t.filter({source.filter_expression})\n"

        if source.aggregations:
            agg_parts = []
            for agg in source.aggregations:
                expr = self._bind_expression(agg.expression)
                agg_parts.append(f"        {agg.name}={expr},")
            agg_body = "\n".join(agg_parts)
            func += f'    return t.group_by("{source.group_by_column}").aggregate(\n{agg_body}\n    )\n'
        else:
            func += f'    return t.group_by("{source.group_by_column}").aggregate()\n'

        return func

    def _generate_union_source_functions(self, source: UnionSource) -> str:
        """Generate per-source functions and a union aggregator."""
        name = self.entity_name
        parts = []

        # Generate a function for each sub-source
        sub_names = []
        for _i, sub_source in enumerate(source.sources):
            src_label = sub_source.table
            sub_name = f"source_{name}_{src_label}"
            sub_names.append(sub_name)

            if not sub_source.duckdb_path:
                raise ValueError(
                    f"TableSource '{sub_source.table}' requires duckdb_path for code generation"
                )
            duckdb_path = sub_source.duckdb_path

            # Build per-source normalization suffix
            suffix = ""
            if sub_source.field_mappings:
                # field_mappings: {source_col: unified_col}
                # Ibis .rename(): {new_name: old_name} i.e. {unified_col: source_col}
                rename_map = {v: k for k, v in sub_source.field_mappings.items()}
                suffix += f".rename({rename_map!r})"
            if sub_source.literal_columns:
                literal_args = ", ".join(
                    f"{col}=ibis.literal({val!r})"
                    for col, val in sub_source.literal_columns.items()
                )
                suffix += f".mutate({literal_args})"

            func = f'''
def {sub_name}(conn: ibis.BaseBackend, backend: str) -> ibis.Table:
    """Read {src_label} for {name} union."""
    if backend == "duckdb":
        try:
            raw = conn.table("{sub_name}")
        except Exception:
            parquet_path = os.path.expanduser("{duckdb_path}")
            raw = conn.read_parquet(parquet_path)
        return raw{suffix}
    elif backend == "bigquery":
        return conn.table("{sub_source.table}", database=("{sub_source.project}", "{sub_source.dataset}")){suffix}
    else:
        raise ValueError(f"Unsupported backend: {{backend}}")
'''
            parts.append(func)

        # Generate union aggregator with schema normalization
        call_lines = ",\n        ".join(f"{sn}(conn, backend)" for sn in sub_names)
        union_func = f'''
def source_{name}(conn: ibis.BaseBackend, backend: str) -> ibis.Table:
    """Union all {name} sources."""
    parts = [
        {call_lines},
    ]
    # Normalize schemas: add typed null columns for any missing fields before union
    all_columns: dict[str, object] = {{}}
    for p in parts:
        for col in p.columns:
            if col not in all_columns:
                all_columns[col] = p[col].type()
    normalized = []
    for p in parts:
        missing = set(all_columns) - set(p.columns)
        if missing:
            p = p.mutate(**{{col: ibis.literal(None).cast(all_columns[col]) for col in sorted(missing)}})
        normalized.append(p.select(*sorted(all_columns)))
    return ibis.union(*normalized)
'''
        parts.append(union_func)

        return "\n".join(parts)

    def _generate_derived_source_function(self, source: DerivedSource) -> str:
        """Generate multi-input source function for identity graph join."""
        config = source.identity_graph_config
        assert config is not None, "DerivedSource requires identity_graph_config"
        name = self.entity_name
        ig_sources = config.sources
        priority = config.priority
        match_key = config.match_key

        # Identify shared fields (unified field name appears in 2+ sources)
        field_to_sources: dict[str, list[str]] = {}
        for src in ig_sources:
            for unified_name in src.fields:
                field_to_sources.setdefault(unified_name, []).append(src.name)
        shared_fields = sorted(f for f, srcs in field_to_sources.items() if len(srcs) > 1)

        # Unique fields (appear in exactly 1 source)
        unique_fields: dict[str, str] = {}
        for src in ig_sources:
            for unified_name in src.fields:
                if unified_name not in shared_fields:
                    unique_fields[unified_name] = src.name

        lines: list[str] = []
        lines.append(f"\ndef source_{name}(sources: dict) -> ibis.Table:")
        lines.append(f'    """Merge {name} from identity graph sources."""')

        # Extract tables
        for src in ig_sources:
            lines.append(f'    t_{src.name} = sources["{src.name}"]')
        lines.append("")

        # Rename columns to unified names
        for src in ig_sources:
            rename_map: dict[str, str] = {}
            if src.match_key_field != match_key:
                rename_map[match_key] = src.match_key_field
            for unified, source_col in src.fields.items():
                if source_col != unified:
                    rename_map[unified] = source_col
            if rename_map:
                lines.append(f"    t_{src.name} = t_{src.name}.rename({rename_map!r})")
        lines.append("")

        # Pre-join mutate: auto columns
        for src in ig_sources:
            parts_list: list[str] = []
            parts_list.append(f"_{src.name}_match_key=t_{src.name}.{match_key}")
            if src.id_field:
                parts_list.append(f"{src.name}_id=t_{src.name}.{src.id_field}")
            if src.date_field:
                parts_list.append(f"first_seen_{src.name}=t_{src.name}.{src.date_field}")
            args = ", ".join(parts_list)
            lines.append(f"    t_{src.name} = t_{src.name}.mutate({args})")
        lines.append("")

        # Cascading FULL OUTER JOIN with .select() after each
        for i in range(1, len(ig_sources)):
            left = f"t_{ig_sources[0].name}" if i == 1 else "result"
            right = f"t_{ig_sources[i].name}"

            lines.append(f"    result = {left}.outer_join(")
            lines.append(f"        {right}, {left}.{match_key} == {right}.{match_key},")
            lines.append('        lname="", rname="_right",')
            lines.append("    ).select(")

            # Build the select list
            select_items: list[str] = []

            # Coalesced match key
            select_items.append(
                f"        {match_key}=ibis.coalesce({left}.{match_key}, {right}.{match_key})"
            )

            # Shared fields - per-source variants
            for field_name in shared_fields:
                for j in range(i + 1):
                    src = ig_sources[j]
                    if field_name in src.fields:
                        col_name = f"{field_name}_{src.name}"
                        if j < i and i > 1:
                            # Already in result from prior join
                            select_items.append(f"        {col_name}=result.{col_name}")
                        elif j < i and i == 1:
                            # First join, left side is original table
                            select_items.append(
                                f"        {col_name}=t_{ig_sources[0].name}.{field_name}"
                            )
                        else:
                            # Right side of current join
                            select_items.append(f"        {col_name}={right}.{field_name}")

            # Unique fields
            for field_name, src_name in sorted(unique_fields.items()):
                src_idx = next(
                    idx for idx, s in enumerate(ig_sources) if s.name == src_name
                )
                if src_idx < i:
                    ref = f"t_{ig_sources[0].name}" if i == 1 else "result"
                    select_items.append(f"        {field_name}={ref}.{field_name}")
                elif src_idx == i:
                    select_items.append(f"        {field_name}={right}.{field_name}")
                # else: not yet joined, skip

            # Auto columns (match key copies, IDs, dates)
            for j in range(i + 1):
                src = ig_sources[j]
                if j < i:
                    ref = f"t_{ig_sources[0].name}" if (j == 0 and i == 1) else "result"
                else:
                    ref = right
                mk_col = f"_{src.name}_match_key"
                select_items.append(f"        {mk_col}={ref}.{mk_col}")
                if src.id_field:
                    id_col = f"{src.name}_id"
                    select_items.append(f"        {id_col}={ref}.{id_col}")
                if src.date_field:
                    date_col = f"first_seen_{src.name}"
                    select_items.append(f"        {date_col}={ref}.{date_col}")

            lines.append(",\n".join(select_items) + ",")
            lines.append("    )")

            # Cache intermediate results to break Ibis expression chain
            # (avoids IntegrityError on cascading joins with 3+ sources)
            if i < len(ig_sources) - 1:
                lines.append("    result = result.cache()")
        lines.append("")

        # PriorityCoalesce
        coalesce_parts: list[str] = []
        for field_name in shared_fields:
            sources_with_field = [
                s
                for s in priority
                if any(
                    field_name in src.fields
                    for src in ig_sources
                    if src.name == s
                )
            ]
            if len(sources_with_field) >= 2:
                chain = f"result.{field_name}_{sources_with_field[0]}"
                for s in sources_with_field[1:]:
                    chain += f".fill_null(result.{field_name}_{s})"
                coalesce_parts.append(f"        {field_name}={chain}")
        if coalesce_parts:
            lines.append("    result = result.mutate(")
            lines.append(",\n".join(coalesce_parts) + ",")
            lines.append("    )")
        lines.append("")

        # Source flags
        flag_parts = [
            f"        is_{src.name}=result._{src.name}_match_key.notnull()"
            for src in ig_sources
        ]
        lines.append("    result = result.mutate(")
        lines.append(",\n".join(flag_parts) + ",")
        lines.append("    )")
        lines.append("")

        # Drop intermediate columns
        drop_list: list[str] = []
        for field_name in shared_fields:
            for src in ig_sources:
                if field_name in src.fields:
                    drop_list.append(f'"{field_name}_{src.name}"')
        for src in ig_sources:
            drop_list.append(f'"_{src.name}_match_key"')
        lines.append(f"    result = result.drop({', '.join(drop_list)})")
        lines.append("")
        lines.append("    return result")

        return "\n".join(lines)

    def generate_module(self) -> str:
        """Generate complete Python module source code."""
        parts = [
            self._generate_imports(),
            self._generate_source_functions(),
        ]

        if self.entity.layers.prep:
            parts.append(self._generate_prep_function())

        if self.entity.layers.dimension:
            parts.append(self._generate_dimension_function())

        if self.entity.layers.snapshot and self.entity.layers.snapshot.enabled:
            parts.append(self._generate_snapshot_function())

        if self.entity.layers.activity:
            parts.append(self._generate_activity_function())

        if self.entity.layers.analytics:
            parts.append(self._generate_analytics_function())

        return "\n\n".join(p for p in parts if p)

    def _map_ibis_type(self, type_str: str) -> str:
        """Map generic type names to Ibis type strings."""
        return IBIS_TYPE_MAP.get(type_str.lower(), type_str.lower())

    def _generate_source_transforms(self) -> str:
        """Generate Ibis code for SourceTransforms (renames, casts, divides, multiplies)."""
        source = self.entity.source
        transforms = getattr(source, "transforms", None)
        if not transforms:
            return ""

        lines = []

        # 1. Renames
        if transforms.renames:
            rename_args = ", ".join(
                f'{r.to_name}="{r.from_name}"' for r in transforms.renames
            )
            lines.append(f"    t = t.rename({rename_args})")

        # 2. Type casts
        if transforms.type_casts:
            for tc in transforms.type_casts:
                ibis_type = self._map_ibis_type(tc.target_type)
                lines.append(f'    t = t.mutate({tc.field}=t.{tc.field}.cast("{ibis_type}"))')

        # 3. Divides
        if transforms.divides:
            for d in transforms.divides:
                new_name = f"{d.field}{d.suffix}"
                ibis_type = self._map_ibis_type(d.target_type)
                lines.append(
                    f'    t = t.mutate({new_name}=(t.{d.field} / {d.divisor}).cast("{ibis_type}"))'
                )

        # 4. Multiplies
        if transforms.multiplies:
            for m in transforms.multiplies:
                new_name = f"{m.field}{m.suffix}"
                ibis_type = self._map_ibis_type(m.target_type)
                lines.append(
                    f'    t = t.mutate({new_name}=(t.{m.field} * {m.multiplier}).cast("{ibis_type}"))'
                )

        return "\n".join(lines) + "\n" if lines else ""

    def _generate_prep_function(self) -> str:
        """Generate prep layer transformation function."""
        layer = self.entity.layers.prep
        computed_cols = layer.computed_columns if layer else []

        func = f"""
def prep_{self.entity_name}(source_{self.entity_name}: ibis.Table) -> ibis.Table:
    \"\"\"Prep layer transformations for {self.entity_name}.\"\"\"
    t = source_{self.entity_name}
"""

        # Source transforms (renames, casts, divides, multiplies)
        transforms_code = self._generate_source_transforms()
        if transforms_code:
            func += "\n" + transforms_code

        # Computed columns
        if computed_cols:
            func += "\n    # Computed columns\n"
            func += "    prep = t.mutate(\n"
            for i, col in enumerate(computed_cols):
                comma = "," if i < len(computed_cols) - 1 else ","
                expr = self._bind_expression(col.expression)
                if hasattr(col, "description") and col.description:
                    func += f"        # {col.description}\n"
                func += f"        {col.name}={expr}{comma}\n"
            func += "    )\n"
            func += "    return prep\n"
        else:
            func += "    return t\n"

        return func

    def _generate_dimension_function(self) -> str:
        """Generate dimension layer function."""
        layer = self.entity.layers.dimension

        # Merge core_computed + layer computed columns
        core_computed = list(self.entity.core_computed or [])
        layer_computed = list(layer.computed_columns or []) if layer else []
        computed_cols = core_computed + layer_computed

        # Determine input parameter name
        if self.entity.layers.prep:
            input_name = f"prep_{self.entity_name}"
        else:
            input_name = f"source_{self.entity_name}"

        func = f"""
def dim_{self.entity_name}({input_name}: ibis.Table) -> ibis.Table:
    \"\"\"Dimension layer for {self.entity_name}.\"\"\"
    t = {input_name}
"""

        if computed_cols:
            func += "\n    # Apply dimension computed columns\n"
            func += "    dim = t.mutate(\n"
            for i, col in enumerate(computed_cols):
                comma = "," if i < len(computed_cols) - 1 else ","
                expr = self._bind_expression(col.expression)
                if hasattr(col, "description") and col.description:
                    func += f"        # {col.description}\n"
                func += f"        {col.name}={expr}{comma}\n"
            func += "    )\n"
            func += "    return dim\n"
        else:
            func += "    return t\n"

        return func

    def _generate_snapshot_function(self) -> str:
        """Generate snapshot layer function that delegates to apply_snapshot()."""
        snapshot = self.entity.layers.snapshot
        assert snapshot is not None, "Entity must have a snapshot layer"
        date_col = snapshot.date_column

        # Parse deduplication ordering
        dedup_parts = snapshot.deduplication_order_by.strip().split()
        dedup_col = dedup_parts[0]
        dedup_desc = True
        if len(dedup_parts) > 1:
            dedup_desc = dedup_parts[1].upper() != "ASC"

        # Format natural_key for generated code
        nk = snapshot.natural_key
        if isinstance(nk, str):
            nk_repr = f'"{nk}"'
        else:
            nk_repr = repr(nk)

        include_validity = snapshot.include_validity_range

        # Determine input: prefer dim, then prep, then source
        if self.entity.layers.dimension:
            input_name = f"dim_{self.entity_name}"
        elif self.entity.layers.prep:
            input_name = f"prep_{self.entity_name}"
        else:
            input_name = f"source_{self.entity_name}"

        func = f"""
def snapshot_{self.entity_name}({input_name}: ibis.Table) -> ibis.Table:
    \"\"\"Snapshot layer: daily snapshot with {date_col} column.\"\"\"
    from fyrnheim.engine.snapshot import apply_snapshot

    return apply_snapshot(
        {input_name},
        natural_key={nk_repr},
        date_column="{date_col}",
        dedup_order_by="{dedup_col}",
        dedup_descending={dedup_desc},
        include_validity_range={include_validity},
    )
"""
        return func

    def _generate_analytics_function(self) -> str:
        """Generate analytics aggregation function."""
        analytics = self.entity.layers.analytics
        if analytics is None:
            return ""

        # Determine input: prefer dim, then prep, then source
        if self.entity.layers.dimension:
            input_name = f"dim_{self.entity_name}"
        elif self.entity.layers.prep:
            input_name = f"prep_{self.entity_name}"
        else:
            input_name = f"source_{self.entity_name}"

        date_expr = self._bind_expression(analytics.date_expression)

        # Build aggregation expressions
        agg_lines = []
        for m in analytics.metrics:
            expr = self._bind_expression(m.expression)
            agg_lines.append(f"        {m.name}={expr},")

        # Build group_by arguments
        group_parts = ['"_date"']
        for d in analytics.dimensions:
            group_parts.append(f'"{d}"')
        group_by = ", ".join(group_parts)

        func = f"""
def analytics_{self.entity_name}({input_name}: ibis.Table) -> ibis.Table:
    \"\"\"Analytics aggregation for {self.entity_name}.\"\"\"
    t = {input_name}
    t = t.mutate(_date={date_expr})
    return t.group_by({group_by}).aggregate(
{chr(10).join(agg_lines)}
    )
"""
        return func

    def _generate_activity_function(self) -> str:
        """Generate activity layer function that produces an event stream."""
        activity = self.entity.layers.activity
        if activity is None:
            return ""

        # Determine input: prefer dim, then prep, then source
        if self.entity.layers.dimension:
            input_name = f"dim_{self.entity_name}"
        elif self.entity.layers.prep:
            input_name = f"prep_{self.entity_name}"
        else:
            input_name = f"source_{self.entity_name}"

        entity_id = activity.entity_id_field
        identity = activity.identity_field

        # Build per-trigger SELECT parts
        parts_code = []
        for at in activity.types:
            if at.trigger == "row_appears":
                part = (
                    f"\n    # Activity: {at.name} (row_appears)\n"
                    f"    act_{at.name} = t.select(\n"
                    f'        entity_id=t.{entity_id}.cast("string"),\n'
                    f'        identity=t.{identity}.cast("string"),\n'
                    f"        ts=t.{at.timestamp_field},\n"
                    f'        activity_type=ibis.literal("{at.name}"),\n'
                    f"    )"
                )
                parts_code.append(part)
            elif at.trigger == "status_becomes":
                values_repr = repr(at.values)
                part = (
                    f"\n    # Activity: {at.name} (status_becomes on {at.field})\n"
                    f"    _filtered = t.filter(\n"
                    f"        t.{at.field}.isin({values_repr})\n"
                    f"    )\n"
                    f"    act_{at.name} = _filtered.select(\n"
                    f'        entity_id=_filtered.{entity_id}.cast("string"),\n'
                    f'        identity=_filtered.{identity}.cast("string"),\n'
                    f"        ts=_filtered.{at.timestamp_field},\n"
                    f'        activity_type=ibis.literal("{at.name}"),\n'
                    f"    )"
                )
                parts_code.append(part)
            elif at.trigger == "field_changes":
                part = (
                    f"\n    # Activity: {at.name} (field_changes on {at.field})\n"
                    f"    _prev = t.{at.field}.lag().over(\n"
                    f"        ibis.window(group_by=t.{entity_id}, order_by=t.{at.timestamp_field})\n"
                    f"    )\n"
                    f"    _changed = t.filter(\n"
                    f"        _prev.notnull() & (_prev != t.{at.field})\n"
                    f"    )\n"
                    f"    act_{at.name} = _changed.select(\n"
                    f'        entity_id=_changed.{entity_id}.cast("string"),\n'
                    f'        identity=_changed.{identity}.cast("string"),\n'
                    f"        ts=_changed.{at.timestamp_field},\n"
                    f'        activity_type=ibis.literal("{at.name}"),\n'
                    f"    )"
                )
                parts_code.append(part)

        activity_names = [f"act_{at.name}" for at in activity.types]
        if len(activity_names) == 1:
            union_line = f"    return {activity_names[0]}"
        else:
            union_args = ", ".join(activity_names)
            union_line = f"    return ibis.union({union_args})"

        func = f"""
def activity_{self.entity_name}({input_name}: ibis.Table) -> ibis.Table:
    \"\"\"Activity stream for {self.entity_name}.\"\"\"
    t = {input_name}
{"".join(parts_code)}

{union_line}
"""
        return func

    def write_module(self, output_dir: Path) -> Path:
        """Write generated module to disk. Returns path to written file."""
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        output_path = output_dir / f"{self.entity_name}_transforms.py"
        code = self.generate_module()
        output_path.write_text(code)
        return output_path
