//! The iterative correction protocol engine.
//!
//! Implements the "propose-verify-refine" loop:
//! 1. Validate the SQL
//! 2. If valid, run optional lint/safety checks, return success
//! 3. If invalid, attempt to fix via the fixer engine
//! 4. Re-validate the fixed SQL
//! 5. Track all attempts (never retry the same fix twice)
//! 6. After max_iterations or no progress, return best attempt

use std::collections::HashSet;
use std::time::Instant;

use crate::fixer::engine::fix;
use crate::fixer::types::FixConfig;
use crate::linter::engine::lint;
use crate::schema::types::SchemaDefinition;
use crate::validator::engine::validate;

use super::types::*;

/// Run the iterative correction protocol on a SQL query.
///
/// This is the primary entry point. It validates, fixes, re-validates, and
/// optionally lints and safety-checks the SQL in a loop until:
/// - The SQL is valid and passes all checks, or
/// - `max_iterations` is exhausted, or
/// - No further progress can be made.
///
/// # Example
///
/// ```rust,no_run
/// use altimate_core::correction::engine::correct;
/// use altimate_core::correction::types::CorrectionConfig;
/// use altimate_core::SchemaDefinition;
///
/// # async fn example() {
/// let schema = SchemaDefinition::from_yaml_file("schema.yaml").unwrap();
/// let result = correct("SELECT * FROM usrs", &schema, &CorrectionConfig::default()).await;
/// println!("Status: {}", result.status);
/// if let Some(sql) = &result.corrected_sql {
///     println!("Corrected: {}", sql);
/// }
/// # }
/// ```
pub async fn correct(
    sql: &str,
    schema: &SchemaDefinition,
    config: &CorrectionConfig,
) -> CorrectionResult {
    let start = Instant::now();
    let mut current_sql = sql.to_string();
    let mut iterations = Vec::new();
    let mut tried_fixes: HashSet<String> = HashSet::new();

    // Step 1: Initial validation
    let initial_result = validate(&current_sql, schema).await;

    if initial_result.valid {
        // Already valid — compute quality score and return
        let score = compute_quality_score(&current_sql, schema, config);
        let elapsed = start.elapsed().as_millis() as u64;
        return CorrectionResult {
            status: CorrectionStatus::AlreadyValid,
            original_sql: sql.to_string(),
            corrected_sql: None,
            iterations,
            final_validation: serde_json::to_value(&initial_result).ok(),
            final_score: score,
            total_time_ms: elapsed,
        };
    }

    // Step 2: Correction loop
    let mut best_sql = current_sql.clone();
    let mut best_error_count = initial_result.errors.len();

    for iteration_num in 1..=config.max_iterations {
        // Validate current SQL
        let result = validate(&current_sql, schema).await;

        if result.valid {
            // Fixed! Compute score and return.
            let score = compute_quality_score(&current_sql, schema, config);
            let elapsed = start.elapsed().as_millis() as u64;
            return CorrectionResult {
                status: CorrectionStatus::Fixed,
                original_sql: sql.to_string(),
                corrected_sql: Some(current_sql),
                iterations,
                final_validation: serde_json::to_value(&result).ok(),
                final_score: score,
                total_time_ms: elapsed,
            };
        }

        let error_messages: Vec<String> = result.errors.iter().map(|e| e.message.clone()).collect();

        // Create a fix key to avoid retrying the same fix
        let fix_key = format!("{}:{}", current_sql, error_messages.join("|"));

        if tried_fixes.contains(&fix_key) {
            // Already tried this exact fix — skip
            iterations.push(CorrectionIteration {
                iteration: iteration_num,
                input_sql: current_sql.clone(),
                validation_errors: error_messages,
                fix_applied: None,
                fix_description: Some("Skipped: already tried this fix".to_string()),
                result: IterationResult::Skipped,
            });
            break;
        }
        tried_fixes.insert(fix_key);

        // Attempt to fix via the fixer engine
        let fix_config = FixConfig {
            min_confidence: config.min_confidence,
            max_iterations: 1, // Single fix pass per correction iteration
            revalidate: false, // We'll re-validate ourselves
        };

        let fix_result = fix(&current_sql, schema, &fix_config).await;

        if !fix_result.fixed {
            // No fix available
            iterations.push(CorrectionIteration {
                iteration: iteration_num,
                input_sql: current_sql.clone(),
                validation_errors: error_messages,
                fix_applied: None,
                fix_description: Some("No automatic fix available".to_string()),
                result: IterationResult::Skipped,
            });
            break;
        }

        let fixed_sql = fix_result.fixed_sql.clone();
        let fix_descriptions: Vec<String> = fix_result
            .fixes_applied
            .iter()
            .map(|f| f.explanation.clone())
            .collect();

        // Re-validate the fixed SQL
        let post_result = validate(&fixed_sql, schema).await;

        let iteration_result = if post_result.valid {
            IterationResult::Fixed
        } else if post_result.errors.len() < result.errors.len() {
            IterationResult::StillInvalid
        } else if post_result.errors.len() > result.errors.len() {
            IterationResult::NewErrors
        } else {
            // Same number of errors — check if they're different errors
            let old_msgs: HashSet<&str> =
                result.errors.iter().map(|e| e.message.as_str()).collect();
            let new_msgs: HashSet<&str> = post_result
                .errors
                .iter()
                .map(|e| e.message.as_str())
                .collect();
            if old_msgs == new_msgs {
                IterationResult::Skipped
            } else {
                IterationResult::StillInvalid
            }
        };

        iterations.push(CorrectionIteration {
            iteration: iteration_num,
            input_sql: current_sql.clone(),
            validation_errors: error_messages,
            fix_applied: Some(fixed_sql.clone()),
            fix_description: Some(fix_descriptions.join("; ")),
            result: iteration_result,
        });

        // Track the best result
        if post_result.errors.len() < best_error_count {
            best_error_count = post_result.errors.len();
            best_sql = fixed_sql.clone();
        }

        // Update current SQL for next iteration
        current_sql = fixed_sql;

        // Stop if we introduced new errors (don't go further down a bad path)
        if iteration_result == IterationResult::NewErrors {
            // Revert to best SQL
            current_sql = best_sql.clone();
            break;
        }

        // Stop if no progress (same errors, iteration skipped)
        if iteration_result == IterationResult::Skipped {
            break;
        }
    }

    // Final validation of the best SQL
    let final_result = validate(&current_sql, schema).await;
    let status = if final_result.valid {
        CorrectionStatus::Fixed
    } else if current_sql != sql {
        CorrectionStatus::PartialFix
    } else {
        CorrectionStatus::Unfixable
    };

    let score = compute_quality_score(&current_sql, schema, config);
    let elapsed = start.elapsed().as_millis() as u64;

    CorrectionResult {
        status,
        original_sql: sql.to_string(),
        corrected_sql: if current_sql != sql {
            Some(current_sql)
        } else {
            None
        },
        iterations,
        final_validation: serde_json::to_value(&final_result).ok(),
        final_score: score,
        total_time_ms: elapsed,
    }
}

