from __future__ import annotations
from typing import TYPE_CHECKING, TypeVar, Literal as PyLiteral
from dataclasses import dataclass, field

if TYPE_CHECKING:
    from .sql_compiler import ColumnInfo
    from pandas import DataFrame

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

def _indent(s: str, level: int = 1) -> str:
    """Indent each line of a string."""
    prefix = "  " * level
    return "\n".join(prefix + line for line in s.split("\n"))

def _hash_value(value: object) -> int:
    """Recursively compute a content-based hash for any value."""
    if value is None:
        return hash(None)
    if isinstance(value, IRNode):
        return value.content_hash()
    if isinstance(value, (list, tuple)):
        return hash(tuple(_hash_value(item) for item in value))
    if isinstance(value, dict):
        return hash(tuple(sorted((_hash_value(k), _hash_value(v)) for k, v in value.items())))
    # For objects with name/alias attributes (like TableInfo, ColumnInfo)
    name = getattr(value, "name", None)
    if name is not None:
        return hash(("name", name))
    alias = getattr(value, "alias", None)
    if alias is not None:
        return hash(("alias", alias))
    # Try normal hash, fall back to string repr
    try:
        return hash(value)
    except TypeError:
        return hash(repr(value))

#------------------------------------------------------
# IRNode
#------------------------------------------------------

class IRNode:
    def _repr_fields(self) -> list[tuple[str, object]]:
        """Override in subclasses to specify fields for repr."""
        return []

    def _hash_fields(self) -> list[tuple[str, object]]:
        """Override in subclasses to specify fields for hashing. Defaults to _repr_fields."""
        return self._repr_fields()

    def content_hash(self) -> int:
        """Compute a stable content-based hash for this IR node. Cached after first call."""
        cached = self.__dict__.get('_cached_hash')
        if cached is not None:
            return cached
        fields = self._hash_fields()
        field_hashes = tuple(
            (name, _hash_value(value)) for name, value in fields
        )
        h = hash((self.__class__.__name__, field_hashes))
        self.__dict__['_cached_hash'] = h
        return h

    def __hash__(self) -> int:
        return self.content_hash()

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, IRNode):
            return False
        if self.__class__ != other.__class__:
            return False
        return self._hash_fields() == other._hash_fields()

    def __repr__(self) -> str:
        class_name = self.__class__.__name__
        fields = self._repr_fields()
        if not fields:
            return f"{class_name}()"

        # Check if we can use compact single-line format (2 or fewer args, no nested structures)
        def is_simple(value: object) -> bool:
            if isinstance(value, list):
                return not value  # empty list is simple
            if isinstance(value, IRNode):
                return False
            return True

        if len(fields) <= 2 and all(is_simple(v) for _, v in fields):
            parts = [f"{name}={repr(value)}" for name, value in fields]
            return f"{class_name}({', '.join(parts)})"

        parts = []
        for name, value in fields:
            if isinstance(value, list):
                if not value:
                    parts.append(f"{name}=[]")
                else:
                    items = ",\n".join(repr(item) for item in value)
                    parts.append(f"{name}=[\n{_indent(items)}\n]")
            elif isinstance(value, IRNode):
                parts.append(f"{name}={repr(value)}")
            else:
                parts.append(f"{name}={repr(value)}")

        if len(parts) == 1 and "\n" not in parts[0]:
            return f"{class_name}({parts[0]})"

        inner = ",\n".join(parts)
        return f"{class_name}(\n{_indent(inner)}\n)"

