//! SMT solver integration for proof automation

use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::process::{Command, Stdio};
use std::io::Write;

/// SMT solver backend
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SmtBackend {
    Z3,
    CVC5,
    Yices2,
    MathSAT,
}

impl SmtBackend {
    /// Get command for solver
    fn command(&self) -> &str {
        match self {
            Self::Z3 => "z3",
            Self::CVC5 => "cvc5",
            Self::Yices2 => "yices-smt2",
            Self::MathSAT => "mathsat",
        }
    }
    
    /// Get command arguments
    fn args(&self) -> Vec<&str> {
        match self {
            Self::Z3 => vec!["-in", "-smt2"],
            Self::CVC5 => vec!["--lang", "smt2", "--incremental"],
            Self::Yices2 => vec!["--incremental"],
            Self::MathSAT => vec!["-input=smt2"],
        }
    }
}

/// SMT solver interface
pub struct SmtSolver {
    backend: SmtBackend,
    timeout_ms: u64,
    assertions: Vec<String>,
    declarations: HashMap<String, String>,
}

impl SmtSolver {
    /// Create new SMT solver
    pub fn new(backend: SmtBackend) -> Self {
        Self {
            backend,
            timeout_ms: 5000,
            assertions: Vec::new(),
            declarations: HashMap::new(),
        }
    }
    
    /// Set timeout
    pub fn set_timeout(&mut self, timeout_ms: u64) {
        self.timeout_ms = timeout_ms;
    }
    
    /// Declare a variable
    pub fn declare_var(&mut self, name: &str, sort: &str) {
        self.declarations.insert(
            name.to_string(),
            format!("(declare-fun {name} () {sort})")
        );
    }
    
    /// Declare a function
    pub fn declare_fun(&mut self, name: &str, params: &[&str], ret: &str) {
        let params_str = params.join(" ");
        self.declarations.insert(
            name.to_string(),
            format!("(declare-fun {name} ({params_str}) {ret})")
        );
    }
    
    /// Add an assertion
    pub fn assert(&mut self, expr: &str) {
        self.assertions.push(format!("(assert {expr})"));
    }
    
    /// Check satisfiability
    pub fn check_sat(&self) -> Result<SmtResult> {
        let query = self.build_query("(check-sat)");
        self.execute_query(&query)
    }
    
    /// Get model (if sat)
    pub fn get_model(&self) -> Result<Option<HashMap<String, String>>> {
        let query = self.build_query("(check-sat)\n(get-model)");
        let result = self.execute_query(&query)?;
        
        if result == SmtResult::Sat {
            Ok(Some(self.parse_model(&query)?))
        } else {
            Ok(None)
        }
    }
    
    /// Check validity (prove formula)
    pub fn prove(&self, formula: &str) -> Result<SmtResult> {
        let mut solver = self.clone();
        solver.assert(&format!("(not {formula})"));
        
        match solver.check_sat()? {
            SmtResult::Unsat => Ok(SmtResult::Valid),
            SmtResult::Sat => Ok(SmtResult::Invalid),
            other => Ok(other),
        }
    }
    
    /// Build complete SMT-LIB2 query
    fn build_query(&self, command: &str) -> String {
        let mut query = String::new();
        
        query.push_str("(set-logic ALL)\n");
        query.push_str(&format!("(set-option :timeout {})\n", self.timeout_ms));
        
        for decl in self.declarations.values() {
            query.push_str(decl);
            query.push('\n');
        }
        
        for assertion in &self.assertions {
            query.push_str(assertion);
            query.push('\n');
        }
        
        query.push_str(command);
        query.push('\n');
        query.push_str("(exit)\n");
        
        query
    }
    
    /// Execute query via solver process
    fn execute_query(&self, query: &str) -> Result<SmtResult> {
        let mut child = Command::new(self.backend.command())
            .args(self.backend.args())
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()?;
        
        if let Some(mut stdin) = child.stdin.take() {
            stdin.write_all(query.as_bytes())?;
        }
        
        let output = child.wait_with_output()?;
        let stdout = String::from_utf8_lossy(&output.stdout);
        
        self.parse_result(&stdout)
    }
    