/// Compute a quality score for a SQL query.
///
/// Runs lint and (optionally) safety checks and combines scores:
/// - syntax: 0 or 1 (from validation)
/// - lint: 1.0 minus a penalty per finding
/// - safety: 1.0 if no policy violations, reduced otherwise
/// - complexity: based on query structure
fn compute_quality_score(
    sql: &str,
    schema: &SchemaDefinition,
    config: &CorrectionConfig,
) -> QualityScore {
    let lint_score = if config.enable_lint_checks {
        let lint_result = lint(sql, schema);
        compute_lint_score(&lint_result)
    } else {
        1.0
    };

    let complexity_score = compute_complexity_score(sql);

    let mut score = QualityScore {
        syntax_valid: true,
        lint_score,
        safety_score: 1.0,
        complexity_score,
        ..Default::default()
    };

    score.compute_overall();
    score
}

/// Compute a lint-based score.
///
/// Starts at 1.0 and deducts per finding:
/// - Error: -0.2
/// - Warning: -0.1
/// - Info: -0.05
fn compute_lint_score(result: &crate::linter::types::LintResult) -> f64 {
    use crate::linter::types::LintSeverity;
    let mut score = 1.0_f64;
    for finding in &result.findings {
        match finding.severity {
            LintSeverity::Error => score -= 0.2,
            LintSeverity::Warning => score -= 0.1,
            LintSeverity::Info => score -= 0.05,
        }
    }
    score.max(0.0)
}

