"""Python logic extractor using stdlib ast module.

Targets pandas-style ETL scripts (groupby, merge, agg, filter, query).
Also detects embedded SQL strings and delegates to SQLLogicExtractor.
"""
from __future__ import annotations

import ast
import logging
import re
from typing import List

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

logger = logging.getLogger(__name__)

# Regex to detect SQL strings in Python code
_SQL_PATTERN = re.compile(
    r"""(?:SELECT|INSERT|UPDATE|DELETE|CREATE|WITH)\s""",
    re.IGNORECASE,
)

# Pandas aggregation function names -> our measure types
_PANDAS_AGG_MAP = {
    "sum": MeasureAggregation.SUM,
    "count": MeasureAggregation.COUNT,
    "mean": MeasureAggregation.AVG,
    "avg": MeasureAggregation.AVG,
    "min": MeasureAggregation.MIN,
    "max": MeasureAggregation.MAX,
    "nunique": MeasureAggregation.COUNT_DISTINCT,
    "std": MeasureAggregation.CUSTOM,
    "var": MeasureAggregation.CUSTOM,
    "median": MeasureAggregation.CUSTOM,
}


class _BusinessLogicVisitor(ast.NodeVisitor):
    """AST visitor that collects pandas business logic patterns."""

    def __init__(self):
        self.measures: List[Measure] = []
        self.filters: List[FilterPredicate] = []
        self.joins: List[JoinClause] = []
        self.grain_columns: List[GrainColumn] = []
        self.dependencies: List[Dependency] = []
        self.embedded_sql: List[str] = []
        self.assignments: List[str] = []

    # ---- .groupby(...).agg(...) ----

    def visit_Call(self, node: ast.Call):
        func_name = self._get_method_name(node)

        if func_name == "groupby":
            self._handle_groupby(node)
        elif func_name == "merge":
            self._handle_merge(node)
        elif func_name == "agg":
            self._handle_agg(node)
        elif func_name == "query":
            self._handle_query(node)
        elif func_name == "filter":
            self._handle_filter(node)
        elif func_name in ("read_csv", "read_excel", "read_sql", "read_parquet"):
            self._handle_read(node, func_name)

        self.generic_visit(node)

    def visit_Assign(self, node: ast.Assign):
        """Detect business logic in variable assignments like revenue = df['a'] * df['b']."""
        if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name):
            var_name = node.targets[0].id
            # Check for arithmetic on DataFrame columns
            if self._has_subscript_ops(node.value):
                expr = ast.dump(node.value)
                self.measures.append(Measure(
                    name=var_name,
                    expression=f"{var_name} = <computed>",
                    aggregation=MeasureAggregation.CUSTOM,
                    confidence=0.5,
                ))
        self.generic_visit(node)

    def visit_Constant(self, node: ast.Constant):
        """Detect SQL strings embedded in Python."""
        if isinstance(node.value, str) and len(node.value) > 20:
            if _SQL_PATTERN.search(node.value):
                self.embedded_sql.append(node.value)
        self.generic_visit(node)

    # ---- Handlers ----

    def _handle_groupby(self, node: ast.Call):
        """Extract GROUP BY columns from df.groupby(...)."""
        for arg in node.args:
            cols = self._extract_string_values(arg)
            for col in cols:
                self.grain_columns.append(GrainColumn(column=col))

    def _handle_merge(self, node: ast.Call):
        """Extract join info from pd.merge(...) or df.merge(...)."""
        left_on = ""
        right_on = ""
        how = "inner"

        for kw in node.keywords:
            if kw.arg == "on":
                cols = self._extract_string_values(kw.value)
                if cols:
                    left_on = cols[0]
                    right_on = cols[0]
            elif kw.arg == "left_on":
                cols = self._extract_string_values(kw.value)
                if cols:
                    left_on = cols[0]
            elif kw.arg == "right_on":
                cols = self._extract_string_values(kw.value)
                if cols:
                    right_on = cols[0]
            elif kw.arg == "how":
                if isinstance(kw.value, ast.Constant):
                    how = str(kw.value.value)

        jt_map = {"inner": JoinType.INNER, "left": JoinType.LEFT,
                   "right": JoinType.RIGHT, "outer": JoinType.FULL, "cross": JoinType.CROSS}

        self.joins.append(JoinClause(
            left_table="<dataframe>",
            right_table="<dataframe>",
            join_type=jt_map.get(how, JoinType.INNER),
            left_key=left_on,
            right_key=right_on,
        ))

    def _handle_agg(self, node: ast.Call):
        """Extract aggregation measures from .agg({...}) or .agg(func)."""
        for arg in node.args:
            if isinstance(arg, ast.Dict):
                for key, val in zip(arg.keys, arg.values):
                    col_name = key.value if isinstance(key, ast.Constant) else ""
                    agg_name = val.value if isinstance(val, ast.Constant) else ""
                    agg_type = _PANDAS_AGG_MAP.get(str(agg_name), MeasureAggregation.CUSTOM)
                    if col_name:
                        self.measures.append(Measure(
                            name=f"{agg_name}({col_name})",
                            expression=f"{agg_name}({col_name})",
                            aggregation=agg_type,
                            source_columns=[col_name],
                            confidence=0.7,
                        ))
            elif isinstance(arg, ast.Constant):
                agg_name = str(arg.value)
                agg_type = _PANDAS_AGG_MAP.get(agg_name, MeasureAggregation.CUSTOM)
                self.measures.append(Measure(
                    name=agg_name,
                    expression=agg_name,
                    aggregation=agg_type,
                    confidence=0.5,
                ))

        # Also check keyword arguments like .agg(revenue='sum')
        for kw in node.keywords:
            if kw.arg and isinstance(kw.value, ast.Constant):
                agg_name = str(kw.value.value)
                agg_type = _PANDAS_AGG_MAP.get(agg_name, MeasureAggregation.CUSTOM)
                self.measures.append(Measure(
                    name=f"{agg_name}({kw.arg})",
                    expression=f"{agg_name}({kw.arg})",
                    aggregation=agg_type,
                    source_columns=[kw.arg],
                    confidence=0.7,
                ))

    def _handle_query(self, node: ast.Call):
        """Extract filter from df.query('col > 5')."""
        for arg in node.args:
            if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
                self.filters.append(FilterPredicate(
                    column="",
                    source_clause=arg.value,
                    confidence=0.6,
                ))

    def _handle_filter(self, node: ast.Call):
        """Extract filter from df.filter(items=[...]) or df.filter(like='...')."""
        for kw in node.keywords:
            if kw.arg == "like" and isinstance(kw.value, ast.Constant):
                self.filters.append(FilterPredicate(
                    column="",
                    source_clause=f"filter(like='{kw.value.value}')",
                    confidence=0.5,
                ))

    def _handle_read(self, node: ast.Call, func_name: str):
        """Track data source reads as dependencies."""
        for arg in node.args:
            if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
                self.dependencies.append(Dependency(
                    name=arg.value,
                    dep_type="file" if func_name != "read_sql" else "sql",
                ))
                break

    # ---- Utilities ----

    def _get_method_name(self, node: ast.Call) -> str:
        if isinstance(node.func, ast.Attribute):
            return node.func.attr
        if isinstance(node.func, ast.Name):
            return node.func.id
        return ""

    def _extract_string_values(self, node: ast.expr) -> List[str]:
        """Extract string values from a constant or list of constants."""
        if isinstance(node, ast.Constant) and isinstance(node.value, str):
            return [node.value]
        if isinstance(node, ast.List):
            return [
                elt.value for elt in node.elts
                if isinstance(elt, ast.Constant) and isinstance(elt.value, str)
            ]
        return []

    def _has_subscript_ops(self, node: ast.expr) -> bool:
        """Check if an expression contains DataFrame-style subscript operations."""
        if isinstance(node, ast.Subscript):
            return True
        if isinstance(node, ast.BinOp):
            return self._has_subscript_ops(node.left) or self._has_subscript_ops(node.right)
        return False


