"""Fact table harmonization — align columns across systems and generate UNION SQL."""
from __future__ import annotations

import logging
import re
from typing import Any, Dict, List, Optional, Set, Tuple

from .contracts import (
    AlignmentStrategy,
    FactAlignment,
    GrainComparison,
    GrainLevel,
    GrainRelationship,
    _new_id,
)
from .grain_analyzer import GrainAnalyzer

logger = logging.getLogger(__name__)

# Aggregate function defaults by measure-name heuristic
_AGG_HEURISTICS: List[Tuple[re.Pattern, str]] = [
    (re.compile(r"(count|cnt|num)", re.I), "SUM"),
    (re.compile(r"(avg|average|mean)", re.I), "AVG"),
    (re.compile(r"(max|maximum|ceiling)", re.I), "MAX"),
    (re.compile(r"(min|minimum|floor)", re.I), "MIN"),
    (re.compile(r"(rate|ratio|pct|percent)", re.I), "AVG"),
]

_DEFAULT_AGG = "SUM"


class FactHarmonizer:
    """Align fact columns across systems and generate UNION SQL."""

    def __init__(self, fuzzy_threshold: int = 80):
        self._fuzzy_threshold = fuzzy_threshold

    def align_facts(
        self,
        comparison: GrainComparison,
        columns_a: List[str],
        columns_b: List[str],
    ) -> FactAlignment:
        """Create an alignment plan for two fact tables.

        Args:
            comparison: GrainComparison from GrainAnalyzer
            columns_a: Column names from system A's fact table
            columns_b: Column names from system B's fact table
        """
        # Determine target grain
        target_grain = self._resolve_target_grain(comparison)

        # Map columns between the two systems
        col_map, unmapped_a, unmapped_b = self._match_columns(columns_a, columns_b)

        # Determine aggregate expressions for grain bridging
        agg_exprs: Dict[str, str] = {}
        if comparison.alignment_strategy == AlignmentStrategy.AGGREGATE_TO_COARSER:
            finer_cols = columns_a if comparison.relationship == GrainRelationship.A_FINER else columns_b
            finer_system = comparison.system_a if comparison.relationship == GrainRelationship.A_FINER else comparison.system_b
            for col in finer_cols:
                if self._is_measure_column(col):
                    agg = self._infer_aggregation(col)
                    agg_exprs[col] = f"{agg}({col})"

        alignment = FactAlignment(
            source_facts=[
                {"system": comparison.system_a, "table": comparison.table_a, "grain_level": comparison.grain_a.grain_level.value},
                {"system": comparison.system_b, "table": comparison.table_b, "grain_level": comparison.grain_b.grain_level.value},
            ],
            target_grain=target_grain,
            column_mapping=col_map,
            unmapped_columns={
                comparison.system_a: unmapped_a,
                comparison.system_b: unmapped_b,
            },
            aggregate_expressions=agg_exprs,
        )
        return alignment

    def generate_union_sql(
        self,
        alignment: FactAlignment,
        target_table: str = "",
        schema_prefix: str = "",
    ) -> str:
        """Generate UNION ALL SQL from a fact alignment plan.

        Args:
            alignment: FactAlignment from align_facts()
            target_table: Name for the enterprise fact table
            schema_prefix: Optional schema prefix (e.g. "EDW.")
        """
        if not alignment.source_facts or not alignment.column_mapping:
            return "-- No alignment data available"

        # Build unified column list
        enterprise_cols = sorted(alignment.column_mapping.keys())
        source_system_col = "'__SOURCE_SYSTEM__'"

        parts: List[str] = []
        for src in alignment.source_facts:
            system = src["system"]
            table = src["table"]
            grain = src.get("grain_level", "")
            fq_table = f"{schema_prefix}{table}" if schema_prefix else table

            select_cols: List[str] = []
            for ent_col in enterprise_cols:
                src_cols = alignment.column_mapping.get(ent_col, [])
                matched = self._find_system_col(src_cols, system)
                if matched:
                    # Check if aggregation needed
                    agg_expr = alignment.aggregate_expressions.get(matched)
                    if agg_expr and grain != alignment.target_grain.value:
                        select_cols.append(f"  {agg_expr} AS {ent_col}")
                    else:
                        select_cols.append(f"  {matched} AS {ent_col}")
                else:
                    select_cols.append(f"  NULL AS {ent_col}")

            select_cols.append(f"  '{system}' AS SOURCE_SYSTEM")
            select_list = ",\n".join(select_cols)

            # Build GROUP BY if aggregating
            needs_agg = (
                alignment.aggregate_expressions
                and grain != alignment.target_grain.value
            )
            group_by = ""
            if needs_agg:
                non_agg_cols = []
                for ent_col in enterprise_cols:
                    src_cols = alignment.column_mapping.get(ent_col, [])
                    matched = self._find_system_col(src_cols, system)
                    if matched and matched not in alignment.aggregate_expressions:
                        non_agg_cols.append(matched)
                if non_agg_cols:
                    group_by = "\nGROUP BY " + ", ".join(non_agg_cols)

            parts.append(f"SELECT\n{select_list}\nFROM {fq_table}{group_by}")

        union_sql = "\n\nUNION ALL\n\n".join(parts)

        if target_table:
            fq_target = f"{schema_prefix}{target_table}" if schema_prefix else target_table
            return f"CREATE OR REPLACE TABLE {fq_target} AS\n{union_sql};"
        return union_sql + ";"

    def _match_columns(
        self,
        cols_a: List[str],
        cols_b: List[str],
    ) -> Tuple[Dict[str, List[str]], List[str], List[str]]:
        """Match columns between two systems.

        Returns:
            (column_mapping, unmapped_a, unmapped_b)
            column_mapping: enterprise_col -> [sys_a:col, sys_b:col]
        """
        mapping: Dict[str, List[str]] = {}
        used_a: Set[str] = set()
        used_b: Set[str] = set()

        norm_a = {self._normalize_col(c): c for c in cols_a}
        norm_b = {self._normalize_col(c): c for c in cols_b}

        # Pass 1: Exact normalized match
        for na, orig_a in norm_a.items():
            if na in norm_b:
                ent_col = na.upper()
                mapping[ent_col] = [f"A:{orig_a}", f"B:{norm_b[na]}"]
                used_a.add(orig_a)
                used_b.add(norm_b[na])

        # Pass 2: Fuzzy match for remaining
        remaining_a = {self._normalize_col(c): c for c in cols_a if c not in used_a}
        remaining_b = {self._normalize_col(c): c for c in cols_b if c not in used_b}

        try:
            from rapidfuzz import fuzz
            for na, orig_a in remaining_a.items():
                best_score = 0
                best_nb = ""
                for nb in remaining_b:
                    score = fuzz.ratio(na, nb)
                    if score > best_score:
                        best_score = score
                        best_nb = nb
                if best_score >= self._fuzzy_threshold and best_nb:
                    ent_col = na.upper()
                    mapping[ent_col] = [f"A:{orig_a}", f"B:{remaining_b[best_nb]}"]
                    used_a.add(orig_a)
                    used_b.add(remaining_b[best_nb])
        except ImportError:
            logger.debug("rapidfuzz not available, skipping fuzzy column matching")

        unmapped_a = [c for c in cols_a if c not in used_a]
        unmapped_b = [c for c in cols_b if c not in used_b]
        return mapping, unmapped_a, unmapped_b

    @staticmethod
    def _normalize_col(col: str) -> str:
        """Normalize column name for comparison."""
        norm = re.sub(r"^(DIM_|FACT_|ENT_|STG_|SRC_|FK_|PK_)", "", col.upper())
        norm = re.sub(r"(_ID|_KEY|_SK|_NK)$", "", norm)
        return norm

    @staticmethod
    def _is_measure_column(col: str) -> bool:
        """Heuristic: is this column a measure (not a dimension key)?"""
        upper = col.upper()
        if upper.endswith(("_ID", "_KEY", "_SK", "_NK", "_CODE", "_TYPE")):
            return False
        if upper.endswith(("_DATE", "_DT", "_TIMESTAMP", "_TS")):
            return False
        measure_words = ("AMOUNT", "AMT", "QTY", "QUANTITY", "REVENUE", "COST",
                         "PRICE", "RATE", "COUNT", "TOTAL", "SUM", "VALUE", "VOL",
                         "VOLUME", "WEIGHT", "BALANCE", "FEE", "TAX", "DISCOUNT")
        return any(w in upper for w in measure_words)

    @staticmethod
    def _infer_aggregation(col: str) -> str:
        """Infer aggregation function from column name."""
        for pattern, agg in _AGG_HEURISTICS:
            if pattern.search(col):
                return agg
        return _DEFAULT_AGG

    @staticmethod
    def _find_system_col(src_cols: List[str], system_prefix: str) -> str:
        """Find the column for a specific system from mapping entries like 'A:col_name'."""
        prefix_map = {"A": 0, "B": 1}
        for entry in src_cols:
            if ":" in entry:
                pfx, col = entry.split(":", 1)
                if pfx == system_prefix or pfx in prefix_map:
                    return col
        return ""

    def _resolve_target_grain(self, comparison: GrainComparison) -> GrainLevel:
        """Resolve target grain for alignment."""
        if comparison.alignment_strategy == AlignmentStrategy.DIRECT_UNION:
            return comparison.grain_a.grain_level
        return GrainAnalyzer.coarser_grain(
            comparison.grain_a.grain_level,
            comparison.grain_b.grain_level,
        )
