//! SMT solver integration.
//!
//! Provides an interface to SMT solvers for constraint solving.

use crate::ast::{BinaryOperator, Expression, Literal, QuantifierKind, UnaryOperator};
use crate::error::VerifierResult;
use crate::symbolic::{SymbolicState, SymbolicValue};
use crate::VarType;

/// SMT solver result.
pub enum SmtResult {
    /// Formula is satisfiable with model.
    Sat(Vec<(String, String)>),
    /// Formula is unsatisfiable.
    Unsat,
    /// Unknown (timeout, etc.).
    Unknown(String),
}

/// Encoded SMT expression (string representation).
#[derive(Clone, Debug)]
pub struct SmtExpr(String);

impl SmtExpr {
    pub fn new(s: &str) -> Self {
        SmtExpr(s.to_string())
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }
}

/// SMT solver context.
pub struct SmtContext {
    /// Timeout in milliseconds.
    timeout_ms: u64,
    /// Variable declarations.
    declarations: Vec<String>,
    /// Assertions.
    assertions: Vec<String>,
}

impl SmtContext {
    /// Creates a new SMT context.
    pub fn new(timeout_ms: u64) -> Self {
        Self {
            timeout_ms,
            declarations: Vec::new(),
            assertions: Vec::new(),
        }
    }

    /// Get the timeout in milliseconds
    pub fn timeout_ms(&self) -> u64 {
        self.timeout_ms
    }

    /// Add a variable declaration
    pub fn declare(&mut self, name: &str, sort: &str) {
        self.declarations.push(format!("(declare-fun {} () {})", name, sort));
    }

    /// Add an assertion
    pub fn assert(&mut self, expr: &SmtExpr) {
        self.assertions.push(format!("(assert {})", expr.as_str()));
    }

    /// Get all declarations
    pub fn declarations(&self) -> &[String] {
        &self.declarations
    }

    /// Get all assertions
    pub fn assertions(&self) -> &[String] {
        &self.assertions
    }

    /// Generate SMT-LIB script
    pub fn to_smtlib(&self) -> String {
        let mut script = String::new();
        script.push_str("; SMT-LIB2 script\n");
        script.push_str(&format!("; Timeout: {}ms\n", self.timeout_ms));
        script.push_str("(set-logic ALL)\n");

        for decl in &self.declarations {
            script.push_str(decl);
            script.push('\n');
        }

        for assertion in &self.assertions {
            script.push_str(assertion);
            script.push('\n');
        }

        script.push_str("(check-sat)\n");
        script.push_str("(get-model)\n");
        script
    }

    /// Clear all declarations and assertions
    pub fn reset(&mut self) {
        self.declarations.clear();
        self.assertions.clear();
    }

    /// Encodes a symbolic state as SMT constraints.
    pub fn encode_state(&self, state: &SymbolicState) -> VerifierResult<Vec<SmtExpr>> {
        let mut constraints = Vec::new();

        // Encode path condition
        for cond in state.path_condition() {
            let encoded = self.encode_symbolic_value(cond)?;
            constraints.push(encoded);
        }

        Ok(constraints)
    }

    /// Encodes an AST expression as SMT.
    pub fn encode_expression(&self, expr: &Expression) -> VerifierResult<SmtExpr> {
        let smt_str = self.expr_to_smt(expr)?;
        Ok(SmtExpr::new(&smt_str))
    }

