"""DAX logic extractor using regex-based parsing.

Extracts business logic from Power BI DAX expressions:
- Measure definitions (CALCULATE, SUMX, COUNTROWS, etc.)
- Filter expressions (FILTER, ALL, ALLEXCEPT, VALUES)
- Table/column references as dependencies
- Dimension references as grain columns
"""
from __future__ import annotations

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

from ..constants import (
    FilterOperator,
    MeasureAggregation,
    SourceType,
)
from ..contracts import (
    Dependency,
    FilterPredicate,
    GrainColumn,
    LogicArtifact,
    LogicObjects,
    Measure,
)

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# DAX aggregation function mapping
# ---------------------------------------------------------------------------

_DAX_AGG_MAP: Dict[str, MeasureAggregation] = {
    "SUM": MeasureAggregation.SUM,
    "SUMX": MeasureAggregation.SUM,
    "COUNT": MeasureAggregation.COUNT,
    "COUNTA": MeasureAggregation.COUNT,
    "COUNTX": MeasureAggregation.COUNT,
    "COUNTROWS": MeasureAggregation.COUNT,
    "DISTINCTCOUNT": MeasureAggregation.COUNT_DISTINCT,
    "DISTINCTCOUNTX": MeasureAggregation.COUNT_DISTINCT,
    "COUNTBLANK": MeasureAggregation.COUNT,
    "AVERAGE": MeasureAggregation.AVG,
    "AVERAGEX": MeasureAggregation.AVG,
    "AVERAGEA": MeasureAggregation.AVG,
    "MIN": MeasureAggregation.MIN,
    "MINX": MeasureAggregation.MIN,
    "MINA": MeasureAggregation.MIN,
    "MAX": MeasureAggregation.MAX,
    "MAXX": MeasureAggregation.MAX,
    "MAXA": MeasureAggregation.MAX,
    "CALCULATE": MeasureAggregation.CUSTOM,
    "CALCULATETABLE": MeasureAggregation.CUSTOM,
    "DIVIDE": MeasureAggregation.CUSTOM,
    "IF": MeasureAggregation.CUSTOM,
    "SWITCH": MeasureAggregation.CUSTOM,
}

# ---------------------------------------------------------------------------
# Regex patterns for DAX parsing
# ---------------------------------------------------------------------------

# Measure definition: MeasureName = <expression>  or  MeasureName := <expression>
_MEASURE_DEF_PATTERN = re.compile(
    r"(?:MEASURE\s+)?([A-Za-z_][\w\s]*?)\s*[:=]=?\s*\n?\s*((?:CALCULATE|SUMX?|COUNTX?|COUNTA?|COUNTROWS|DISTINCTCOUNT|AVERAGE[AX]?|MIN[AX]?|MAX[AX]?|DIVIDE|IF|SWITCH|CALCULATETABLE)\s*\(.*)",
    re.IGNORECASE | re.DOTALL,
)

# Simpler: name = FUNC(...)
_NAMED_MEASURE_PATTERN = re.compile(
    r"^\s*(?:MEASURE\s+(?:\[?[\w\s]+\]?\s*\[)?\s*)?(\[?[A-Za-z_][\w\s]*?\]?)\s*[:=]=?\s*(.+)",
    re.MULTILINE,
)

# DAX aggregation call: FUNC(args)
_AGG_CALL_PATTERN = re.compile(
    r"\b(SUM|SUMX|COUNT|COUNTA|COUNTX|COUNTROWS|DISTINCTCOUNT|DISTINCTCOUNTX|COUNTBLANK|"
    r"AVERAGE|AVERAGEX|AVERAGEA|MIN|MINX|MINA|MAX|MAXX|MAXA|CALCULATE|CALCULATETABLE|"
    r"DIVIDE|IF|SWITCH)\s*\(",
    re.IGNORECASE,
)

# Table reference: 'TableName'[Column] or TableName[Column] or just 'TableName'
_TABLE_COL_PATTERN = re.compile(
    r"'?([A-Za-z_][\w\s]*?)'?\[([A-Za-z_][\w\s]*?)\]",
)

# Standalone table reference: 'Table Name' (in FILTER, RELATEDTABLE, etc.)
_TABLE_REF_PATTERN = re.compile(
    r"(?:FILTER|RELATEDTABLE|ALL|ALLEXCEPT|VALUES|DISTINCT|RELATED|SUMMARIZE|"
    r"ADDCOLUMNS|SELECTCOLUMNS|GROUPBY|TOPN|EARLIER|EARLIEST|LOOKUPVALUE|"
    r"CALCULATETABLE|NATURALLEFOUTERJOIN|NATURALINNERJOIN|CROSSJOIN|UNION|INTERSECT|EXCEPT)"
    r"\s*\(\s*'?([A-Za-z_][\w\s]*?)'?\s*[,\)]",
    re.IGNORECASE,
)