/// Compute a complexity score from SQL structure.
///
/// Returns 1.0 for simple queries, decreasing for complex ones.
fn compute_complexity_score(sql: &str) -> f64 {
    let upper = sql.to_uppercase();
    let mut penalty = 0.0_f64;

    // Count JOINs
    let join_count = upper.matches(" JOIN ").count();
    penalty += join_count as f64 * 0.1;

    // Subqueries
    let subquery_indicators = upper.matches("SELECT").count();
    if subquery_indicators > 1 {
        penalty += (subquery_indicators - 1) as f64 * 0.15;
    }

    // Window functions
    if upper.contains(" OVER(") || upper.contains(" OVER (") {
        penalty += 0.1;
    }

    // UNION
    if upper.contains(" UNION ") {
        penalty += 0.1;
    }

    // CROSS JOIN
    if upper.contains("CROSS JOIN") {
        penalty += 0.2;
    }

    (1.0 - penalty).max(0.0)
}

/// Correct a SQL query with default configuration.
///
/// Convenience wrapper around [`correct`] with [`CorrectionConfig::default()`].
pub async fn correct_default(sql: &str, schema: &SchemaDefinition) -> CorrectionResult {
    correct(sql, schema, &CorrectionConfig::default()).await
}

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

    fn ecommerce_schema() -> SchemaDefinition {
        SchemaDefinition::from_yaml(include_str!(
            "../../../../tests/fixtures/schemas/ecommerce.yaml"
        ))
        .expect("Failed to load ecommerce schema fixture")
    }

    // ─── Already valid SQL ───────────────────────────────────

    #[tokio::test]
    async fn test_correct_valid_sql_returns_already_valid() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECT id, email FROM users", &schema).await;
        assert_eq!(result.status, CorrectionStatus::AlreadyValid);
        assert!(result.corrected_sql.is_none());
        assert!(result.iterations.is_empty());
    }

    #[tokio::test]
    async fn test_correct_valid_select_star() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECT * FROM users", &schema).await;
        assert_eq!(result.status, CorrectionStatus::AlreadyValid);
    }

    #[tokio::test]
    async fn test_correct_valid_join_query() {
        let schema = ecommerce_schema();
        let sql = "SELECT u.id, o.id FROM users u JOIN orders o ON u.id = o.user_id";
        let result = correct_default(sql, &schema).await;
        assert_eq!(result.status, CorrectionStatus::AlreadyValid);
    }

    #[tokio::test]
    async fn test_already_valid_has_quality_score() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECT id FROM users", &schema).await;
        assert_eq!(result.status, CorrectionStatus::AlreadyValid);
        assert!(result.final_score.syntax_valid);
        assert!(result.final_score.overall > 0.0);
    }

    #[tokio::test]
    async fn test_already_valid_has_zero_time() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECT id FROM users", &schema).await;
        // Should complete very quickly
        assert!(result.total_time_ms < 5000);
    }

    // ─── Fixable errors ──────────────────────────────────────

    #[tokio::test]
    async fn test_correct_misspelled_table() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECT * FROM usrs", &schema).await;
        assert_eq!(result.status, CorrectionStatus::Fixed);
        assert_eq!(result.corrected_sql.as_deref(), Some("SELECT * FROM users"));
    }

    #[tokio::test]
    async fn test_correct_misspelled_table_records_iteration() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECT * FROM usrs", &schema).await;
        assert!(!result.iterations.is_empty());
        assert_eq!(result.iterations[0].iteration, 1);
        assert!(!result.iterations[0].validation_errors.is_empty());
    }

    #[tokio::test]
    async fn test_correct_misspelled_table_has_fix_description() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECT * FROM usrs", &schema).await;
        assert!(result.iterations[0].fix_applied.is_some());
        assert!(result.iterations[0].fix_description.is_some());
    }

    #[tokio::test]
    async fn test_correct_misspelled_column() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECT emal FROM users", &schema).await;
        // The fixer may or may not handle column fixes — check status is not AlreadyValid
        assert_ne!(result.status, CorrectionStatus::AlreadyValid);
    }

    #[tokio::test]
    async fn test_correct_records_original_sql() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECT * FROM usrs", &schema).await;
        assert_eq!(result.original_sql, "SELECT * FROM usrs");
    }

    #[tokio::test]
    async fn test_correct_fixed_has_quality_score() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECT * FROM usrs", &schema).await;
        assert!(result.final_score.overall > 0.0);
    }

    // ─── Unfixable errors ────────────────────────────────────

    #[tokio::test]
    async fn test_correct_syntax_error_is_unfixable() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECTZ * FROM users", &schema).await;
        assert_eq!(result.status, CorrectionStatus::Unfixable);
        assert!(result.corrected_sql.is_none());
    }

    #[tokio::test]
    async fn test_correct_empty_sql_is_unfixable() {
        let schema = ecommerce_schema();
        let result = correct_default("", &schema).await;
        assert_eq!(result.status, CorrectionStatus::Unfixable);
    }

    #[tokio::test]
    async fn test_correct_garbage_sql_is_unfixable() {
        let schema = ecommerce_schema();
        let result = correct_default("this is not sql at all", &schema).await;
        assert_eq!(result.status, CorrectionStatus::Unfixable);
    }

    #[tokio::test]
    async fn test_correct_unfixable_has_iterations() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECTZ * FROM users", &schema).await;
        // There should be at least one iteration that was skipped
        assert!(!result.iterations.is_empty() || result.status == CorrectionStatus::Unfixable);
    }

    // ─── Max iterations ──────────────────────────────────────

    #[tokio::test]
    async fn test_correct_respects_max_iterations() {
        let schema = ecommerce_schema();
        let config = CorrectionConfig {
            max_iterations: 1,
            ..Default::default()
        };
        let result = correct("SELECT * FROM usrs", &schema, &config).await;
        assert!(result.iterations.len() <= 1);
    }

    #[tokio::test]
    async fn test_correct_max_iterations_zero() {
        let schema = ecommerce_schema();
        let config = CorrectionConfig {
            max_iterations: 0,
            ..Default::default()
        };
        let result = correct("SELECT * FROM usrs", &schema, &config).await;
        assert!(result.iterations.is_empty());
        // With 0 iterations, the final validation after the loop should still run
        assert_ne!(result.status, CorrectionStatus::AlreadyValid);
    }

    #[tokio::test]
    async fn test_correct_max_iterations_large() {
        let schema = ecommerce_schema();
        let config = CorrectionConfig {
            max_iterations: 100,
            ..Default::default()
        };
        let result = correct("SELECT * FROM usrs", &schema, &config).await;
        // Should fix in just 1-2 iterations, not 100
        assert!(result.iterations.len() <= 5);
        assert_eq!(result.status, CorrectionStatus::Fixed);
    }

    // ─── Confidence threshold ────────────────────────────────

    #[tokio::test]
    async fn test_correct_high_confidence_threshold() {
        let schema = ecommerce_schema();
        let config = CorrectionConfig {
            min_confidence: 0.99,
            ..Default::default()
        };
        let result = correct("SELECT * FROM usrs", &schema, &config).await;
        // High threshold may prevent the fix
        assert!(result.iterations.len() >= 1 || result.status == CorrectionStatus::Unfixable);
    }

    #[tokio::test]
    async fn test_correct_low_confidence_threshold() {
        let schema = ecommerce_schema();
        let config = CorrectionConfig {
            min_confidence: 0.1,
            ..Default::default()
        };
        let result = correct("SELECT * FROM usrs", &schema, &config).await;
        assert_eq!(result.status, CorrectionStatus::Fixed);
    }

    // ─── Lint integration ────────────────────────────────────

    #[tokio::test]
    async fn test_correct_with_lint_enabled() {
        let schema = ecommerce_schema();
        let config = CorrectionConfig {
            enable_lint_checks: true,
            ..Default::default()
        };
        let result = correct("SELECT * FROM usrs", &schema, &config).await;
        // Lint findings should affect the quality score
        assert!(result.final_score.lint_score <= 1.0);
    }

    #[tokio::test]
    async fn test_correct_with_lint_disabled() {
        let schema = ecommerce_schema();
        let config = CorrectionConfig {
            enable_lint_checks: false,
            ..Default::default()
        };
        let result = correct("SELECT * FROM usrs", &schema, &config).await;
        // Lint disabled means lint_score defaults to 1.0
        assert_eq!(result.final_score.lint_score, 1.0);
    }

    // ─── Safety scan config ──────────────────────────────────

    #[tokio::test]
    async fn test_correct_with_safety_disabled() {
        let schema = ecommerce_schema();
        let config = CorrectionConfig {
            enable_safety_scan: false,
            ..Default::default()
        };
        let result = correct("SELECT * FROM usrs", &schema, &config).await;
        // Safety score defaults to 1.0 when disabled
        assert_eq!(result.final_score.safety_score, 1.0);
    }

    // ─── Quality score ───────────────────────────────────────

    #[tokio::test]
    async fn test_quality_score_for_simple_query() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECT id FROM users", &schema).await;
        // Simple query should have high complexity_score
        assert!(result.final_score.complexity_score >= 0.8);
    }

    #[tokio::test]
    async fn test_quality_score_overall_range() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECT id FROM users", &schema).await;
        assert!(result.final_score.overall >= 0.0);
        assert!(result.final_score.overall <= 1.0);
    }

    #[tokio::test]
    async fn test_quality_score_compute_overall() {
        let mut score = QualityScore {
            syntax_valid: true,
            lint_score: 1.0,
            safety_score: 1.0,
            complexity_score: 1.0,
            overall: 0.0,
        };
        score.compute_overall();
        assert!((score.overall - 1.0).abs() < f64::EPSILON);
    }

    #[tokio::test]
    async fn test_quality_score_zero_when_invalid() {
        let mut score = QualityScore {
            syntax_valid: false,
            lint_score: 0.0,
            safety_score: 0.0,
            complexity_score: 0.0,
            overall: 0.0,
        };
        score.compute_overall();
        assert!((score.overall - 0.0).abs() < f64::EPSILON);
    }

    #[tokio::test]
    async fn test_quality_score_partial() {
        let mut score = QualityScore {
            syntax_valid: true,
            lint_score: 0.5,
            safety_score: 1.0,
            complexity_score: 0.5,
            overall: 0.0,
        };
        score.compute_overall();
        // 1.0 * 0.3 + 0.5 * 0.2 + 1.0 * 0.2 + 0.5 * 0.3 = 0.3 + 0.1 + 0.2 + 0.15 = 0.75
        assert!((score.overall - 0.75).abs() < f64::EPSILON);
    }

    // ─── Complexity scoring ──────────────────────────────────

    #[test]
    fn test_complexity_score_simple() {
        let score = compute_complexity_score("SELECT id FROM users");
        assert!((score - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_complexity_score_with_join() {
        let score =
            compute_complexity_score("SELECT u.id FROM users u JOIN orders o ON u.id = o.user_id");
        assert!(score < 1.0);
        assert!(score > 0.5);
    }

    #[test]
    fn test_complexity_score_with_subquery() {
        let score = compute_complexity_score(
            "SELECT id FROM users WHERE id IN (SELECT user_id FROM orders)",
        );
        assert!(score < 1.0);
    }

    #[test]
    fn test_complexity_score_cross_join() {
        let score = compute_complexity_score("SELECT * FROM users CROSS JOIN orders");
        assert!(score < 0.8);
    }

    #[test]
    fn test_complexity_score_never_negative() {
        let score = compute_complexity_score(
            "SELECT a FROM (SELECT b FROM (SELECT c FROM (SELECT d FROM t1 \
             CROSS JOIN t2 JOIN t3 ON t2.x = t3.y) UNION SELECT e FROM t4))",
        );
        assert!(score >= 0.0);
    }

    // ─── Lint score computation ──────────────────────────────

    #[test]
    fn test_lint_score_clean() {
        let result = crate::linter::types::LintResult {
            sql: "SELECT id FROM users".to_string(),
            findings: vec![],
            error_count: 0,
            warning_count: 0,
            clean: true,
        };
        let score = compute_lint_score(&result);
        assert!((score - 1.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_lint_score_with_errors() {
        use crate::linter::types::*;
        let result = LintResult {
            sql: "SELECT * FROM users".to_string(),
            findings: vec![LintFinding {
                code: "L001".to_string(),
                rule: "select_star".to_string(),
                severity: LintSeverity::Warning,
                message: "Avoid SELECT *".to_string(),
                suggestion: None,
                line: None,
                column: None,
            }],
            error_count: 0,
            warning_count: 1,
            clean: false,
        };
        let score = compute_lint_score(&result);
        assert!((score - 0.9).abs() < f64::EPSILON);
    }

    #[test]
    fn test_lint_score_never_negative() {
        use crate::linter::types::*;
        let findings: Vec<LintFinding> = (0..20)
            .map(|i| LintFinding {
                code: format!("L{:03}", i),
                rule: "test".to_string(),
                severity: LintSeverity::Error,
                message: "test".to_string(),
                suggestion: None,
                line: None,
                column: None,
            })
            .collect();
        let result = LintResult {
            sql: "test".to_string(),
            findings,
            error_count: 20,
            warning_count: 0,
            clean: false,
        };
        let score = compute_lint_score(&result);
        assert!(score >= 0.0);
    }

    // ─── Serialization ───────────────────────────────────────

    #[tokio::test]
    async fn test_correction_result_serializes_to_json() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECT * FROM usrs", &schema).await;
        let json = serde_json::to_string(&result);
        assert!(json.is_ok());
    }

    #[tokio::test]
    async fn test_correction_config_serializes() {
        let config = CorrectionConfig::default();
        let json = serde_json::to_string(&config).unwrap();
        assert!(json.contains("max_iterations"));
    }

    #[tokio::test]
    async fn test_correction_status_display() {
        assert_eq!(CorrectionStatus::AlreadyValid.to_string(), "already_valid");
        assert_eq!(CorrectionStatus::Fixed.to_string(), "fixed");
        assert_eq!(CorrectionStatus::PartialFix.to_string(), "partial_fix");
        assert_eq!(CorrectionStatus::Unfixable.to_string(), "unfixable");
    }

    #[tokio::test]
    async fn test_iteration_result_display() {
        assert_eq!(IterationResult::Fixed.to_string(), "fixed");
        assert_eq!(IterationResult::StillInvalid.to_string(), "still_invalid");
        assert_eq!(IterationResult::NewErrors.to_string(), "new_errors");
        assert_eq!(IterationResult::Skipped.to_string(), "skipped");
    }

    // ─── End-to-end flows ────────────────────────────────────

    #[tokio::test]
    async fn test_correct_preserves_valid_sql() {
        let schema = ecommerce_schema();
        let sql = "SELECT id, email, created_at FROM users WHERE id = 1";
        let result = correct_default(sql, &schema).await;
        assert_eq!(result.status, CorrectionStatus::AlreadyValid);
        assert!(result.corrected_sql.is_none());
    }

    #[tokio::test]
    async fn test_correct_final_validation_present() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECT * FROM usrs", &schema).await;
        assert!(result.final_validation.is_some());
    }

    #[tokio::test]
    async fn test_correct_time_tracking() {
        let schema = ecommerce_schema();
        let result = correct_default("SELECT * FROM usrs", &schema).await;
        assert!(result.total_time_ms > 0 || result.status == CorrectionStatus::AlreadyValid);
    }
}