    /// Converts an AST expression to SMT-LIB format.
    fn expr_to_smt(&self, expr: &Expression) -> VerifierResult<String> {
        match expr {
            Expression::Literal(lit) => self.literal_to_smt(lit),

            Expression::Variable(name) => Ok(self.sanitize_name(name)),

            Expression::BinaryOp { left, op, right } => {
                let l = self.expr_to_smt(left)?;
                let r = self.expr_to_smt(right)?;
                let op_str = self.binop_to_smt(*op);
                Ok(format!("({} {} {})", op_str, l, r))
            }

            Expression::UnaryOp { op, operand } => {
                let v = self.expr_to_smt(operand)?;
                let op_str = self.unop_to_smt(*op);
                Ok(format!("({} {})", op_str, v))
            }

            Expression::Ternary {
                condition,
                then_expr,
                else_expr,
            } => {
                let c = self.expr_to_smt(condition)?;
                let t = self.expr_to_smt(then_expr)?;
                let e = self.expr_to_smt(else_expr)?;
                Ok(format!("(ite {} {} {})", c, t, e))
            }

            Expression::Quantifier {
                kind,
                var,
                var_type,
                body,
            } => {
                let sort = self.type_to_smt(var_type);
                let body_smt = self.expr_to_smt(body)?;
                let quantifier = match kind {
                    QuantifierKind::ForAll => "forall",
                    QuantifierKind::Exists => "exists",
                };
                Ok(format!(
                    "({} (({} {})) {})",
                    quantifier,
                    self.sanitize_name(var),
                    sort,
                    body_smt
                ))
            }

            Expression::FunctionCall { function, args } => {
                let args_smt: Result<Vec<_>, _> =
                    args.iter().map(|a| self.expr_to_smt(a)).collect();
                let args_str = args_smt?.join(" ");
                Ok(format!("({} {})", self.sanitize_name(function), args_str))
            }

            Expression::MemberAccess { object, member } => {
                let obj = self.expr_to_smt(object)?;
                Ok(format!("{}_{}", obj, self.sanitize_name(member)))
            }

            Expression::Index { array, index } => {
                let arr = self.expr_to_smt(array)?;
                let idx = self.expr_to_smt(index)?;
                Ok(format!("(select {} {})", arr, idx))
            }

            Expression::Old(inner) => {
                let inner_smt = self.expr_to_smt(inner)?;
                Ok(format!("{}_old", inner_smt))
            }

            Expression::Result => Ok("__result".to_string()),

            Expression::Sum { var, range: _, body } => {
                // Sum is modeled as uninterpreted function
                let body_smt = self.expr_to_smt(body)?;
                Ok(format!("(sum_{} {})", self.sanitize_name(var), body_smt))
            }
        }
    }

    /// Converts a literal to SMT-LIB format.
    fn literal_to_smt(&self, lit: &Literal) -> VerifierResult<String> {
        match lit {
            Literal::Bool(b) => Ok(if *b { "true" } else { "false" }.to_string()),
            Literal::Uint(n) => Ok(format!("(_ bv{} 256)", n)),
            Literal::Int(n) => {
                if *n >= 0 {
                    Ok(format!("(_ bv{} 256)", n))
                } else {
                    Ok(format!("(bvneg (_ bv{} 256))", -n))
                }
            }
            Literal::Address(bytes) => {
                let hex: String = bytes.iter().map(|b| format!("{:02x}", b)).collect();
                Ok(format!("#x{}", hex))
            }
            Literal::Bytes(bytes) => {
                let hex: String = bytes.iter().map(|b| format!("{:02x}", b)).collect();
                Ok(format!("#x{}", hex))
            }
            Literal::String(s) => Ok(format!("\"{}\"", s.replace('\"', "\"\""))),
        }
    }