# FILTER expressions: FILTER(Table, condition)
_FILTER_PATTERN = re.compile(
    r"FILTER\s*\(\s*(?:'?([A-Za-z_][\w\s]*?)'?\s*,\s*)(.+?)\)",
    re.IGNORECASE,
)

# Simple filter in CALCULATE: Table[Column] = value or Table[Column] > value
_CALC_FILTER_PATTERN = re.compile(
    r"'?([A-Za-z_][\w\s]*?)'?\[([A-Za-z_][\w\s]*?)\]\s*(=|<>|>=|<=|>|<)\s*(.+?)(?:\s*[,)])",
    re.IGNORECASE,
)

# ALL / ALLEXCEPT patterns
_ALL_PATTERN = re.compile(
    r"\bALL\s*\(\s*'?([A-Za-z_][\w\s]*?)'?(?:\s*\[.+?\])?\s*\)",
    re.IGNORECASE,
)

# SUMMARIZE / GROUPBY grain pattern
_SUMMARIZE_PATTERN = re.compile(
    r"(?:SUMMARIZE|GROUPBY)\s*\(\s*'?([A-Za-z_][\w\s]*?)'?\s*,\s*(.+?)\)",
    re.IGNORECASE | re.DOTALL,
)

# VAR pattern for variable assignments
_VAR_PATTERN = re.compile(
    r"\bVAR\s+(\w+)\s*=\s*(.+?)(?=\bVAR\b|\bRETURN\b|$)",
    re.IGNORECASE | re.DOTALL,
)

# RETURN pattern
_RETURN_PATTERN = re.compile(
    r"\bRETURN\s+(.+?)$",
    re.IGNORECASE | re.DOTALL,
)

# DAX filter operator mapping
_DAX_OP_MAP = {
    "=": FilterOperator.EQ,
    "<>": FilterOperator.NEQ,
    ">": FilterOperator.GT,
    ">=": FilterOperator.GTE,
    "<": FilterOperator.LT,
    "<=": FilterOperator.LTE,
}