class PythonLogicExtractor:
    """Extract business logic from Python source code (pandas focus)."""

    def extract(self, source_code: str, source_path: str = "", source_name: str = "") -> LogicArtifact:
        """Parse Python and extract business logic elements."""
        try:
            tree = ast.parse(source_code)
        except SyntaxError as e:
            logger.warning("Python parse error for %s: %s", source_path, e)
            return LogicArtifact(
                source_type=SourceType.PYTHON,
                source_path=source_path,
                source_name=source_name or source_path,
                raw_source=source_code[:2000],
                confidence=0.0,
                explanation=f"Syntax error: {e}",
            )

        visitor = _BusinessLogicVisitor()
        visitor.visit(tree)

        # Compute confidence
        has_content = bool(visitor.measures or visitor.filters or visitor.joins or visitor.grain_columns)
        confidence = 0.6 if has_content else 0.2

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

        artifact = LogicArtifact(
            source_type=SourceType.PYTHON,
            source_path=source_path,
            source_name=source_name or source_path,
            raw_source=source_code[:2000],
            grain=grain_desc,
            confidence=confidence,
            objects=LogicObjects(
                measures=visitor.measures,
                filters=visitor.filters,
                joins=visitor.joins,
                grain_columns=visitor.grain_columns,
                dependencies=visitor.dependencies,
            ),
            metadata={"embedded_sql_count": len(visitor.embedded_sql)},
        )

        return artifact

    def extract_file(self, file_path: str) -> LogicArtifact:
        """Extract from a Python file on disk."""
        with open(file_path, "r", encoding="utf-8", errors="replace") as f:
            source = f.read()
        return self.extract(source, source_path=file_path)
