//! Abstract Syntax Tree for the verification DSL.
//!
//! The DSL supports:
//! - Contract invariants (properties that hold across all states)
//! - Function pre/post conditions
//! - Temporal properties (eventually, always)

use serde::{Deserialize, Serialize};

use crate::VarType;

/// A parsed contract with specifications.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Contract {
    /// Contract name.
    pub name: String,
    /// State variables.
    pub state_vars: Vec<StateVar>,
    /// Functions.
    pub functions: Vec<Function>,
    /// Specifications.
    pub specs: Vec<Specification>,
}

/// State variable declaration.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct StateVar {
    /// Variable name.
    pub name: String,
    /// Variable type.
    pub var_type: VarType,
    /// Initial value expression (optional).
    pub init: Option<Expression>,
}

/// Function definition.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Function {
    /// Function name.
    pub name: String,
    /// Parameters.
    pub params: Vec<(String, VarType)>,
    /// Return type.
    pub returns: Option<VarType>,
    /// Function body.
    pub body: Vec<Statement>,
    /// Pre-conditions.
    pub requires: Vec<Expression>,
    /// Post-conditions.
    pub ensures: Vec<Expression>,
}

/// Statement in function body.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum Statement {
    /// Variable assignment.
    Assign {
        target: String,
        value: Expression,
    },
    /// If-then-else.
    If {
        condition: Expression,
        then_branch: Vec<Statement>,
        else_branch: Vec<Statement>,
    },
    /// While loop.
    While {
        condition: Expression,
        invariant: Option<Expression>,
        body: Vec<Statement>,
    },
    /// For loop.
    For {
        init: Box<Statement>,
        condition: Expression,
        update: Box<Statement>,
        body: Vec<Statement>,
    },
    /// Return statement.
    Return(Option<Expression>),
    /// Require (revert if false).
    Require(Expression, Option<String>),
    /// Assert (for verification).
    Assert(Expression),
    /// Assume (hint to verifier).
    Assume(Expression),
    /// Emit event.
    Emit {
        event: String,
        args: Vec<Expression>,
    },
    /// External call.
    Call {
        target: Expression,
        function: String,
        args: Vec<Expression>,
        value: Option<Expression>,
    },
    /// Block of statements.
    Block(Vec<Statement>),
}

/// Expression in the DSL.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum Expression {
    /// Literal value.
    Literal(Literal),
    /// Variable reference.
    Variable(String),
    /// Binary operation.
    BinaryOp {
        left: Box<Expression>,
        op: BinaryOperator,
        right: Box<Expression>,
    },
    /// Unary operation.
    UnaryOp {
        op: UnaryOperator,
        operand: Box<Expression>,
    },
    /// Ternary conditional.
    Ternary {
        condition: Box<Expression>,
        then_expr: Box<Expression>,
        else_expr: Box<Expression>,
    },
    /// Function call.
    FunctionCall {
        function: String,
        args: Vec<Expression>,
    },
    /// Member access (e.g., msg.sender).
    MemberAccess {
        object: Box<Expression>,
        member: String,
    },
    /// Array/mapping index.
    Index {
        array: Box<Expression>,
        index: Box<Expression>,
    },
    /// Quantifier (forall, exists).
    Quantifier {
        kind: QuantifierKind,
        var: String,
        var_type: VarType,
        body: Box<Expression>,
    },
    /// Old value (in post-conditions).
    Old(Box<Expression>),
    /// Result value (in post-conditions).
    Result,
    /// Sum over range (for array properties).
    Sum {
        var: String,
        range: Box<Expression>,
        body: Box<Expression>,
    },
}

/// Literal values.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum Literal {
    Bool(bool),
    Uint(u128),
    Int(i128),
    Address([u8; 20]),
    Bytes(Vec<u8>),
    String(String),
}

/// Binary operators.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum BinaryOperator {
    // Arithmetic
    Add,
    Sub,
    Mul,
    Div,
    Mod,
    Exp,
    // Comparison
    Eq,
    Ne,
    Lt,
    Le,
    Gt,
    Ge,
    // Logical
    And,
    Or,
    Implies,
    Iff,
    // Bitwise
    BitAnd,
    BitOr,
    BitXor,
    Shl,
    Shr,
}

impl BinaryOperator {
    /// Returns true if this is a comparison operator.
    pub fn is_comparison(&self) -> bool {
        matches!(
            self,
            BinaryOperator::Eq
                | BinaryOperator::Ne
                | BinaryOperator::Lt
                | BinaryOperator::Le
                | BinaryOperator::Gt
                | BinaryOperator::Ge
        )
    }

    /// Returns true if this is a logical operator.
    pub fn is_logical(&self) -> bool {
        matches!(
            self,
            BinaryOperator::And | BinaryOperator::Or | BinaryOperator::Implies | BinaryOperator::Iff
        )
    }
}

/// Unary operators.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum UnaryOperator {
    Not,
    Neg,
    BitNot,
}

/// Quantifier kinds.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum QuantifierKind {
    ForAll,
    Exists,
}

/// Specification types.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum Specification {
    /// Contract invariant.
    Invariant(Invariant),
    /// Temporal or custom property.
    Property(Property),
    /// Function annotation.
    Annotation(Annotation),
}

/// Contract invariant.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Invariant {
    /// Invariant name.
    pub name: String,
    /// The invariant expression.
    pub expr: Expression,
    /// Description.
    pub description: Option<String>,
}

/// Custom property.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Property {
    /// Property name.
    pub name: String,
    /// Property kind.
    pub kind: PropertyKind,
    /// The property expression.
    pub expr: Expression,
    /// Description.
    pub description: Option<String>,
}

/// Property kinds.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum PropertyKind {
    /// Safety property (always holds).
    Safety,
    /// Liveness property (eventually holds).
    Liveness,
    /// Reachability (some state is reachable).
    Reachability,
    /// Custom property.
    Custom,
}

/// Function annotation.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Annotation {
    /// Function name this annotation applies to.
    pub function: String,
    /// Pre-conditions.
    pub requires: Vec<Expression>,
    /// Post-conditions.
    pub ensures: Vec<Expression>,
    /// Modifies clause.
    pub modifies: Vec<String>,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_binary_operator_classification() {
        assert!(BinaryOperator::Eq.is_comparison());
        assert!(BinaryOperator::Lt.is_comparison());
        assert!(!BinaryOperator::Add.is_comparison());

        assert!(BinaryOperator::And.is_logical());
        assert!(BinaryOperator::Implies.is_logical());
        assert!(!BinaryOperator::Mul.is_logical());
    }

    #[test]
    fn test_expression_construction() {
        let expr = Expression::BinaryOp {
            left: Box::new(Expression::Variable("x".to_string())),
            op: BinaryOperator::Add,
            right: Box::new(Expression::Literal(Literal::Uint(1))),
        };

        if let Expression::BinaryOp { op, .. } = expr {
            assert_eq!(op, BinaryOperator::Add);
        } else {
            panic!("Expected BinaryOp");
        }
    }

    #[test]
    fn test_invariant_creation() {
        let inv = Invariant {
            name: "positive_balance".to_string(),
            expr: Expression::BinaryOp {
                left: Box::new(Expression::Variable("balance".to_string())),
                op: BinaryOperator::Ge,
                right: Box::new(Expression::Literal(Literal::Uint(0))),
            },
            description: Some("Balance is always non-negative".to_string()),
        };

        assert_eq!(inv.name, "positive_balance");
    }
}