class DAXLogicExtractor:
    """Extract business logic from Power BI DAX expressions.

    Handles:
    - Measure definitions (named or inline)
    - CALCULATE with filter arguments
    - Iterator functions (SUMX, COUNTX, AVERAGEX, etc.)
    - FILTER, ALL, ALLEXCEPT context modifiers
    - SUMMARIZE/GROUPBY grain detection
    - Table[Column] reference extraction
    """

    def extract(
        self,
        source: str,
        source_path: str = "",
        source_name: str = "",
    ) -> LogicArtifact:
        """Parse DAX source and extract all business logic elements."""
        if not source or not source.strip():
            return LogicArtifact(
                source_type=SourceType.DAX,
                source_path=source_path,
                source_name=source_name,
                raw_source="",
                confidence=0.0,
                explanation="Empty DAX source",
            )

        try:
            measures = self._extract_measures(source)
            filters = self._extract_filters(source)
            grain_columns = self._extract_grain(source)
            dependencies = self._extract_dependencies(source)
            confidence = self._compute_confidence(measures, filters, grain_columns, dependencies)

            grain_desc = ""
            if grain_columns:
                grain_desc = ", ".join(g.column for g in grain_columns)

            return LogicArtifact(
                source_type=SourceType.DAX,
                source_path=source_path,
                source_name=source_name,
                raw_source=source[:2000],
                grain=grain_desc,
                confidence=confidence,
                objects=LogicObjects(
                    measures=measures,
                    filters=filters,
                    grain_columns=grain_columns,
                    dependencies=dependencies,
                ),
                explanation=self._build_explanation(measures, filters, grain_columns, dependencies),
            )
        except Exception as e:
            logger.warning("DAX parse error for %s: %s", source_path or source_name, e)
            return LogicArtifact(
                source_type=SourceType.DAX,
                source_path=source_path,
                source_name=source_name,
                raw_source=source[:2000],
                confidence=0.0,
                explanation=f"Parse error: {e}",
            )

    def extract_file(self, file_path: str) -> LogicArtifact:
        """Extract logic from a DAX file on disk."""
        from pathlib import Path

        path = Path(file_path)
        if not path.exists():
            return LogicArtifact(
                source_type=SourceType.DAX,
                source_path=file_path,
                confidence=0.0,
                explanation=f"File not found: {file_path}",
            )

        source = path.read_text(encoding="utf-8", errors="replace")
        return self.extract(source, source_path=file_path, source_name=path.stem)

    # ------------------------------------------------------------------
    # Measure extraction
    # ------------------------------------------------------------------

    def _extract_measures(self, source: str) -> List[Measure]:
        """Extract DAX measures from source."""
        measures: List[Measure] = []
        seen_names: Set[str] = set()

        # 1. Try named measure definitions (Name = FUNC(...))
        for match in _NAMED_MEASURE_PATTERN.finditer(source):
            name = match.group(1).strip().strip("[]")
            expr = match.group(2).strip()

            if name.lower() in seen_names or not expr:
                continue

            # Find the aggregation function
            agg = MeasureAggregation.CUSTOM
            agg_match = _AGG_CALL_PATTERN.search(expr)
            if agg_match:
                func_name = agg_match.group(1).upper()
                agg = _DAX_AGG_MAP.get(func_name, MeasureAggregation.CUSTOM)

            # Extract source columns from Table[Column] references
            cols = self._extract_column_refs(expr)

            measures.append(Measure(
                name=name,
                expression=expr[:500],
                aggregation=agg,
                source_columns=[c for _, c in cols],
                source_table=cols[0][0] if cols else "",
                confidence=0.8 if agg != MeasureAggregation.CUSTOM else 0.6,
            ))
            seen_names.add(name.lower())

        # 2. If no named measures found, look for standalone aggregation calls
        if not measures:
            for match in _AGG_CALL_PATTERN.finditer(source):
                func_name = match.group(1).upper()
                agg = _DAX_AGG_MAP.get(func_name, MeasureAggregation.CUSTOM)

                # Skip CALCULATE/IF/SWITCH as standalone (they're wrappers)
                if func_name in ("CALCULATE", "CALCULATETABLE", "IF", "SWITCH"):
                    continue

                # Get the argument region (rough extraction)
                start = match.end()
                depth = 1
                end = start
                for i in range(start, min(start + 500, len(source))):
                    if source[i] == "(":
                        depth += 1
                    elif source[i] == ")":
                        depth -= 1
                        if depth == 0:
                            end = i
                            break

                arg_text = source[start:end].strip()
                cols = self._extract_column_refs(arg_text)

                measure_name = f"{func_name.lower()}_{cols[0][1]}" if cols else func_name.lower()
                if measure_name.lower() in seen_names:
                    continue

                measures.append(Measure(
                    name=measure_name,
                    expression=f"{func_name}({arg_text[:200]})",
                    aggregation=agg,
                    source_columns=[c for _, c in cols],
                    source_table=cols[0][0] if cols else "",
                    confidence=0.65,
                ))
                seen_names.add(measure_name.lower())

        return measures

    # ------------------------------------------------------------------
    # Filter extraction
    # ------------------------------------------------------------------

    def _extract_filters(self, source: str) -> List[FilterPredicate]:
        """Extract filter predicates from FILTER() and CALCULATE() arguments."""
        filters: List[FilterPredicate] = []
        seen: Set[str] = set()

        # 1. FILTER(Table, condition)
        for match in _FILTER_PATTERN.finditer(source):
            condition = match.group(2).strip()
            # Parse the condition for comparison operators
            for op_match in _CALC_FILTER_PATTERN.finditer(condition + ")"):
                col = op_match.group(2).strip()
                op_str = op_match.group(3).strip()
                val = op_match.group(4).strip().rstrip(",)")
                key = f"{col}|{op_str}|{val}"
                if key not in seen:
                    filters.append(FilterPredicate(
                        column=col,
                        operator=_DAX_OP_MAP.get(op_str, FilterOperator.EQ),
                        value=val,
                        source_clause=condition[:200],
                    ))
                    seen.add(key)

        # 2. Simple filters in CALCULATE(expr, filter1, filter2, ...)
        for match in _CALC_FILTER_PATTERN.finditer(source):
            col = match.group(2).strip()
            op_str = match.group(3).strip()
            val = match.group(4).strip().rstrip(",)")
            key = f"{col}|{op_str}|{val}"
            if key not in seen:
                filters.append(FilterPredicate(
                    column=col,
                    operator=_DAX_OP_MAP.get(op_str, FilterOperator.EQ),
                    value=val,
                    source_clause=f"{match.group(1).strip()}[{col}] {op_str} {val}",
                ))
                seen.add(key)

        return filters

    # ------------------------------------------------------------------
    # Grain extraction
    # ------------------------------------------------------------------

    def _extract_grain(self, source: str) -> List[GrainColumn]:
        """Extract grain columns from SUMMARIZE/GROUPBY and dimension refs."""
        grain_columns: List[GrainColumn] = []
        seen: Set[str] = set()

        # 1. SUMMARIZE(Table, Table[Col1], Table[Col2], ...)
        for match in _SUMMARIZE_PATTERN.finditer(source):
            table = match.group(1).strip()
            args = match.group(2)
            for col_match in _TABLE_COL_PATTERN.finditer(args):
                col_table = col_match.group(1).strip()
                col_name = col_match.group(2).strip()
                key = f"{col_table}.{col_name}".lower()
                if key not in seen:
                    grain_columns.append(GrainColumn(
                        column=col_name,
                        table=col_table,
                    ))
                    seen.add(key)

        # 2. ALLEXCEPT(Table, Table[Col]) — columns NOT removed = grain
        for match in re.finditer(
            r"ALLEXCEPT\s*\(\s*'?(\w[\w\s]*?)'?\s*,\s*(.+?)\)",
            source,
            re.IGNORECASE,
        ):
            args = match.group(2)
            for col_match in _TABLE_COL_PATTERN.finditer(args):
                col_table = col_match.group(1).strip()
                col_name = col_match.group(2).strip()
                key = f"{col_table}.{col_name}".lower()
                if key not in seen:
                    grain_columns.append(GrainColumn(
                        column=col_name,
                        table=col_table,
                    ))
                    seen.add(key)

        return grain_columns

    # ------------------------------------------------------------------
    # Dependency extraction
    # ------------------------------------------------------------------

    def _extract_dependencies(self, source: str) -> List[Dependency]:
        """Extract table references from DAX source."""
        tables: Set[str] = set()

        # Table[Column] references
        for match in _TABLE_COL_PATTERN.finditer(source):
            table = match.group(1).strip()
            if table and not table.upper().startswith(("VAR", "RETURN")):
                tables.add(table)

        # Standalone table references in functions
        for match in _TABLE_REF_PATTERN.finditer(source):
            table = match.group(1).strip()
            if table:
                tables.add(table)

        # ALL(Table) references
        for match in _ALL_PATTERN.finditer(source):
            table = match.group(1).strip()
            if table:
                tables.add(table)

        return [
            Dependency(name=t, dep_type="table")
            for t in sorted(tables)
        ]

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _extract_column_refs(self, text: str) -> List[Tuple[str, str]]:
        """Extract (table, column) pairs from Table[Column] references."""
        refs = []
        seen = set()
        for match in _TABLE_COL_PATTERN.finditer(text):
            table = match.group(1).strip()
            col = match.group(2).strip()
            key = f"{table}.{col}".lower()
            if key not in seen:
                refs.append((table, col))
                seen.add(key)
        return refs

    def _compute_confidence(
        self,
        measures: List[Measure],
        filters: List[FilterPredicate],
        grain: List[GrainColumn],
        deps: List[Dependency],
    ) -> float:
        """Compute overall confidence score."""
        score = 0.0
        if measures:
            score += 0.35
            # Bonus for named measures (vs inferred)
            named = sum(1 for m in measures if m.name and not m.name.startswith(("sum_", "count_", "avg_")))
            if named > 0:
                score += 0.10
        if filters:
            score += 0.15
        if grain:
            score += 0.15
        if deps:
            score += 0.15
        # Cap at 1.0
        return min(round(score, 2), 1.0)

    def _build_explanation(
        self,
        measures: List[Measure],
        filters: List[FilterPredicate],
        grain: List[GrainColumn],
        deps: List[Dependency],
    ) -> str:
        parts = []
        if measures:
            names = [m.name or m.expression[:30] for m in measures[:5]]
            parts.append(f"{len(measures)} measure(s): {', '.join(names)}")
        if filters:
            parts.append(f"{len(filters)} filter(s)")
        if grain:
            cols = [g.column for g in grain[:5]]
            parts.append(f"grain: {', '.join(cols)}")
        if deps:
            tables = [d.name for d in deps[:5]]
            parts.append(f"tables: {', '.join(tables)}")
        return "; ".join(parts) if parts else "No business logic elements detected"
