//! Chain-of-thought reasoning framework with step validation and error recovery

use crate::error::{ReasoningError, Result};
use async_trait::async_trait;
use indexmap::IndexMap;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::{Duration, Instant};
use uuid::Uuid;

/// Evidence supporting a reasoning step
#[derive(Debug, Clone)]
pub struct Evidence {
    pub id: String,
    pub source_type: String,
    pub description: String,
    pub reliability: f64, // 0.0 to 1.0
    pub timestamp: Instant,
}

impl Evidence {
    pub fn new(source_type: impl Into<String>, description: impl Into<String>) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            source_type: source_type.into(),
            description: description.into(),
            reliability: 1.0,
            timestamp: Instant::now(),
        }
    }

    pub fn with_reliability(mut self, reliability: f64) -> Self {
        self.reliability = reliability.clamp(0.0, 1.0);
        self
    }
}

/// Chain-of-thought reasoning step
#[derive(Debug, Clone)]
pub struct ReasoningStep {
    pub id: String,
    pub description: String,
    pub evidence: Vec<Evidence>,
    pub conclusion: String,
    pub confidence: f64, // 0.0 to 1.0
    pub step_type: StepType,
    pub dependencies: HashSet<String>, // IDs of steps this depends on
    pub created_at: Instant,
    pub validation_status: ValidationStatus,
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum StepType {
    Premise,
    Deduction,
    Induction,
    Abduction,
    Hypothesis,
    Observation,
    Conclusion,
}

#[derive(Debug, Clone)]
pub enum ValidationStatus {
    Pending,
    Valid,
    Invalid { reason: String },
    Warning { message: String },
}

impl ReasoningStep {
    pub fn new(
        id: impl Into<String>,
        description: impl Into<String>,
        evidence: Vec<Evidence>,
        conclusion: impl Into<String>,
    ) -> Self {
        let id = id.into();
        let evidence_deps: HashSet<String> = evidence
            .iter()
            .filter(|e| e.source_type.starts_with("from_step"))
            .map(|e| e.description.clone())
            .collect();

        Self {
            id,
            description: description.into(),
            evidence,
            conclusion: conclusion.into(),
            confidence: 1.0,
            step_type: StepType::Deduction,
            dependencies: evidence_deps,
            created_at: Instant::now(),
            validation_status: ValidationStatus::Pending,
        }
    }

    pub fn set_confidence(&mut self, confidence: f64) {
        self.confidence = confidence.clamp(0.0, 1.0);
    }

    pub fn set_step_type(&mut self, step_type: StepType) {
        self.step_type = step_type;
    }

    pub fn add_evidence(&mut self, evidence: Evidence) {
        if evidence.source_type.starts_with("from_step") {
            self.dependencies.insert(evidence.description.clone());
        }
        self.evidence.push(evidence);
    }

    pub fn get_evidence_strength(&self) -> f64 {
        if self.evidence.is_empty() {
            return 0.0;
        }

        let total_reliability: f64 = self.evidence.iter().map(|e| e.reliability).sum();
        total_reliability / self.evidence.len() as f64
    }
}

/// Step validation trait
#[async_trait]
pub trait StepValidation {
    async fn validate_step(
        &self,
        step: &ReasoningStep,
        context: &ReasoningContext,
    ) -> Result<ValidationStatus>;
}

/// Rule-based step validator
pub struct StepValidator {
    contradiction_rules: HashMap<String, Box<dyn Fn(&ReasoningStep) -> bool + Send + Sync>>,
    evidence_requirements: HashMap<StepType, usize>,
    confidence_thresholds: HashMap<StepType, f64>,
}

impl Default for StepValidator {
    fn default() -> Self {
        Self::new()
    }
}

impl StepValidator {
    pub fn new() -> Self {
        let mut evidence_requirements = HashMap::new();
        evidence_requirements.insert(StepType::Premise, 1);
        evidence_requirements.insert(StepType::Deduction, 2);
        evidence_requirements.insert(StepType::Induction, 3);
        evidence_requirements.insert(StepType::Hypothesis, 1);
        evidence_requirements.insert(StepType::Observation, 1);
        evidence_requirements.insert(StepType::Conclusion, 1);

        let mut confidence_thresholds = HashMap::new();
        confidence_thresholds.insert(StepType::Premise, 0.8);
        confidence_thresholds.insert(StepType::Deduction, 0.7);
        confidence_thresholds.insert(StepType::Conclusion, 0.6);

        Self {
            contradiction_rules: HashMap::new(),
            evidence_requirements,
            confidence_thresholds,
        }
    }

