from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import StrEnum


class AstNode(ABC):
    """
    Abstract base class for all elements of an AgentSpeak program.
    """

    @abstractmethod
    def _to_agentspeak(self) -> str:
        """
        Generates the AgentSpeak code string.
        """
        pass

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


class AstExpression(AstNode, ABC):
    """
    Intermediate class for anything that can be used in a logical expression.
    """

    def __and__(self, other: ExprCoalescible) -> AstBinaryOp:
        return AstBinaryOp(self, BinaryOperatorType.AND, _coalesce_expr(other))

    def __or__(self, other: ExprCoalescible) -> AstBinaryOp:
        return AstBinaryOp(self, BinaryOperatorType.OR, _coalesce_expr(other))

    def __invert__(self) -> AstLogicalExpression:
        if isinstance(self, AstLogicalExpression):
            self.negated = not self.negated
            return self
        return AstLogicalExpression(self, negated=True)


type ExprCoalescible = AstExpression | str | int | float


def _coalesce_expr(value: ExprCoalescible) -> AstExpression:
    if isinstance(value, AstExpression):
        return value
    if isinstance(value, str):
        return AstString(value)
    if isinstance(value, (int, float)):
        return AstNumber(value)
    raise TypeError(f"Cannot coalesce type {type(value)} into an AstTerm.")


@dataclass
class AstTerm(AstExpression, ABC):
    """
    Base class for terms appearing inside literals.
    """

    def __ge__(self, other: ExprCoalescible) -> AstBinaryOp:
        return AstBinaryOp(self, BinaryOperatorType.GREATER_EQUALS, _coalesce_expr(other))

    def __gt__(self, other: ExprCoalescible) -> AstBinaryOp:
        return AstBinaryOp(self, BinaryOperatorType.GREATER_THAN, _coalesce_expr(other))

    def __le__(self, other: ExprCoalescible) -> AstBinaryOp:
        return AstBinaryOp(self, BinaryOperatorType.LESS_EQUALS, _coalesce_expr(other))

    def __lt__(self, other: ExprCoalescible) -> AstBinaryOp:
        return AstBinaryOp(self, BinaryOperatorType.LESS_THAN, _coalesce_expr(other))

    def __eq__(self, other: ExprCoalescible) -> AstBinaryOp:
        return AstBinaryOp(self, BinaryOperatorType.EQUALS, _coalesce_expr(other))

    def __ne__(self, other: ExprCoalescible) -> AstBinaryOp:
        return AstBinaryOp(self, BinaryOperatorType.NOT_EQUALS, _coalesce_expr(other))


@dataclass(eq=False)
class AstAtom(AstTerm):
    """
    Represents a grounded atom in AgentSpeak (e.g., lowercase constants).
    """

    value: str

    def _to_agentspeak(self) -> str:
        return self.value.lower()


@dataclass(eq=False)
class AstVar(AstTerm):
    """
    Represents an ungrounded variable in AgentSpeak (e.g., capitalized names).
    """

    name: str

    def _to_agentspeak(self) -> str:
        return self.name.capitalize()


@dataclass(eq=False)
class AstNumber(AstTerm):
    """
    Represents a numeric constant in AgentSpeak.
    """

    value: int | float

    def _to_agentspeak(self) -> str:
        return str(self.value)


@dataclass(eq=False)
class AstString(AstTerm):
    """
    Represents a string literal in AgentSpeak.
    """

    value: str

    def _to_agentspeak(self) -> str:
        return f'"{self.value}"'


@dataclass(eq=False)
class AstLiteral(AstTerm):
    """
    Represents a literal (functor and terms) in AgentSpeak.
    """

    functor: str
    terms: list[AstTerm] = field(default_factory=list)

    def _to_agentspeak(self) -> str:
        if not self.terms:
            return self.functor
        args = ", ".join(map(str, self.terms))
        return f"{self.functor}({args})"


class BinaryOperatorType(StrEnum):
    AND = "&"
    OR = "|"
    GREATER_THAN = ">"
    LESS_THAN = "<"
    EQUALS = "=="
    NOT_EQUALS = "\\=="
    GREATER_EQUALS = ">="
    LESS_EQUALS = "<="