#------------------------------------------------------
# Fragment
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class Fragment(IRNode):
    selects: list[IRNode] = field(default_factory=list)
    scans: list[IRNode] = field(default_factory=list)
    joins: list[Join] = field(default_factory=list)
    filters: list[IRNode] = field(default_factory=list)  # Filter, Exists, Not, etc.
    group_by: list[IRNode] = field(default_factory=list)
    aggregates: list[AggregatePlaceholder] = field(default_factory=list)

    def classify(self) -> PyLiteral["expression", "filtered", "query"]:
        if not self.aggregates and len(self.scans) == 0 and len(self.joins) == 0:
            if not self.filters:
                return "expression"
            return "filtered"
        return "query"

    def clone(self) -> Fragment:
        """Create a shallow copy of this fragment."""
        return Fragment(
            selects=list(self.selects),
            scans=list(self.scans),
            joins=list(self.joins),
            filters=list(self.filters),
            group_by=list(self.group_by),
            aggregates=list(self.aggregates),
        )

    def _repr_fields(self) -> list[tuple[str, object]]:
        fields = []
        if self.selects:
            fields.append(("selects", self.selects))
        if self.scans:
            fields.append(("scans", self.scans))
        if self.joins:
            fields.append(("joins", self.joins))
        if self.filters:
            fields.append(("filters", self.filters))
        if self.group_by:
            fields.append(("group_by", self.group_by))
        if self.aggregates:
            fields.append(("aggregates", self.aggregates))
        return fields

#------------------------------------------------------
# Expression
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class Expression(IRNode):
    op: str = field()
    args: list[IRNode] = field(default_factory=list)

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [("op", self.op), ("args", self.args)]

#------------------------------------------------------
# TableFunction
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class TableFunction(IRNode):
    op: str = field()
    args: list[IRNode] = field(default_factory=list)

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [("op", self.op), ("args", self.args)]

#------------------------------------------------------
# WindowExpression
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class WindowExpression(IRNode):
    op: str = field()
    args: list[IRNode] = field(default_factory=list)
    partition_by: list[IRNode] = field(default_factory=list)
    order_by: list[IRNode] = field(default_factory=list)

    def _repr_fields(self) -> list[tuple[str, object]]:
        fields: list[tuple[str, object]] = [("op", self.op), ("args", self.args)]
        if self.partition_by:
            fields.append(("partition_by", self.partition_by))
        if self.order_by:
            fields.append(("order_by", self.order_by))
        return fields
#------------------------------------------------------
# Literal
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class Literal(IRNode):
    value: object = field()
    type: str = field()

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [("value", self.value), ("type", self.type)]

#------------------------------------------------------
# Filter
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class Filter(IRNode):
    op: str = field()
    left: IRNode = field()
    right: IRNode = field()

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [("op", self.op), ("left", self.left), ("right", self.right)]

#------------------------------------------------------
# Join
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class Join(IRNode):
    target: IRNode = field()  # TableRef for simple tables, Alias for subqueries
    incoming: set[TableRef] = field(default_factory=set)
    filters: list[Filter] = field(default_factory=list)
    type: str = field(default="INNER")

    def _repr_fields(self) -> list[tuple[str, object]]:
        fields = [("incoming", self.incoming), ("target", self.target)]
        if self.filters:
            fields.append(("filters", self.filters))
        return fields

#------------------------------------------------------
# TableRef
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class TableRef(IRNode):
    """Reference to a table or view in FROM clause."""
    name: str = field()
    alias: str = field(default="")

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [("name", self.name), ("alias", self.alias)]

    def _hash_fields(self) -> list[tuple[str, object]]:
        return [("name", self.name), ("alias", self.alias)]

#------------------------------------------------------
# ColumnRef
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class ColumnRef(IRNode):
    table: TableRef = field()
    name: str = field()  # column name (required)
    column: ColumnInfo | None = field(default=None)  # optional ColumnInfo for type info

    def _repr_fields(self) -> list[tuple[str, object]]:
        table_alias = self.table.alias if self.table.alias else self.table.name
        return [("column", f"{table_alias}.{self.name}")]

    def _hash_fields(self) -> list[tuple[str, object]]:
        return [("table", self.table), ("name", self.name)]