    pub fn add_contradiction_rule<F>(&mut self, name: impl Into<String>, rule: F)
    where
        F: Fn(&ReasoningStep) -> bool + Send + Sync + 'static,
    {
        self.contradiction_rules.insert(name.into(), Box::new(rule));
    }
}

#[async_trait]
impl StepValidation for StepValidator {
    async fn validate_step(
        &self,
        step: &ReasoningStep,
        _context: &ReasoningContext,
    ) -> Result<ValidationStatus> {
        // Check evidence requirements
        let required_evidence = self
            .evidence_requirements
            .get(&step.step_type)
            .unwrap_or(&1);

        if step.evidence.len() < *required_evidence {
            return Ok(ValidationStatus::Invalid {
                reason: format!(
                    "Step requires at least {} evidence items, but only {} provided",
                    required_evidence,
                    step.evidence.len()
                ),
            });
        }

        // Check confidence threshold
        if let Some(threshold) = self.confidence_thresholds.get(&step.step_type)
            && step.confidence < *threshold
        {
            return Ok(ValidationStatus::Warning {
                message: format!(
                    "Step confidence {} is below recommended threshold {}",
                    step.confidence, threshold
                ),
            });
        }

        // Check contradiction rules
        for (rule_name, rule_fn) in &self.contradiction_rules {
            if rule_fn(step) {
                return Ok(ValidationStatus::Invalid {
                    reason: format!("Step violates contradiction rule: {rule_name}"),
                });
            }
        }

        // Check evidence strength
        let evidence_strength = step.get_evidence_strength();
        if evidence_strength < 0.5 {
            return Ok(ValidationStatus::Warning {
                message: format!("Low evidence strength: {evidence_strength:.2}"),
            });
        }

        Ok(ValidationStatus::Valid)
    }
}

/// Reasoning context for chain-of-thought processing
#[derive(Debug, Clone)]
pub struct ReasoningContext {
    pub steps: IndexMap<String, ReasoningStep>,
    pub variables: HashMap<String, String>,
    pub metadata: HashMap<String, serde_json::Value>,
    pub created_at: Instant,
}

impl ReasoningContext {
    pub fn new() -> Self {
        Self {
            steps: IndexMap::new(),
            variables: HashMap::new(),
            metadata: HashMap::new(),
            created_at: Instant::now(),
        }
    }

    pub fn set_variable(&mut self, key: impl Into<String>, value: impl Into<String>) {
        self.variables.insert(key.into(), value.into());
    }

    pub fn get_variable(&self, key: &str) -> Option<&str> {
        self.variables.get(key).map(|s| s.as_str())
    }

    pub fn add_step(&mut self, step: ReasoningStep) {
        self.steps.insert(step.id.clone(), step);
    }

    pub fn get_step(&self, id: &str) -> Option<&ReasoningStep> {
        self.steps.get(id)
    }

    pub fn validate_dependencies(&self) -> Result<()> {
        for step in self.steps.values() {
            for dep_id in &step.dependencies {
                if !self.steps.contains_key(dep_id) {
                    return Err(ReasoningError::InsufficientEvidence {
                        step_id: step.id.clone(),
                    }
                    .into());
                }
            }
        }
        Ok(())
    }

    pub fn detect_circular_dependencies(&self) -> Result<()> {
        let mut visited = HashSet::new();
        let mut in_progress = HashSet::new();

        for step_id in self.steps.keys() {
            if !visited.contains(step_id) {
                self.dfs_cycle_detection(step_id, &mut visited, &mut in_progress)?;
            }
        }

        Ok(())
    }

