"""
This program has been developed by students from the bachelor Computer Science at Utrecht
University within the Software Project course.
© Copyright Utrecht University (Department of Information and Computing Sciences)
"""

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.

    This class serves as the foundation for all AgentSpeak abstract syntax tree (AST) nodes.
    It defines the core interface that all AST nodes must implement to generate AgentSpeak code.
    """

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

        This method converts the AST node into its corresponding
        AgentSpeak source code representation.

        :return: The AgentSpeak code string representation of this node.
        """
        pass

    def __str__(self) -> str:
        """
        Returns the string representation of this AST node.

        This method provides a convenient way to get the AgentSpeak code representation
        by delegating to the _to_agentspeak method.

        :return: The AgentSpeak code string representation of this node.
        """
        return self._to_agentspeak()


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

    This class extends AstNode to provide common functionality for all expressions
    that can be used in logical operations within AgentSpeak programs.
    """

    def __and__(self, other: ExprCoalescible) -> AstBinaryOp:
        """
        Creates a logical AND operation between this expression and another.

        This method allows for operator overloading of the & operator to create
        binary logical operations in a more intuitive syntax.

        :param other: The right-hand side expression to combine with this one.
        :return: A new AstBinaryOp representing the logical AND operation.
        """
        return AstBinaryOp(self, BinaryOperatorType.AND, _coalesce_expr(other))

    def __or__(self, other: ExprCoalescible) -> AstBinaryOp:
        """
        Creates a logical OR operation between this expression and another.

        This method allows for operator overloading of the | operator to create
        binary logical operations in a more intuitive syntax.

        :param other: The right-hand side expression to combine with this one.
        :return: A new AstBinaryOp representing the logical OR operation.
        """
        return AstBinaryOp(self, BinaryOperatorType.OR, _coalesce_expr(other))

    def __invert__(self) -> AstLogicalExpression:
        """
        Creates a logical negation of this expression.

        This method allows for operator overloading of the ~ operator to create
        negated expressions. If the expression is already a logical expression,
        it toggles the negation flag. Otherwise, it wraps the expression in a
        new AstLogicalExpression with negation set to True.

        :return: An AstLogicalExpression representing the negated form of this expression.
        """
        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).

    Atoms are the simplest form of terms in AgentSpeak, representing concrete,
    unchanging values. They are typically used as constants in logical expressions.

    :ivar value: The string value of this atom, which will be converted to lowercase
                 in the AgentSpeak representation.
    """

    value: str

    def _to_agentspeak(self) -> str:
        """
        Converts this atom to its AgentSpeak string representation.

        Atoms are represented in lowercase in AgentSpeak to distinguish them
        from variables (which are capitalized).

        :return: The lowercase string representation of this atom.
        """
        return self.value.lower()


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

    Variables in AgentSpeak are placeholders that can be bound to specific values
    during execution. They are distinguished from atoms by their capitalization.

    :ivar name: The name of this variable, which will be capitalized in the
                AgentSpeak representation.
    """

    name: str

    def _to_agentspeak(self) -> str:
        """
        Converts this variable to its AgentSpeak string representation.

        Variables are represented with capitalized names in AgentSpeak to distinguish
        them from atoms (which are lowercase).

        :return: The capitalized string representation of this variable.
        """
        return self.name.capitalize()


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

    Numeric constants can be either integers or floating-point numbers and are
    used in logical expressions and comparisons.

    :ivar value: The numeric value of this constant (can be int or float).
    """

    value: int | float

    def _to_agentspeak(self) -> str:
        """
        Converts this numeric constant to its AgentSpeak string representation.

        :return: The string representation of the numeric value.
        """
        return str(self.value)


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

    String literals are used to represent textual data and are enclosed in
    double quotes in the AgentSpeak representation.

    :ivar value: The string content of this literal.
    """

    value: str

    def _to_agentspeak(self) -> str:
        """
        Converts this string literal to its AgentSpeak string representation.

        String literals are enclosed in double quotes in AgentSpeak.

        :return: The string literal enclosed in double quotes.
        """
        return f'"{self.value}"'


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

    Literals are the fundamental building blocks of AgentSpeak programs, consisting
    of a functor (predicate name) and an optional list of terms (arguments).

    :ivar functor: The name of the predicate or function.
    :ivar terms: A list of terms (arguments) for this literal. Defaults to an empty list.
    """

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

    def _to_agentspeak(self) -> str:
        """
        Converts this literal to its AgentSpeak string representation.

        If the literal has no terms, it returns just the functor name.
        Otherwise, it returns the functor followed by the terms in parentheses.

        :return: The AgentSpeak string representation of this literal.
        """
        if not self.terms:
            return self.functor
        args = ", ".join(map(str, self.terms))
        return f"{self.functor}({args})"