#------------------------------------------------------
# Aggregate
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class AggregatePlaceholder(IRNode):
    ref: IRNode = field()
    op: str = field(default="COUNT")
    args: list[IRNode] = field(default_factory=list)
    body: Fragment = field(default_factory=Fragment)
    projection: list[ColumnRef] = field(default_factory=list)
    group: list[ColumnRef] = field(default_factory=list)
    optional: bool = field(default=False)

    def classify(self) -> str:
        return self.body.classify()

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [
            ("ref", self.ref),
            ("op", self.op),
            ("args", self.args),
            ("body", self.body),
            ("projection", self.projection),
            ("group", self.group),
            ("optional", self.optional),
        ]

#------------------------------------------------------
# Alias
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class Alias(IRNode):
    node: IRNode = field()
    alias: str = field()

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [("node", self.node), ("alias", self.alias)]

#------------------------------------------------------
# Lateral
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class Lateral(IRNode):
    """LATERAL subquery that can reference tables from the outer query."""
    node: IRNode = field()  # The subquery (usually an Alias wrapping a Fragment or UnionAll)

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [("node", self.node)]

#------------------------------------------------------
# And
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class And(IRNode):
    conditions: list[IRNode] = field(default_factory=list)

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [("conditions", self.conditions)]

#------------------------------------------------------
# Or
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class Or(IRNode):
    conditions: list[IRNode] = field(default_factory=list)

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [("conditions", self.conditions)]

#------------------------------------------------------
# CaseWhen
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class CaseWhen(IRNode):
    """CASE WHEN cond1 THEN result1 WHEN cond2 THEN result2 ... [ELSE else_result] END"""
    branches: list[tuple[IRNode, IRNode]] = field(default_factory=list)  # [(condition, result), ...]
    else_result: IRNode | None = field(default=None)

    def _repr_fields(self) -> list[tuple[str, object]]:
        fields: list[tuple[str, object]] = [("branches", self.branches)]
        if self.else_result is not None:
            fields.append(("else_result", self.else_result))
        return fields

#------------------------------------------------------
# Exists
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class Exists(IRNode):
    """EXISTS (subquery)"""
    subquery: IRNode = field()

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [("subquery", self.subquery)]

#------------------------------------------------------
# Not
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class Not(IRNode):
    """NOT (condition)"""
    condition: IRNode = field()

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [("condition", self.condition)]

#------------------------------------------------------
# Distinct
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class Distinct(IRNode):
    """DISTINCT wrapper for aggregate arguments."""
    node: IRNode = field()

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [("node", self.node)]

#------------------------------------------------------
# RowId
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class RowId(IRNode):
    """Reference to a table's rowid (for external tables)."""
    table: TableRef = field()

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [("table", self.table)]

#------------------------------------------------------
# Star
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class Star(IRNode):
    """Represents SELECT * or table.* in SQL."""
    table: TableRef | None = field(default=None)

    def _repr_fields(self) -> list[tuple[str, object]]:
        if self.table:
            return [("table", self.table.alias or self.table.name)]
        return []

#------------------------------------------------------
# Values
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class Values(IRNode):
    """VALUES clause with inline data."""
    data: DataFrame = field()
    alias: str = field(default="")
    columns: list[ColumnInfo] = field(default_factory=list)

    def _repr_fields(self) -> list[tuple[str, object]]:
        fields = []
        if self.alias:
            fields.append(("alias", self.alias))
        if self.columns:
            fields.append(("columns", [col.name for col in self.columns]))
        fields.append(("data", self.data))
        return fields

#------------------------------------------------------
# CreateView
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class CreateView(IRNode):
    """CREATE VIEW statement."""
    name: str = field()
    query: IRNode = field()  # Fragment, UnionAll, GroupByMax, etc.

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [("name", self.name), ("query", self.query)]

#------------------------------------------------------
# Union
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class Union(IRNode):
    """UNION (distinct) of multiple queries."""
    queries: list[Fragment] = field(default_factory=list)

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [("queries", self.queries)]