    fn dfs_cycle_detection(
        &self,
        step_id: &str,
        visited: &mut HashSet<String>,
        in_progress: &mut HashSet<String>,
    ) -> Result<()> {
        if in_progress.contains(step_id) {
            return Err(ReasoningError::CircularReasoning.into());
        }

        if visited.contains(step_id) {
            return Ok(());
        }

        in_progress.insert(step_id.to_string());

        if let Some(step) = self.steps.get(step_id) {
            for dep_id in &step.dependencies {
                self.dfs_cycle_detection(dep_id, visited, in_progress)?;
            }
        }

        in_progress.remove(step_id);
        visited.insert(step_id.to_string());

        Ok(())
    }
}

impl Default for ReasoningContext {
    fn default() -> Self {
        Self::new()
    }
}

/// Reasoning execution result
#[derive(Debug)]
pub struct ReasoningResult {
    pub is_valid: bool,
    pub conclusion: String,
    pub steps_executed: usize,
    pub overall_confidence: f64,
    pub execution_time: Duration,
    pub errors: Vec<String>,
    pub warnings: Vec<String>,
    pub recovered_from_errors: bool,
    pub recovery_count: usize,
}

/// Chain-of-thought reasoning engine
pub struct ChainOfThought {
    name: String,
    context: ReasoningContext,
    validator: Option<Arc<dyn StepValidation + Send + Sync>>,
    timeout: Option<Duration>,
    error_recovery_enabled: bool,
    recovery_steps: HashMap<String, ReasoningStep>,
    expected_template_steps: Option<usize>,
    template_name: Option<String>,
}

impl std::fmt::Debug for ChainOfThought {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ChainOfThought")
            .field("name", &self.name)
            .field("context", &self.context)
            .field("validator", &self.validator.is_some())
            .field("timeout", &self.timeout)
            .field("error_recovery_enabled", &self.error_recovery_enabled)
            .field("recovery_steps", &self.recovery_steps.len())
            .field("expected_template_steps", &self.expected_template_steps)
            .field("template_name", &self.template_name)
            .finish()
    }
}

