//! Query pattern guardrail rules.
//!
//! These rules restrict which SQL statement types are allowed or blocked
//! based on the policy configuration.

use crate::explainer::types::QueryExplanation;
use crate::schema::types::SchemaDefinition;

use super::super::policy::PolicyConfig;
use super::super::types::{PolicyCategory, PolicyViolation, ViolationDetail, ViolationSeverity};
use super::{GuardrailRule, RuleResult};

/// Extract the effective statement type from a SQL string.
///
/// Handles leading whitespace, single-line comments (`--`),
/// block comments (`/* */`), and CTEs (`WITH ... AS ... SELECT`).
/// Returns the uppercase first keyword (e.g., "SELECT", "INSERT").
fn extract_statement_type(sql: &str) -> String {
    let mut s = sql.trim();

    // Strip leading single-line comments
    while s.starts_with("--") {
        if let Some(newline_pos) = s.find('\n') {
            s = s[newline_pos + 1..].trim();
        } else {
            // Entire string is a comment
            return String::new();
        }
    }

    // Strip leading block comments
    while s.starts_with("/*") {
        if let Some(end_pos) = s.find("*/") {
            s = s[end_pos + 2..].trim();
        } else {
            // Unclosed block comment
            return String::new();
        }
    }

    // Extract the first keyword
    let first_word: String = s
        .chars()
        .take_while(|c| c.is_ascii_alphanumeric() || *c == '_')
        .collect();

    let keyword = first_word.to_uppercase();

    // CTEs start with WITH but the effective statement is whatever follows
    // For guardrail purposes, WITH ... SELECT is a SELECT
    if keyword == "WITH" {
        return "SELECT".to_string();
    }

    keyword
}

/// Rule: allowed_statements — only allow specific statement types.
pub struct AllowedStatementsRule;

impl GuardrailRule for AllowedStatementsRule {
    fn name(&self) -> &str {
        "allowed_statements"
    }

    fn category(&self) -> PolicyCategory {
        PolicyCategory::QueryPatterns
    }

    fn evaluate(
        &self,
        sql: &str,
        _explanation: &QueryExplanation,
        _schema: &SchemaDefinition,
        config: &PolicyConfig,
    ) -> RuleResult {
        let mut violations = Vec::new();

        let qp = match &config.query_patterns {
            Some(qp) => qp,
            None => {
                return RuleResult {
                    violations,
                    warnings: Vec::new(),
                };
            }
        };

        if let Some(allowed) = &qp.allowed_statements {
            let stmt_type = extract_statement_type(sql);
            if !stmt_type.is_empty() {
                let is_allowed = allowed.iter().any(|a| a.to_uppercase() == stmt_type);
                if !is_allowed {
                    violations.push(PolicyViolation {
                        rule: self.name().to_string(),
                        category: self.category(),
                        severity: ViolationSeverity::Error,
                        message: format!("{stmt_type} statements are not allowed by policy"),
                        remediation: Some(format!(
                            "Only the following statement types are allowed: {}",
                            allowed.join(", ")
                        )),
                        detail: ViolationDetail::BlockedStatement {
                            statement: stmt_type,
                        },
                    });
                }
            }
        }

        RuleResult {
            violations,
            warnings: Vec::new(),
        }
    }
}

/// Rule: block_statements — block specific statement types.
pub struct BlockStatementsRule;

impl GuardrailRule for BlockStatementsRule {
    fn name(&self) -> &str {
        "block_statements"
    }

    fn category(&self) -> PolicyCategory {
        PolicyCategory::QueryPatterns
    }

    fn evaluate(
        &self,
        sql: &str,
        _explanation: &QueryExplanation,
        _schema: &SchemaDefinition,
        config: &PolicyConfig,
    ) -> RuleResult {
        let mut violations = Vec::new();

        let qp = match &config.query_patterns {
            Some(qp) => qp,
            None => {
                return RuleResult {
                    violations,
                    warnings: Vec::new(),
                };
            }
        };

        if let Some(blocked) = &qp.block_statements {
            let stmt_type = extract_statement_type(sql);
            if !stmt_type.is_empty() {
                let is_blocked = blocked.iter().any(|b| b.to_uppercase() == stmt_type);
                if is_blocked {
                    violations.push(PolicyViolation {
                        rule: self.name().to_string(),
                        category: self.category(),
                        severity: ViolationSeverity::Error,
                        message: format!("{stmt_type} statements are blocked by policy"),
                        remediation: Some(format!(
                            "The following statement types are blocked: {}",
                            blocked.join(", ")
                        )),
                        detail: ViolationDetail::BlockedStatement {
                            statement: stmt_type,
                        },
                    });
                }
            }
        }

        RuleResult {
            violations,
            warnings: Vec::new(),
        }
    }
}

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

    #[test]
    fn test_extract_select() {
        assert_eq!(extract_statement_type("SELECT * FROM users"), "SELECT");
    }

    #[test]
    fn test_extract_with_whitespace() {
        assert_eq!(extract_statement_type("  \n  SELECT 1"), "SELECT");
    }

    #[test]
    fn test_extract_with_comment() {
        assert_eq!(extract_statement_type("-- comment\nSELECT 1"), "SELECT");
    }

    #[test]
    fn test_extract_with_block_comment() {
        assert_eq!(
            extract_statement_type("/* block comment */ SELECT 1"),
            "SELECT"
        );
    }

    #[test]
    fn test_extract_cte() {
        assert_eq!(
            extract_statement_type("WITH cte AS (SELECT 1) SELECT * FROM cte"),
            "SELECT"
        );
    }

    #[test]
    fn test_extract_insert() {
        assert_eq!(
            extract_statement_type("INSERT INTO users (id) VALUES (1)"),
            "INSERT"
        );
    }

    #[test]
    fn test_extract_delete() {
        assert_eq!(
            extract_statement_type("DELETE FROM users WHERE id = 1"),
            "DELETE"
        );
    }

    #[test]
    fn test_extract_case_insensitive() {
        assert_eq!(extract_statement_type("select * from users"), "SELECT");
    }
}