@dataclass
class AstBinaryOp(AstExpression):
    """
    Represents a binary logical or relational operation in AgentSpeak.
    """

    left: AstExpression
    operator: BinaryOperatorType
    right: AstExpression

    def __post_init__(self):
        self.left = _as_logical(self.left)
        self.right = _as_logical(self.right)

    def _to_agentspeak(self) -> str:
        l_str = str(self.left)
        r_str = str(self.right)

        assert isinstance(self.left, AstLogicalExpression)
        assert isinstance(self.right, AstLogicalExpression)

        if isinstance(self.left.expression, AstBinaryOp) or self.left.negated:
            l_str = f"({l_str})"
        if isinstance(self.right.expression, AstBinaryOp) or self.right.negated:
            r_str = f"({r_str})"

        return f"{l_str} {self.operator.value} {r_str}"


@dataclass
class AstLogicalExpression(AstExpression):
    """
    Represents a logical expression, potentially negated, in AgentSpeak.
    """

    expression: AstExpression
    negated: bool = False

    def _to_agentspeak(self) -> str:
        expr_str = str(self.expression)
        if isinstance(self.expression, AstBinaryOp) and self.negated:
            expr_str = f"({expr_str})"
        return f"{'not ' if self.negated else ''}{expr_str}"


def _as_logical(expr: AstExpression) -> AstLogicalExpression:
    if isinstance(expr, AstLogicalExpression):
        return expr
    return AstLogicalExpression(expr)


class StatementType(StrEnum):
    EMPTY = ""
    DO_ACTION = "."
    ACHIEVE_GOAL = "!"
    TEST_GOAL = "?"
    ADD_BELIEF = "+"
    REMOVE_BELIEF = "-"
    REPLACE_BELIEF = "-+"


@dataclass
class AstStatement(AstNode):
    """
    A statement that can appear inside a plan.
    """

    type: StatementType
    expression: AstExpression

    def _to_agentspeak(self) -> str:
        return f"{self.type.value}{self.expression}"


@dataclass
class AstRule(AstNode):
    """
    Represents an inference rule in AgentSpeak. If there is no condition, it always holds.
    """

    result: AstExpression
    condition: AstExpression | None = None

    def __post_init__(self):
        if self.condition is not None:
            self.condition = _as_logical(self.condition)

    def _to_agentspeak(self) -> str:
        if not self.condition:
            return f"{self.result}."
        return f"{self.result} :- {self.condition}."


class TriggerType(StrEnum):
    ADDED_BELIEF = "+"
    # REMOVED_BELIEF = "-"  # TODO
    # MODIFIED_BELIEF = "^"  # TODO
    ADDED_GOAL = "+!"
    # REMOVED_GOAL = "-!"  # TODO


@dataclass
class AstPlan(AstNode):
    """
    Represents a plan in AgentSpeak, consisting of a trigger, context, and body.
    """

    type: TriggerType
    trigger_literal: AstExpression
    context: list[AstExpression]
    body: list[AstStatement]

    def _to_agentspeak(self) -> str:
        assert isinstance(self.trigger_literal, AstLiteral)

        indent = " " * 6
        colon = "  :   "
        arrow = "  <-  "

        lines = []

        lines.append(f"{self.type.value}{self.trigger_literal}")

        if self.context:
            lines.append(colon + f" &\n{indent}".join(str(c) for c in self.context))

        if self.body:
            lines.append(arrow + f";\n{indent}".join(str(s) for s in self.body) + ".")

        lines.append("")

        return "\n".join(lines)


@dataclass
class AstProgram(AstNode):
    """
    Represents a full AgentSpeak program, consisting of rules and plans.
    """

    rules: list[AstRule] = field(default_factory=list)
    plans: list[AstPlan] = field(default_factory=list)

    def _to_agentspeak(self) -> str:
        lines = []
        lines.extend(map(str, self.rules))

        lines.extend(["", ""])
        lines.extend(map(str, self.plans))

        return "\n".join(lines)