impl ChainOfThought {
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            context: ReasoningContext::new(),
            validator: None,
            timeout: None,
            error_recovery_enabled: false,
            recovery_steps: HashMap::new(),
            expected_template_steps: None,
            template_name: None,
        }
    }

    pub fn with_context(name: impl Into<String>, context: ReasoningContext) -> Self {
        Self {
            name: name.into(),
            context,
            validator: None,
            timeout: None,
            error_recovery_enabled: false,
            recovery_steps: HashMap::new(),
            expected_template_steps: None,
            template_name: None,
        }
    }

    pub fn set_validator(&mut self, validator: impl StepValidation + Send + Sync + 'static) {
        self.validator = Some(Arc::new(validator));
    }

    pub fn set_timeout(&mut self, timeout: Duration) {
        self.timeout = Some(timeout);
    }

    pub fn enable_error_recovery(&mut self, enabled: bool) {
        self.error_recovery_enabled = enabled;
    }

    pub fn get_context(&self) -> &ReasoningContext {
        &self.context
    }

    pub fn expected_steps(&self) -> usize {
        self.expected_template_steps
            .unwrap_or(self.context.steps.len())
    }

    pub fn get_template_name(&self) -> Option<&str> {
        self.template_name.as_deref()
    }

    pub async fn add_step(&mut self, step: ReasoningStep) -> Result<()> {
        // Validate dependencies exist
        for dep_id in &step.dependencies {
            if !self.context.steps.contains_key(dep_id) {
                return Err(ReasoningError::InsufficientEvidence {
                    step_id: step.id.clone(),
                }
                .into());
            }
        }

        self.context.add_step(step);
        Ok(())
    }

    pub async fn add_step_as_branch(&mut self, step: ReasoningStep, parent_id: &str) -> Result<()> {
        if !self.context.steps.contains_key(parent_id) {
            return Err(ReasoningError::InsufficientEvidence {
                step_id: step.id.clone(),
            }
            .into());
        }

        self.context.add_step(step);
        Ok(())
    }

    pub async fn add_recovery_step(
        &mut self,
        failed_step_id: impl Into<String>,
        recovery_step: ReasoningStep,
    ) -> Result<()> {
        self.recovery_steps
            .insert(failed_step_id.into(), recovery_step);
        Ok(())
    }

    pub async fn execute(&mut self) -> Result<ReasoningResult> {
        let start_time = Instant::now();
        let errors = Vec::new();
        let mut warnings = Vec::new();
        let mut recovery_count = 0;
        let mut recovered_from_errors = false;

        // Validate overall structure
        if let Err(e) = self.context.validate_dependencies() {
            if self.error_recovery_enabled {
                // Try to recover using recovery steps
                recovery_count += 1;
                recovered_from_errors = true;
            } else {
                return Err(e);
            }
        }

        self.context.detect_circular_dependencies()?;

        // Execute validation if validator is present
        if let Some(validator) = &self.validator {
            for step in self.context.steps.values() {
                match validator.validate_step(step, &self.context).await? {
                    ValidationStatus::Invalid { reason: _ } => {
                        if self.error_recovery_enabled && self.recovery_steps.contains_key(&step.id)
                        {
                            recovery_count += 1;
                            recovered_from_errors = true;
                        } else {
                            return Err(ReasoningError::StepValidationFailed {
                                step_id: step.id.clone(),
                            }
                            .into());
                        }
                    }
                    ValidationStatus::Warning { message } => {
                        warnings.push(message);
                    }
                    ValidationStatus::Valid => {}
                    ValidationStatus::Pending => {}
                }
            }
        }

        // Check for evidence insufficiency
        for step in self.context.steps.values() {
            if step.evidence.is_empty() && !matches!(step.step_type, StepType::Premise) {
                return Err(ReasoningError::InsufficientEvidence {
                    step_id: step.id.clone(),
                }
                .into());
            }
        }

        // Calculate overall confidence
        let total_confidence: f64 = self
            .context
            .steps
            .values()
            .map(|s| s.confidence * s.get_evidence_strength())
            .sum();
        let overall_confidence = if self.context.steps.is_empty() {
            0.0
        } else {
            total_confidence / self.context.steps.len() as f64
        };

        // Get final conclusion
        let conclusion = self
            .context
            .steps
            .values()
            .last()
            .map(|s| s.conclusion.clone())
            .unwrap_or_default();

        let execution_time = start_time.elapsed();

        Ok(ReasoningResult {
            is_valid: errors.is_empty(),
            conclusion,
            steps_executed: self.context.steps.len(),
            overall_confidence,
            execution_time,
            errors,
            warnings,
            recovered_from_errors,
            recovery_count,
        })
    }
}

/// Template-based reasoning engine
pub struct ReasoningEngine {
    templates: HashMap<String, Vec<String>>,
    default_validator: Arc<StepValidator>,
}

impl ReasoningEngine {
    pub fn new() -> Self {
        Self {
            templates: HashMap::new(),
            default_validator: Arc::new(StepValidator::new()),
        }
    }

    pub async fn register_template(
        &mut self,
        name: impl Into<String>,
        steps: Vec<String>,
    ) -> Result<()> {
        self.templates.insert(name.into(), steps);
        Ok(())
    }

    pub async fn create_chain_from_template(
        &self,
        template_name: &str,
        context: ReasoningContext,
    ) -> Result<ChainOfThought> {
        let template_steps = self.templates.get(template_name).ok_or_else(|| {
            ReasoningError::InvalidStep(format!("Template not found: {template_name}"))
        })?;

        let mut chain = ChainOfThought::with_context(template_name, context);
        chain.expected_template_steps = Some(template_steps.len());
        chain.template_name = Some(template_name.to_string());
        chain.set_validator(StepValidator::new());

        Ok(chain)
    }
}

impl Default for ReasoningEngine {
    fn default() -> Self {
        Self::new()
    }
}