    /// Parse solver result
    fn parse_result(&self, output: &str) -> Result<SmtResult> {
        if output.contains("sat") && !output.contains("unsat") {
            Ok(SmtResult::Sat)
        } else if output.contains("unsat") {
            Ok(SmtResult::Unsat)
        } else if output.contains("unknown") {
            Ok(SmtResult::Unknown)
        } else if output.contains("timeout") {
            Ok(SmtResult::Timeout)
        } else {
            Ok(SmtResult::Error(format!("Unexpected output: {output}")))
        }
    }
    
    /// Parse model from solver output
    fn parse_model(&self, _output: &str) -> Result<HashMap<String, String>> {
        Ok(HashMap::new())
    }
}

impl Clone for SmtSolver {
    fn clone(&self) -> Self {
        Self {
            backend: self.backend,
            timeout_ms: self.timeout_ms,
            assertions: self.assertions.clone(),
            declarations: self.declarations.clone(),
        }
    }
}

/// SMT query
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SmtQuery {
    pub formula: String,
    pub variables: Vec<(String, String)>,
    pub assumptions: Vec<String>,
}

impl SmtQuery {
    /// Create new query
    pub fn new(formula: &str) -> Self {
        Self {
            formula: formula.to_string(),
            variables: Vec::new(),
            assumptions: Vec::new(),
        }
    }
    
    /// Add variable
    pub fn add_var(&mut self, name: &str, sort: &str) {
        self.variables.push((name.to_string(), sort.to_string()));
    }
    
    /// Add assumption
    pub fn add_assumption(&mut self, expr: &str) {
        self.assumptions.push(expr.to_string());
    }
    
    /// Execute query
    pub fn execute(&self, backend: SmtBackend) -> Result<SmtResult> {
        let mut solver = SmtSolver::new(backend);
        
        for (name, sort) in &self.variables {
            solver.declare_var(name, sort);
        }
        
        for assumption in &self.assumptions {
            solver.assert(assumption);
        }
        
        solver.prove(&self.formula)
    }
}

/// SMT solver result
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum SmtResult {
    Sat,
    Unsat,
    Valid,
    Invalid,
    Unknown,
    Timeout,
    Error(String),
}

impl SmtResult {
    /// Check if result indicates success
    pub fn is_success(&self) -> bool {
        matches!(self, Self::Valid | Self::Unsat)
    }
    
    /// Get human-readable description
    pub fn description(&self) -> &str {
        match self {
            Self::Sat => "Satisfiable",
            Self::Unsat => "Unsatisfiable",
            Self::Valid => "Valid",
            Self::Invalid => "Invalid",
            Self::Unknown => "Unknown",
            Self::Timeout => "Timeout",
            Self::Error(_) => "Error",
        }
    }
}

/// High-level proof automation
pub struct ProofAutomation {
    backend: SmtBackend,
    cache: HashMap<String, SmtResult>,
}

impl ProofAutomation {
    /// Create new automation
    pub fn new(backend: SmtBackend) -> Self {
        Self {
            backend,
            cache: HashMap::new(),
        }
    }
    
    /// Prove implication
    pub fn prove_implication(&mut self, antecedent: &str, consequent: &str) -> Result<SmtResult> {
        let formula = format!("(=> {antecedent} {consequent})");
        self.prove(&formula)
    }
    
    /// Prove equivalence
    pub fn prove_equivalence(&mut self, left: &str, right: &str) -> Result<SmtResult> {
        let formula = format!("(= {left} {right})");
        self.prove(&formula)
    }
    
    /// Prove with caching
    fn prove(&mut self, formula: &str) -> Result<SmtResult> {
        if let Some(cached) = self.cache.get(formula) {
            return Ok(cached.clone());
        }
        
        let query = SmtQuery::new(formula);
        let result = query.execute(self.backend)?;
        self.cache.insert(formula.to_string(), result.clone());
        
        Ok(result)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_smt_query_construction() {
        let mut query = SmtQuery::new("(> x 0)");
        query.add_var("x", "Int");
        query.add_assumption("(< x 10)");
        
        assert_eq!(query.formula, "(> x 0)");
        assert_eq!(query.variables.len(), 1);
        assert_eq!(query.assumptions.len(), 1);
    }
    
    #[test]
    fn test_solver_initialization() {
        let solver = SmtSolver::new(SmtBackend::Z3);
        assert_eq!(solver.backend, SmtBackend::Z3);
        assert_eq!(solver.timeout_ms, 5000);
    }
}