class BinaryOperatorType(StrEnum):
    """
    Enumeration of binary operator types used in AgentSpeak expressions.

    These operators are used to create binary operations between expressions,
    including logical operations (AND, OR) and comparison operations.
    """

    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.

    Binary operations combine two expressions using a logical or comparison operator.
    They are used to create complex logical conditions in AgentSpeak programs.

    :ivar left: The left-hand side expression of the operation.
    :ivar operator: The binary operator type (AND, OR, comparison operators, etc.).
    :ivar right: The right-hand side expression of the operation.
    """

    left: AstExpression
    operator: BinaryOperatorType
    right: AstExpression

    def __post_init__(self):
        """
        Post-initialization processing to ensure proper expression types.

        This method wraps the left and right expressions in AstLogicalExpression
        instances if they aren't already, ensuring consistent handling throughout
        the AST.
        """
        self.left = _as_logical(self.left)
        self.right = _as_logical(self.right)

    def _to_agentspeak(self) -> str:
        """
        Converts this binary operation to its AgentSpeak string representation.

        The method handles proper parenthesization of sub-expressions to maintain
        correct operator precedence and readability.

        :return: The AgentSpeak string representation of this binary operation.
        """
        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.

    Logical expressions can be either positive or negated and form the basis
    of conditions and beliefs in AgentSpeak programs.

    :ivar expression: The underlying expression being evaluated.
    :ivar negated: Boolean flag indicating whether this expression is negated.
    """

    expression: AstExpression
    negated: bool = False

    def _to_agentspeak(self) -> str:
        """
        Converts this logical expression to its AgentSpeak string representation.

        If the expression is negated, it prepends 'not ' to the expression string.
        For complex expressions (binary operations), it adds parentheses when negated
        to maintain correct logical interpretation.

        :return: The AgentSpeak string representation of this logical expression.
        """
        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:
    """
    Converts an expression to a logical expression if it isn't already.

    This helper function ensures that expressions are properly wrapped in
    AstLogicalExpression instances, which is necessary for consistent handling
    of logical operations in the AST.

    :param expr: The expression to convert.
    :return: The expression wrapped in an AstLogicalExpression if it wasn't already.
    """
    if isinstance(expr, AstLogicalExpression):
        return expr
    return AstLogicalExpression(expr)


class StatementType(StrEnum):
    """
    Enumeration of statement types that can appear in AgentSpeak plans.

    These statement types define the different kinds of actions and operations
    that can be performed within the body of an AgentSpeak plan.
    """

    EMPTY = ""
    """Empty statement (no operation, used when evaluating a plan to true)."""

    DO_ACTION = "."
    """Execute an action defined in Python."""

    ACHIEVE_GOAL = "!"
    """Achieve a goal (add a goal to be accomplished)."""

    TEST_GOAL = "?"
    """Test a goal (check if a goal can be achieved)."""

    ADD_BELIEF = "+"
    """Add a belief to the belief base."""

    REMOVE_BELIEF = "-"
    """Remove a belief from the belief base."""

    REPLACE_BELIEF = "-+"
    """Replace a belief in the belief base."""


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

    Statements are the executable units within AgentSpeak plans. They consist
    of a statement type (defining the operation) and an expression (defining
    what to operate on).

    :ivar type: The type of statement (action, goal, belief operation, etc.).
    :ivar expression: The expression that this statement operates on.
    """

    type: StatementType
    expression: AstExpression

    def _to_agentspeak(self) -> str:
        """
        Converts this statement to its AgentSpeak string representation.

        The representation consists of the statement type prefix followed by
        the expression.

        :return: The AgentSpeak string representation of this statement.
        """
        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.

    Rules define logical implications in AgentSpeak programs. They consist of a
    result (conclusion) and an optional condition (premise). When the condition
    holds, the result is inferred to be true.

    :ivar result: The conclusion or result of this rule.
    :ivar condition: The premise or condition for this rule (optional).
    """

    result: AstExpression
    condition: AstExpression | None = None

    def __post_init__(self):
        """
        Post-initialization processing to ensure proper expression types.

        If a condition is provided, this method wraps it in an AstLogicalExpression
        to ensure consistent handling throughout the AST.
        """
        if self.condition is not None:
            self.condition = _as_logical(self.condition)

    def _to_agentspeak(self) -> str:
        """
        Converts this rule to its AgentSpeak string representation.

        If no condition is specified, the rule is represented as a simple fact.
        If a condition is specified, it's represented as an implication (result :- condition).

        :return: The AgentSpeak string representation of this rule.
        """
        if not self.condition:
            return f"{self.result}."
        return f"{self.result} :- {self.condition}."


class TriggerType(StrEnum):
    """
    Enumeration of trigger types for AgentSpeak plans.

    Trigger types define what kind of events can activate an AgentSpeak plan.
    Currently, the system supports triggers for added beliefs and added goals.
    """

    ADDED_BELIEF = "+"
    """Trigger when a belief is added to the belief base."""

    # REMOVED_BELIEF = "-"  # TODO
    # MODIFIED_BELIEF = "^"  # TODO

    ADDED_GOAL = "+!"
    """Trigger when a goal is added to be achieved."""

    # REMOVED_GOAL = "-!"  # TODO


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

    Plans define the reactive behavior of agents in AgentSpeak. They specify what
    actions to take when certain conditions are met (trigger and context).

    :ivar type: The type of trigger that activates this plan.
    :ivar trigger_literal: The specific event or condition that triggers this plan.
    :ivar context: A list of conditions that must hold for this plan to be applicable.
    :ivar body: A list of statements to execute when this plan is triggered.
    """

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

    def _to_agentspeak(self) -> str:
        """
        Converts this plan to its AgentSpeak string representation.

        The representation follows the standard AgentSpeak plan format:
        trigger_type + trigger_literal
        : context_conditions
        <- body_statements.

        :return: The AgentSpeak string representation of this plan.
        """
        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.

    This is the root node of the AgentSpeak AST, containing all the rules
    and plans that define the agent's behavior.

    :ivar rules: A list of inference rules in this program.
    :ivar plans: A list of reactive plans in this program.
    """

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

    def _to_agentspeak(self) -> str:
        """
        Converts this program to its AgentSpeak string representation.

        The representation consists of all rules followed by all plans,
        separated by blank lines for readability.

        :return: The complete AgentSpeak source code for this program.
        """
        lines = []
        lines.extend(map(str, self.rules))

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

        return "\n".join(lines)
