"""SQL logic extractor using sqlglot AST parsing."""
from __future__ import annotations

import logging
from typing import List, Optional

import sqlglot
from sqlglot import exp

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

logger = logging.getLogger(__name__)

# Map sqlglot aggregation names to our enum
_AGG_MAP = {
    "SUM": MeasureAggregation.SUM,
    "COUNT": MeasureAggregation.COUNT,
    "AVG": MeasureAggregation.AVG,
    "MIN": MeasureAggregation.MIN,
    "MAX": MeasureAggregation.MAX,
}

# Map sqlglot join types
_JOIN_MAP = {
    "JOIN": JoinType.INNER,
    "INNER JOIN": JoinType.INNER,
    "LEFT JOIN": JoinType.LEFT,
    "LEFT OUTER JOIN": JoinType.LEFT,
    "RIGHT JOIN": JoinType.RIGHT,
    "RIGHT OUTER JOIN": JoinType.RIGHT,
    "FULL JOIN": JoinType.FULL,
    "FULL OUTER JOIN": JoinType.FULL,
    "CROSS JOIN": JoinType.CROSS,
}

# Map comparison operators
_OP_MAP = {
    "EQ": FilterOperator.EQ,
    "NEQ": FilterOperator.NEQ,
    "GT": FilterOperator.GT,
    "GTE": FilterOperator.GTE,
    "LT": FilterOperator.LT,
    "LTE": FilterOperator.LTE,
    "IN": FilterOperator.IN,
    "LIKE": FilterOperator.LIKE,
    "IS": FilterOperator.EQ,
}


