"""
PIT Formula Base - Core Expression Classes
============================================
Contains the foundational formula expression classes used by both
pit_keywords.py (PITKeywordsMixin) and pit_formula.py (F class, DateTable).

This module exists to break the circular dependency:
  pit_formula.py → pit_keywords.py → pit_formula.py

Instead, both modules import from this shared base:
  pit_formula_base.py ← pit_keywords.py
  pit_formula_base.py ← pit_formula.py
"""

from abc import ABC, abstractmethod
from typing import Union

# ==================== Formula Expression System ====================


class FormulaExpression(ABC):
    """Base class for formula expressions that can be combined"""

    @abstractmethod
    def to_string(self) -> str:
        pass

    def __str__(self) -> str:
        return self.to_string()

    # Arithmetic operators
    def __add__(self, other) -> "BinaryOp":
        return BinaryOp(self, "+", _wrap(other))

    def __radd__(self, other) -> "BinaryOp":
        return BinaryOp(_wrap(other), "+", self)

    def __sub__(self, other) -> "BinaryOp":
        return BinaryOp(self, "-", _wrap(other))

    def __rsub__(self, other) -> "BinaryOp":
        return BinaryOp(_wrap(other), "-", self)

    def __mul__(self, other) -> "BinaryOp":
        return BinaryOp(self, "*", _wrap(other))

    def __rmul__(self, other) -> "BinaryOp":
        return BinaryOp(_wrap(other), "*", self)

    def __truediv__(self, other) -> "BinaryOp":
        return BinaryOp(self, "/", _wrap(other))

    def __rtruediv__(self, other) -> "BinaryOp":
        return BinaryOp(_wrap(other), "/", self)

    def __neg__(self) -> "UnaryOp":
        return UnaryOp("-", self)

    # Comparison operators
    def __gt__(self, other) -> "BinaryOp":
        return BinaryOp(self, ">", _wrap(other))

    def __lt__(self, other) -> "BinaryOp":
        return BinaryOp(self, "<", _wrap(other))

    def __ge__(self, other) -> "BinaryOp":
        return BinaryOp(self, ">=", _wrap(other))

    def __le__(self, other) -> "BinaryOp":
        return BinaryOp(self, "<=", _wrap(other))

    def __eq__(self, other) -> "BinaryOp":
        return BinaryOp(self, "==", _wrap(other))

    def __ne__(self, other) -> "BinaryOp":
        return BinaryOp(self, "!=", _wrap(other))


class Literal(FormulaExpression):
    """Represents a literal value (number or string)"""

    def __init__(self, value: Union[int, float, str]):
        self.value = value

    def to_string(self) -> str:
        if isinstance(self.value, str) and not self.value.replace(".", "").replace("-", "").replace("_", "").isdigit():
            return self.value
        return str(self.value)


class FunctionCall(FormulaExpression):
    """Represents a function call"""

    def __init__(
        self,
        func_name: str,
        *args: Union[FormulaExpression, str, int, float],
        _braceless: bool = False,
        _infix: bool = False,
    ):
        self.func_name = func_name
        self.args = [_wrap(arg) for arg in args]
        self._braceless = _braceless  # For keywords like Pay/Receive that can be used without braces
        self._infix = _infix  # For infix operators like And/Or

    def to_string(self) -> str:
        if self._infix:
            # Infix operator: "A And B" or "A Or B Or C"
            if len(self.args) < 2:
                raise ValueError(f"Infix operator {self.func_name} requires at least 2 arguments")
            operator = f" {self.func_name} "
            return operator.join(arg.to_string() for arg in self.args)
        elif self._braceless:
            if not self.args:
                # Brace-less operator with no args: "And", "Or" (legacy, shouldn't be used)
                return self.func_name
            else:
                # Brace-less prefix keyword with args: "Pay amount", "Receive amount"
                args_str = " ".join(arg.to_string() for arg in self.args)
                return f"{self.func_name} {args_str}"
        # Regular function call with parentheses
        args_str = ", ".join(arg.to_string() for arg in self.args)
        return f"{self.func_name}({args_str})"


class BinaryOp(FormulaExpression):
    """Represents a binary operation"""

    def __init__(self, left: FormulaExpression, op: str, right: FormulaExpression):
        self.left = left
        self.op = op
        self.right = right

    def to_string(self) -> str:
        # Add parentheses for clarity
        left_str = self.left.to_string()
        right_str = self.right.to_string()

        # Don't double-wrap already parenthesized expressions
        if not (left_str.startswith("(") and left_str.endswith(")")):
            if isinstance(self.left, BinaryOp):
                left_str = f"({left_str})"
        if not (right_str.startswith("(") and right_str.endswith(")")):
            if isinstance(self.right, BinaryOp):
                right_str = f"({right_str})"

        return f"{left_str} {self.op} {right_str}"


class UnaryOp(FormulaExpression):
    """Represents a unary operation"""

    def __init__(self, op: str, operand: FormulaExpression):
        self.op = op
        self.operand = operand

    def to_string(self) -> str:
        return f"{self.op}({self.operand.to_string()})"


# ==================== Braceless Keywords Registry ====================
# Some PIT keywords can be used without parentheses (typically operators)
# These are exceptions to the "always use ()" rule
BRACELESS_KEYWORDS = {
    "and_",  # && - Logical AND infix operator (PEP 8 trailing underscore)
    "or_",  # || - Logical OR infix operator (PEP 8 trailing underscore)
    "pay",  # Pay - prefix keyword without braces: Pay payment_flow
    "receive",  # Receive - prefix keyword without braces: Receive payment_flow
    # Add more braceless keywords here as needed
}


def is_braceless_keyword(func) -> bool:
    """Check if a function is a registered braceless keyword."""
    if not callable(func):
        return False
    func_name = getattr(func, "__name__", "")
    return func_name in BRACELESS_KEYWORDS


def _wrap(value) -> FormulaExpression:
    """Wrap a value in a FormulaExpression if it isn't already"""
    if isinstance(value, FormulaExpression):
        return value

    # Check if user forgot to call a PIT keyword method (missing parentheses)
    if callable(value):
        # Check if it's a PIT keyword method
        func_name = getattr(value, "__name__", str(value))

        # Allow braceless keywords - these can be passed without calling
        # (e.g., And, Or as operators that don't need parentheses in PIT syntax)
        if is_braceless_keyword(value):
            # Call the keyword to get its FunctionCall (they take no args)
            return value()

        if "PITKeywordsMixin" in str(value) or hasattr(value, "__self__"):
            raise TypeError(
                f"PIT keyword method '{func_name}' was not called. "
                f"Did you forget the parentheses ()? \n"
                f"Use: F.{func_name}() instead of F.{func_name}"
            )
        else:
            raise TypeError(
                f"Cannot use callable '{func_name}' in formula. "
                f"If this is a PIT keyword, make sure to call it with parentheses ()."
            )

    # Handle Enum types - extract .value for clean string output
    from enum import Enum

    if isinstance(value, Enum):
        return Literal(value.value)
    return Literal(value)


# Control what gets imported with "from pit_formula_base import *"
__all__ = [
    "FormulaExpression",
    "Literal",
    "FunctionCall",
    "BinaryOp",
    "UnaryOp",
    "BRACELESS_KEYWORDS",
    "is_braceless_keyword",
    "_wrap",
]