    /// Converts a binary operator to SMT-LIB.
    fn binop_to_smt(&self, op: BinaryOperator) -> &'static str {
        match op {
            BinaryOperator::Add => "bvadd",
            BinaryOperator::Sub => "bvsub",
            BinaryOperator::Mul => "bvmul",
            BinaryOperator::Div => "bvudiv",
            BinaryOperator::Mod => "bvurem",
            BinaryOperator::Exp => "bvexp", // Would need custom function
            BinaryOperator::Eq => "=",
            BinaryOperator::Ne => "distinct",
            BinaryOperator::Lt => "bvult",
            BinaryOperator::Le => "bvule",
            BinaryOperator::Gt => "bvugt",
            BinaryOperator::Ge => "bvuge",
            BinaryOperator::And => "and",
            BinaryOperator::Or => "or",
            BinaryOperator::Implies => "=>",
            BinaryOperator::Iff => "=",
            BinaryOperator::BitAnd => "bvand",
            BinaryOperator::BitOr => "bvor",
            BinaryOperator::BitXor => "bvxor",
            BinaryOperator::Shl => "bvshl",
            BinaryOperator::Shr => "bvlshr",
        }
    }

    /// Converts a unary operator to SMT-LIB.
    fn unop_to_smt(&self, op: UnaryOperator) -> &'static str {
        match op {
            UnaryOperator::Not => "not",
            UnaryOperator::Neg => "bvneg",
            UnaryOperator::BitNot => "bvnot",
        }
    }

    /// Converts a type to SMT sort.
    fn type_to_smt(&self, var_type: &VarType) -> String {
        match var_type {
            VarType::Uint(bits) => format!("(_ BitVec {})", bits),
            VarType::Int(bits) => format!("(_ BitVec {})", bits),
            VarType::Bool => "Bool".to_string(),
            VarType::Address => "(_ BitVec 160)".to_string(),
            VarType::Bytes(n) => format!("(_ BitVec {})", n * 8),
            VarType::DynBytes => "(Array Int (_ BitVec 8))".to_string(),
            VarType::String => "String".to_string(),
            VarType::Array(inner) => {
                format!("(Array Int {})", self.type_to_smt(inner))
            }
            VarType::Mapping(key, val) => {
                format!(
                    "(Array {} {})",
                    self.type_to_smt(key),
                    self.type_to_smt(val)
                )
            }
            VarType::Struct(name, _) => name.clone(),
        }
    }

    /// Sanitizes a name for SMT-LIB (removes invalid characters).
    fn sanitize_name(&self, name: &str) -> String {
        name.replace('.', "_")
            .replace('[', "_")
            .replace(']', "_")
            .replace('(', "_")
            .replace(')', "_")
    }

    /// Encodes a symbolic value as SMT.
    fn encode_symbolic_value(&self, value: &SymbolicValue) -> VerifierResult<SmtExpr> {
        let smt_str = self.symbolic_to_smt(value)?;
        Ok(SmtExpr::new(&smt_str))
    }

    /// Converts a symbolic value to SMT-LIB format.
    fn symbolic_to_smt(&self, value: &SymbolicValue) -> VerifierResult<String> {
        match value {
            SymbolicValue::Concrete(lit) => self.literal_to_smt(lit),

            SymbolicValue::Symbol(name, _) => Ok(self.sanitize_name(name)),

            SymbolicValue::BinaryOp { left, op, right } => {
                let l = self.symbolic_to_smt(left)?;
                let r = self.symbolic_to_smt(right)?;
                let op_str = self.binop_to_smt(*op);
                Ok(format!("({} {} {})", op_str, l, r))
            }

            SymbolicValue::UnaryOp { op, operand } => {
                let v = self.symbolic_to_smt(operand)?;
                let op_str = self.unop_to_smt(*op);
                Ok(format!("({} {})", op_str, v))
            }

            SymbolicValue::Ite {
                condition,
                then_val,
                else_val,
            } => {
                let c = self.symbolic_to_smt(condition)?;
                let t = self.symbolic_to_smt(then_val)?;
                let e = self.symbolic_to_smt(else_val)?;
                Ok(format!("(ite {} {} {})", c, t, e))
            }

            SymbolicValue::Select { array, index } => {
                let arr = self.symbolic_to_smt(array)?;
                let idx = self.symbolic_to_smt(index)?;
                Ok(format!("(select {} {})", arr, idx))
            }

            SymbolicValue::Store {
                array,
                index,
                value,
            } => {
                let arr = self.symbolic_to_smt(array)?;
                let idx = self.symbolic_to_smt(index)?;
                let val = self.symbolic_to_smt(value)?;
                Ok(format!("(store {} {} {})", arr, idx, val))
            }
        }
    }

    /// Negates an SMT expression.
    pub fn negate(&self, expr: &SmtExpr) -> VerifierResult<SmtExpr> {
        Ok(SmtExpr::new(&format!("(not {})", expr.as_str())))
    }

    /// Creates conjunction of expressions.
    pub fn and(&self, a: &SmtExpr, b: &SmtExpr) -> VerifierResult<SmtExpr> {
        Ok(SmtExpr::new(&format!("(and {} {})", a.as_str(), b.as_str())))
    }

    /// Creates disjunction of expressions.
    pub fn or(&self, a: &SmtExpr, b: &SmtExpr) -> VerifierResult<SmtExpr> {
        Ok(SmtExpr::new(&format!("(or {} {})", a.as_str(), b.as_str())))
    }

    /// Checks satisfiability of an expression.
    pub fn check_sat(&self, expr: &SmtExpr) -> VerifierResult<SmtResult> {
        // Simulated SMT solving (real implementation would use Z3 or similar)
        // For now, return Unknown for complex expressions, or evaluate simple ones

        let s = expr.as_str();

        // Simple cases
        if s == "true" {
            return Ok(SmtResult::Sat(vec![]));
        }
        if s == "false" {
            return Ok(SmtResult::Unsat);
        }

        // Unknown for complex expressions
        Ok(SmtResult::Unknown(
            "SMT backend not configured (enable z3-backend feature)".to_string(),
        ))
    }

    /// Checks satisfiability with additional constraints.
    pub fn check_sat_with_constraints(
        &self,
        constraints: &[SmtExpr],
        expr: &SmtExpr,
    ) -> VerifierResult<SmtResult> {
        // Combine all constraints
        let mut combined = expr.clone();
        for c in constraints {
            combined = self.and(&combined, c)?;
        }

        self.check_sat(&combined)
    }
}

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

    #[test]
    fn test_smt_context_creation() {
        let ctx = SmtContext::new(30000);
        assert_eq!(ctx.timeout_ms, 30000);
    }

    #[test]
    fn test_literal_encoding() {
        let ctx = SmtContext::new(1000);

        assert_eq!(ctx.literal_to_smt(&Literal::Bool(true)).unwrap(), "true");
        assert_eq!(ctx.literal_to_smt(&Literal::Bool(false)).unwrap(), "false");
        assert_eq!(
            ctx.literal_to_smt(&Literal::Uint(42)).unwrap(),
            "(_ bv42 256)"
        );
    }

    #[test]
    fn test_type_encoding() {
        let ctx = SmtContext::new(1000);

        assert_eq!(ctx.type_to_smt(&VarType::Bool), "Bool");
        assert_eq!(ctx.type_to_smt(&VarType::Uint(256)), "(_ BitVec 256)");
        assert_eq!(ctx.type_to_smt(&VarType::Address), "(_ BitVec 160)");
    }

    #[test]
    fn test_negate() {
        let ctx = SmtContext::new(1000);
        let expr = SmtExpr::new("x");
        let neg = ctx.negate(&expr).unwrap();
        assert_eq!(neg.as_str(), "(not x)");
    }

    #[test]
    fn test_and_or() {
        let ctx = SmtContext::new(1000);
        let a = SmtExpr::new("a");
        let b = SmtExpr::new("b");

        assert_eq!(ctx.and(&a, &b).unwrap().as_str(), "(and a b)");
        assert_eq!(ctx.or(&a, &b).unwrap().as_str(), "(or a b)");
    }

    #[test]
    fn test_check_sat_simple() {
        let ctx = SmtContext::new(1000);

        let t = SmtExpr::new("true");
        assert!(matches!(ctx.check_sat(&t).unwrap(), SmtResult::Sat(_)));

        let f = SmtExpr::new("false");
        assert!(matches!(ctx.check_sat(&f).unwrap(), SmtResult::Unsat));
    }
}