class SQLLogicExtractor:
    """
    Extract business logic from SQL using sqlglot AST.

    Extracts: measures, filters, grain (GROUP BY), joins, dependencies.
    Handles: CTEs, subqueries, window functions, CASE WHEN.
    """

    def __init__(self, dialect: str = "snowflake"):
        self.dialect = dialect

    def extract(self, sql: str, source_path: str = "", source_name: str = "") -> LogicArtifact:
        """Parse SQL and extract all business logic elements."""
        try:
            statements = sqlglot.parse(sql, read=self.dialect)
        except sqlglot.errors.ParseError as e:
            logger.warning("sqlglot parse error for %s: %s", source_path, e)
            return LogicArtifact(
                source_type=SourceType.SQL,
                source_path=source_path,
                source_name=source_name or source_path,
                raw_source=sql[:2000],
                confidence=0.0,
                explanation=f"Parse error: {e}",
            )

        all_measures: List[Measure] = []
        all_filters: List[FilterPredicate] = []
        all_joins: List[JoinClause] = []
        all_grain: List[GrainColumn] = []
        all_deps: List[Dependency] = []

        for stmt in statements:
            if stmt is None:
                continue
            all_measures.extend(self._extract_measures(stmt))
            all_filters.extend(self._extract_filters(stmt))
            all_joins.extend(self._extract_joins(stmt))
            all_grain.extend(self._extract_grain(stmt))
            all_deps.extend(self._extract_dependencies(stmt))

        # Compute confidence based on extraction richness
        confidence = self._compute_confidence(all_measures, all_filters, all_joins, all_grain)

        # Build grain description
        grain_desc = ", ".join(g.column for g in all_grain) if all_grain else ""

        return LogicArtifact(
            source_type=SourceType.SQL,
            source_path=source_path,
            source_name=source_name or source_path,
            raw_source=sql[:2000],
            grain=grain_desc,
            confidence=confidence,
            objects=LogicObjects(
                measures=all_measures,
                filters=all_filters,
                joins=all_joins,
                grain_columns=all_grain,
                dependencies=all_deps,
            ),
        )

    # ------------------------------------------------------------------
    # Measures
    # ------------------------------------------------------------------

    def _extract_measures(self, tree: exp.Expression) -> List[Measure]:
        """Extract aggregate expressions as measures."""
        measures: List[Measure] = []

        for node in tree.walk():
            if not isinstance(node, (exp.Sum, exp.Count, exp.Avg, exp.Min, exp.Max)):
                continue

            agg_name = type(node).__name__.upper()
            agg_type = _AGG_MAP.get(agg_name, MeasureAggregation.CUSTOM)

            # Check for COUNT(DISTINCT ...)
            if isinstance(node, exp.Count) and node.find(exp.Distinct):
                agg_type = MeasureAggregation.COUNT_DISTINCT

            # Get the column(s) being aggregated
            source_cols = []
            for col in node.find_all(exp.Column):
                col_name = col.name
                if col_name:
                    source_cols.append(col_name)

            # Get the source table
            source_table = ""
            col_ref = node.find(exp.Column)
            if col_ref and col_ref.table:
                source_table = col_ref.table

            # Get alias from parent
            alias = ""
            parent = node.parent
            if isinstance(parent, exp.Alias):
                alias = parent.alias

            expression_str = node.sql(dialect=self.dialect)

            measures.append(Measure(
                name=alias or expression_str,
                expression=expression_str,
                aggregation=agg_type,
                source_columns=source_cols,
                source_table=source_table,
                alias=alias,
                confidence=0.9 if alias else 0.7,
            ))

        return measures

    # ------------------------------------------------------------------
    # Filters
    # ------------------------------------------------------------------

    def _extract_filters(self, tree: exp.Expression) -> List[FilterPredicate]:
        """Extract WHERE and HAVING filter predicates."""
        filters: List[FilterPredicate] = []

        for where in tree.find_all(exp.Where):
            filters.extend(self._decompose_condition(where.this, is_having=False))

        for having in tree.find_all(exp.Having):
            filters.extend(self._decompose_condition(having.this, is_having=True))

        return filters

    def _decompose_condition(self, cond: exp.Expression, is_having: bool) -> List[FilterPredicate]:
        """Recursively decompose AND/OR into atomic predicates."""
        predicates: List[FilterPredicate] = []

        if isinstance(cond, (exp.And, exp.Or)):
            predicates.extend(self._decompose_condition(cond.left, is_having))
            predicates.extend(self._decompose_condition(cond.right, is_having))
            return predicates

        # Atomic comparison
        column = ""
        operator = FilterOperator.EQ
        value = ""

        col_ref = cond.find(exp.Column)
        if col_ref:
            column = col_ref.sql(dialect=self.dialect)

        # Determine operator
        op_key = type(cond).__name__.upper()
        operator = _OP_MAP.get(op_key, FilterOperator.EQ)

        if isinstance(cond, exp.In):
            operator = FilterOperator.IN
            in_values = cond.find(exp.Tuple)
            value = in_values.sql(dialect=self.dialect) if in_values else ""
        elif isinstance(cond, exp.Between):
            operator = FilterOperator.BETWEEN
            value = cond.sql(dialect=self.dialect)
        elif isinstance(cond, exp.Is):
            right = cond.find(exp.Null)
            if right:
                if cond.find(exp.Not):
                    operator = FilterOperator.IS_NOT_NULL
                else:
                    operator = FilterOperator.IS_NULL
                value = "NULL"
        elif isinstance(cond, exp.Like):
            operator = FilterOperator.LIKE
            lit = cond.find(exp.Literal)
            value = lit.this if lit else ""
        else:
            # Generic binary comparison
            lit = cond.find(exp.Literal)
            if lit:
                value = lit.this

        predicates.append(FilterPredicate(
            column=column,
            operator=operator,
            value=value,
            source_clause=cond.sql(dialect=self.dialect),
            is_having=is_having,
        ))

        return predicates

    # ------------------------------------------------------------------
    # Joins
    # ------------------------------------------------------------------

    def _extract_joins(self, tree: exp.Expression) -> List[JoinClause]:
        """Extract join clauses with table names and keys."""
        joins: List[JoinClause] = []

        for join_node in tree.find_all(exp.Join):
            # sqlglot uses 'side' for LEFT/RIGHT/FULL and 'kind' for CROSS etc.
            side = getattr(join_node, "side", "") or ""
            kind = getattr(join_node, "kind", "") or ""
            if side:
                join_type_str = f"{side} JOIN".upper()
            elif kind:
                join_type_str = f"{kind} JOIN".upper()
            else:
                join_type_str = "JOIN"
            jt = _JOIN_MAP.get(join_type_str, JoinType.INNER)

            # Right table (the joined table)
            right_table = ""
            table_ref = join_node.find(exp.Table)
            if table_ref:
                right_table = table_ref.name or ""

            # Join condition — stored directly in join_node.args['on']
            on_cond = join_node.args.get("on")
            condition = on_cond.sql(dialect=self.dialect) if on_cond else ""

            # Extract key columns from ON clause
            left_key = ""
            right_key = ""
            if on_cond:
                eq = on_cond if isinstance(on_cond, exp.EQ) else on_cond.find(exp.EQ)
                if eq:
                    cols = list(eq.find_all(exp.Column))
                    if len(cols) >= 2:
                        left_key = cols[0].name or ""
                        right_key = cols[1].name or ""

            # Left table — find the FROM clause
            left_table = ""
            select = tree.find(exp.Select)
            if select:
                from_clause = select.find(exp.From)
                if from_clause:
                    from_table = from_clause.find(exp.Table)
                    if from_table:
                        left_table = from_table.name or ""

            joins.append(JoinClause(
                left_table=left_table,
                right_table=right_table,
                join_type=jt,
                left_key=left_key,
                right_key=right_key,
                condition=condition,
            ))

        return joins

    # ------------------------------------------------------------------
    # Grain (GROUP BY)
    # ------------------------------------------------------------------

    def _extract_grain(self, tree: exp.Expression) -> List[GrainColumn]:
        """Extract GROUP BY columns as grain definition."""
        grain: List[GrainColumn] = []

        for group in tree.find_all(exp.Group):
            for col in group.find_all(exp.Column):
                grain.append(GrainColumn(
                    column=col.name or "",
                    table=col.table or "",
                ))

        return grain

    # ------------------------------------------------------------------
    # Dependencies (tables/views/CTEs referenced)
    # ------------------------------------------------------------------

    def _extract_dependencies(self, tree: exp.Expression) -> List[Dependency]:
        """Extract all table/view/CTE references."""
        deps: List[Dependency] = []
        seen = set()

        # CTEs
        for cte in tree.find_all(exp.CTE):
            name = cte.alias or ""
            if name and name not in seen:
                deps.append(Dependency(name=name, dep_type="cte"))
                seen.add(name)

        # Tables
        for table in tree.find_all(exp.Table):
            name = table.name or ""
            if not name or name in seen:
                continue
            schema = table.db or ""
            database = table.catalog or ""
            deps.append(Dependency(
                name=name,
                schema_name=schema,
                database=database,
                dep_type="table",
            ))
            seen.add(name)

        return deps

    # ------------------------------------------------------------------
    # Confidence scoring
    # ------------------------------------------------------------------

    def _compute_confidence(
        self,
        measures: List[Measure],
        filters: List[FilterPredicate],
        joins: List[JoinClause],
        grain: List[GrainColumn],
    ) -> float:
        """
        Compute extraction confidence based on richness.
        Higher when more elements are extracted successfully.
        """
        score = 0.3  # Base: we parsed something

        if measures:
            score += 0.25
        if filters:
            score += 0.15
        if joins:
            score += 0.15
        if grain:
            score += 0.15

        return min(score, 1.0)