#------------------------------------------------------
# UnionAll
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class UnionAll(IRNode):
    """UNION ALL of multiple queries."""
    queries: list[Fragment] = field(default_factory=list)

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [("queries", self.queries)]

#------------------------------------------------------
# GroupByMax
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class GroupByMax(IRNode):
    """SELECT cols, MAX(other_cols) FROM subquery GROUP BY cols."""
    source: IRNode = field()  # Subquery (e.g., UnionAll)
    group_cols: list[ColumnRef] = field(default_factory=list)  # Columns to GROUP BY
    max_cols: list[ColumnRef] = field(default_factory=list)  # Columns to MAX()

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [
            ("source", self.source),
            ("group_cols", self.group_cols),
            ("max_cols", self.max_cols),
        ]


#------------------------------------------------------
# SemiNaiveLoop
#------------------------------------------------------

@dataclass(repr=False, eq=False)
class SemiNaiveTable(IRNode):
    """Per-table semi-naive loop spec."""
    table_name: str = field()
    base_query: IRNode = field()
    delta_query: IRNode = field()
    merge_query: IRNode = field()
    final_query: IRNode = field()
    accum_view_name: str = field()
    delta_view_name: str = field()
    next_delta_view_name: str = field()
    columns: list[str] = field()

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [
            ("table", self.table_name),
            ("base_query", self.base_query),
            ("delta_query", self.delta_query),
            ("merge_query", self.merge_query),
            ("final_query", self.final_query),
            ("accum_view_name", self.accum_view_name),
            ("delta_view_name", self.delta_view_name),
            ("next_delta_view_name", self.next_delta_view_name),
        ]

@dataclass(repr=False, eq=False)
class SemiNaiveLoop(IRNode):
    """Semi-naive fixpoint evaluation for a mutually recursive group."""
    tables: list[SemiNaiveTable] = field(default_factory=list)

    def _repr_fields(self) -> list[tuple[str, object]]:
        return [
            ("tables", self.tables),
        ]

#------------------------------------------------------
# IRRewriter
#------------------------------------------------------

T = TypeVar('T', bound=IRNode|None)
class IRRewriter:
    """
    Base class for IR tree transformations.

    Subclasses can override rewrite_<nodetype> methods to customize behavior.
    Default behavior recursively rewrites children and reconstructs nodes.

    Example usage:
        class MyRewriter(IRRewriter):
            def rewrite_columnref(self, node: ColumnRef) -> IRNode:
                # Custom handling for ColumnRef nodes
                return node

        rewriter = MyRewriter()
        new_tree = rewriter.rewrite(old_tree)
    """

    def rewrite(self, node: T) -> T:
        """Rewrite an IR node. Dispatches to type-specific methods."""
        if node is None:
            return None # type: ignore

        # Look for type-specific method (e.g., rewrite_expression)
        method_name = f'rewrite_{type(node).__name__.lower()}'
        method = getattr(self, method_name, None)
        if method:
            return method(node)

        # Fall back to default rewrite
        return self._default_rewrite(node)

    def _default_rewrite(self, node: IRNode) -> IRNode:
        """Default: recursively rewrite children and reconstruct node."""
        if isinstance(node, Fragment):
            return Fragment(
                selects=[self.rewrite(s) for s in node.selects],
                scans=[self.rewrite(s) for s in node.scans],
                joins=[self.rewrite(j) for j in node.joins],
                filters=[self.rewrite(f) for f in node.filters],
                group_by=[self.rewrite(g) for g in node.group_by],
                aggregates=[self.rewrite(a) for a in node.aggregates],
            )
        elif isinstance(node, Expression):
            return Expression(
                op=node.op,
                args=[self.rewrite(a) for a in node.args],
            )
        elif isinstance(node, TableFunction):
            return TableFunction(
                op=node.op,
                args=[self.rewrite(a) for a in node.args],
            )
        elif isinstance(node, WindowExpression):
            return WindowExpression(
                op=node.op,
                args=[self.rewrite(a) for a in node.args],
                partition_by=[self.rewrite(p) for p in node.partition_by],
                order_by=[self.rewrite(o) for o in node.order_by],
            )
        elif isinstance(node, Filter):
            return Filter(
                op=node.op,
                left=self.rewrite(node.left),
                right=self.rewrite(node.right),
            )
        elif isinstance(node, Join):
            return Join(
                target=self.rewrite(node.target),
                incoming={self.rewrite(t) for t in node.incoming},
                filters=[self.rewrite(f) for f in node.filters],
                type=node.type,
            )
        elif isinstance(node, AggregatePlaceholder):
            return AggregatePlaceholder(
                ref=self.rewrite(node.ref),
                op=node.op,
                args=[self.rewrite(a) for a in node.args],
                body=self.rewrite(node.body),
                projection=[self.rewrite(c) for c in node.projection],
                group=[self.rewrite(c) for c in node.group],
                optional=node.optional,
            )
        elif isinstance(node, Alias):
            return Alias(
                node=self.rewrite(node.node),
                alias=node.alias,
            )
        elif isinstance(node, Lateral):
            return Lateral(
                node=self.rewrite(node.node),
            )
        elif isinstance(node, And):
            return And(
                conditions=[self.rewrite(c) for c in node.conditions],
            )
        elif isinstance(node, Or):
            return Or(
                conditions=[self.rewrite(c) for c in node.conditions],
            )
        elif isinstance(node, CaseWhen):
            return CaseWhen(
                branches=[(self.rewrite(cond), self.rewrite(result)) for cond, result in node.branches],
                else_result=self.rewrite(node.else_result) if node.else_result is not None else None,
            )
        elif isinstance(node, Exists):
            return Exists(
                subquery=self.rewrite(node.subquery),
            )
        elif isinstance(node, Not):
            return Not(
                condition=self.rewrite(node.condition),
            )
        elif isinstance(node, Distinct):
            return Distinct(
                node=self.rewrite(node.node),
            )
        elif isinstance(node, CreateView):
            return CreateView(
                name=node.name,
                query=self.rewrite(node.query),
            )
        elif isinstance(node, Union):
            return Union(
                queries=[self.rewrite(q) for q in node.queries],
            )
        elif isinstance(node, UnionAll):
            return UnionAll(
                queries=[self.rewrite(q) for q in node.queries],
            )
        elif isinstance(node, GroupByMax):
            return GroupByMax(
                source=self.rewrite(node.source),
                group_cols=[self.rewrite(c) for c in node.group_cols],
                max_cols=[self.rewrite(c) for c in node.max_cols],
            )
        elif isinstance(node, SemiNaiveLoop):
            return SemiNaiveLoop(
                tables=[self.rewrite(t) for t in node.tables],
            )
        elif isinstance(node, SemiNaiveTable):
            return SemiNaiveTable(
                table_name=node.table_name,
                columns=list(node.columns),
                base_query=self.rewrite(node.base_query),
                delta_query=self.rewrite(node.delta_query),
                merge_query=self.rewrite(node.merge_query),
                final_query=self.rewrite(node.final_query),
                accum_view_name=node.accum_view_name,
                delta_view_name=node.delta_view_name,
                next_delta_view_name=node.next_delta_view_name,
            )
        elif isinstance(node, RowId):
            return RowId(
                table=self.rewrite(node.table),
            )
        # Leaf nodes (no children to rewrite): return as-is
        # Literal, TableRef, ColumnRef, ViewRef, Star
        return node

    # Override these in subclasses for custom behavior:
    # def rewrite_columnref(self, node: ColumnRef) -> IRNode: ...
    # def rewrite_expression(self, node: Expression) -> IRNode: ...
    # etc.
