//! PII detection and classification engine.
//!
//! Scans columns for PII patterns by name, tags, and type hints.
//! Provides query-level PII checking and dialect-specific masking suggestions.

use super::types::*;
use crate::schema::types::{ColumnDef, SchemaDefinition, SqlDialect};

/// Classify all columns in a schema for PII.
///
/// Scans every column across all tables, applying name-pattern, tag-based,
/// and type-based detection rules. Returns a [`PiiReport`] with per-column
/// results and an overall risk level.
///
/// # Example
///
/// ```rust,no_run
/// use altimate_core::pii::classify_schema;
/// use altimate_core::SchemaDefinition;
///
/// let schema = SchemaDefinition::from_yaml_file("schema.yaml").unwrap();
/// let report = classify_schema(&schema);
/// println!("PII: {}/{} columns", report.pii_count, report.total_columns);
/// ```
pub fn classify_schema(schema: &SchemaDefinition) -> PiiReport {
    let mut columns = Vec::new();

    for (table_name, table_def) in &schema.tables {
        for col in &table_def.columns {
            let result = classify_column_internal(table_name, col, schema.dialect);
            columns.push(result);
        }
    }

    let pii_count = columns
        .iter()
        .filter(|c| c.classification != PiiClassification::None)
        .count();
    let total_columns = columns.len();
    let risk_level = compute_schema_risk(&columns);

    PiiReport {
        columns,
        pii_count,
        total_columns,
        risk_level,
    }
}

/// Classify a single column for PII.
///
/// This is the public API for classifying an individual column.
pub fn classify_column(table_name: &str, col: &ColumnDef, dialect: SqlDialect) -> PiiColumnResult {
    classify_column_internal(table_name, col, dialect)
}

/// Check if a SQL query accesses any PII columns.
///
/// Extracts column references from the SQL text (using simple identifier
/// matching against the schema) and checks each against PII classification.
///
/// # Example
///
/// ```rust,no_run
/// use altimate_core::pii::check_query_pii;
/// use altimate_core::SchemaDefinition;
///
/// let schema = SchemaDefinition::from_yaml_file("schema.yaml").unwrap();
/// let result = check_query_pii("SELECT email, name FROM users", &schema);
/// if result.accesses_pii {
///     println!("Warning: query accesses PII columns!");
/// }
/// ```
pub fn check_query_pii(sql: &str, schema: &SchemaDefinition) -> PiiQueryResult {
    let sql_lower = sql.to_lowercase();
    let mut pii_columns = Vec::new();
    let mut suggested_alternatives = Vec::new();

    for (table_name, table_def) in &schema.tables {
        // Check if this table is referenced in the SQL
        if !sql_contains_identifier(&sql_lower, &table_name.to_lowercase()) {
            continue;
        }

        for col in &table_def.columns {
            if sql_contains_identifier(&sql_lower, &col.name.to_lowercase()) {
                let classification = classify_column_internal(table_name, col, schema.dialect);
                if classification.classification != PiiClassification::None {
                    pii_columns.push(PiiColumnAccess {
                        table: table_name.clone(),
                        column: col.name.clone(),
                        classification: classification.classification.clone(),
                        suggested_masking: classification.suggested_masking.clone(),
                    });

                    // Generate a masking suggestion for the query
                    if let Some(ref masking) = classification.suggested_masking {
                        suggested_alternatives.push(format!(
                            "Replace `{}.{}` with `{}`",
                            table_name, col.name, masking
                        ));
                    }
                }
            }
        }
    }

    let accesses_pii = !pii_columns.is_empty();
    let risk_level = if pii_columns.is_empty() {
        PiiRiskLevel::None
    } else {
        pii_columns
            .iter()
            .map(|c| c.classification.risk_level())
            .max_by(|a, b| risk_ord(a).cmp(&risk_ord(b)))
            .unwrap_or(PiiRiskLevel::None)
    };

    PiiQueryResult {
        accesses_pii,
        pii_columns,
        suggested_alternatives,
        risk_level,
    }
}

/// Suggest a masking expression for a column in a given dialect.
///
/// Returns `None` if the column is not classified as PII.
pub fn suggest_masking(col: &ColumnDef, dialect: SqlDialect) -> Option<String> {
    let classification = classify_by_name(&col.name);
    let classification = if classification == PiiClassification::None {
        classify_by_tags(&col.tags)
    } else {
        classification
    };

    if classification == PiiClassification::None {
        return None;
    }

    Some(masking_expression(&col.name, &classification, dialect))
}

// --- Internal helpers ---

fn classify_column_internal(
    table_name: &str,
    col: &ColumnDef,
    dialect: SqlDialect,
) -> PiiColumnResult {
    // Priority: tag-based > name-pattern > type-pattern
    let (classification, confidence, method) = classify_column_all_methods(col);

    let suggested_masking = if classification != PiiClassification::None {
        Some(masking_expression(&col.name, &classification, dialect))
    } else {
        None
    };

    PiiColumnResult {
        table: table_name.to_string(),
        column: col.name.clone(),
        classification,
        confidence,
        detection_method: method.to_string(),
        suggested_masking,
    }
}

/// Apply all detection methods and return the highest-confidence match.
fn classify_column_all_methods(col: &ColumnDef) -> (PiiClassification, f64, &'static str) {
    // Tag-based detection (highest priority)
    let tag_class = classify_by_tags(&col.tags);
    if tag_class != PiiClassification::None {
        return (tag_class, 1.0, "tag");
    }

    // Name-pattern detection
    let name_class = classify_by_name(&col.name);
    if name_class != PiiClassification::None {
        let confidence = name_pattern_confidence(&col.name, &name_class);
        return (name_class, confidence, "name_pattern");
    }

    // Type-pattern detection
    let type_class = classify_by_type(&col.data_type);
    if type_class != PiiClassification::None {
        return (type_class, 0.5, "type_pattern");
    }

    (PiiClassification::None, 0.0, "none")
}

/// Classify by semantic tags on the column.
fn classify_by_tags(tags: &[String]) -> PiiClassification {
    let tags_lower: Vec<String> = tags.iter().map(|t| t.to_lowercase()).collect();

    if tags_lower.iter().any(|t| t == "pii_email" || t == "email")
        && tags_lower.contains(&"pii".to_string())
    {
        return PiiClassification::Email;
    }
    if tags_lower.iter().any(|t| t == "pii_phone") {
        return PiiClassification::Phone;
    }
    if tags_lower.iter().any(|t| t == "pii_ssn" || t == "ssn") {
        return PiiClassification::SSN;
    }
    if tags_lower
        .iter()
        .any(|t| t == "pii_financial" || t == "financial")
    {
        return PiiClassification::Financial;
    }
    if tags_lower
        .iter()
        .any(|t| t == "pii_health" || t == "health" || t == "hipaa")
    {
        return PiiClassification::Health;
    }
    if tags_lower.iter().any(|t| t == "pii_address") {
        return PiiClassification::Address;
    }
    if tags_lower.iter().any(|t| t == "pii_name") {
        return PiiClassification::Name;
    }
    if tags_lower.iter().any(|t| t == "pii_dob") {
        return PiiClassification::DateOfBirth;
    }
    if tags_lower.iter().any(|t| t == "pii_ip") {
        return PiiClassification::IPAddress;
    }

    // Generic "pii" tag without specific sub-type
    if tags_lower.contains(&"pii".to_string()) {
        return PiiClassification::Custom("unspecified".to_string());
    }

    PiiClassification::None
}

/// Classify by column name patterns.
fn classify_by_name(name: &str) -> PiiClassification {
    let lower = name.to_lowercase();

    // Email patterns
    if lower == "email"
        || lower == "email_address"
        || lower == "emailaddress"
        || lower == "user_email"
        || lower == "customer_email"
        || lower.ends_with("_email")
    {
        return PiiClassification::Email;
    }

    // Phone patterns
    if lower == "phone"
        || lower == "phone_number"
        || lower == "phonenumber"
        || lower == "mobile"
        || lower == "mobile_number"
        || lower == "cell_phone"
        || lower == "telephone"
        || lower.ends_with("_phone")
    {
        return PiiClassification::Phone;
    }

    // SSN patterns
    if lower == "ssn"
        || lower == "social_security"
        || lower == "social_security_number"
        || lower == "socialsecuritynumber"
        || lower == "tax_id"
        || lower == "national_id"
    {
        return PiiClassification::SSN;
    }

    // Name patterns
    if lower == "first_name"
        || lower == "firstname"
        || lower == "last_name"
        || lower == "lastname"
        || lower == "full_name"
        || lower == "fullname"
        || lower == "user_name"
        || lower == "username"
        || lower == "customer_name"
    {
        return PiiClassification::Name;
    }

    // Address patterns
    if lower == "address"
        || lower == "street_address"
        || lower == "home_address"
        || lower == "mailing_address"
        || lower == "zip_code"
        || lower == "zipcode"
        || lower == "postal_code"
        || lower == "city"
        || lower == "state"
    {
        return PiiClassification::Address;
    }

    // Financial patterns
    if lower == "credit_card"
        || lower == "creditcard"
        || lower == "card_number"
        || lower == "cardnumber"
        || lower == "salary"
        || lower == "income"
        || lower == "bank_account"
        || lower == "account_number"
        || lower == "routing_number"
        || lower == "cvv"
    {
        return PiiClassification::Financial;
    }

    // Date of birth patterns
    if lower == "dob"
        || lower == "date_of_birth"
        || lower == "dateofbirth"
        || lower == "birth_date"
        || lower == "birthdate"
        || lower == "birthday"
    {
        return PiiClassification::DateOfBirth;
    }

    // IP address patterns
    if lower == "ip_address"
        || lower == "ipaddress"
        || lower == "ip"
        || lower == "client_ip"
        || lower == "source_ip"
        || lower == "remote_addr"
    {
        return PiiClassification::IPAddress;
    }

    // Health patterns
    if lower == "diagnosis"
        || lower == "medical_record"
        || lower == "health_condition"
        || lower == "prescription"
        || lower == "patient_id"
    {
        return PiiClassification::Health;
    }

    // Passport
    if lower == "passport" || lower == "passport_number" {
        return PiiClassification::Custom("passport".to_string());
    }

    PiiClassification::None
}

/// Confidence for name-pattern detection based on how specific the match is.
fn name_pattern_confidence(name: &str, classification: &PiiClassification) -> f64 {
    let lower = name.to_lowercase();

    // Exact well-known names get high confidence
    match classification {
        PiiClassification::Email if lower == "email" || lower == "email_address" => 0.95,
        PiiClassification::Phone if lower == "phone" || lower == "phone_number" => 0.95,
        PiiClassification::SSN if lower == "ssn" || lower == "social_security_number" => 0.98,
        PiiClassification::Financial
            if lower == "credit_card" || lower == "card_number" || lower == "salary" =>
        {
            0.95
        }
        PiiClassification::Name
            if lower == "first_name" || lower == "last_name" || lower == "full_name" =>
        {
            0.85
        }
        PiiClassification::DateOfBirth if lower == "dob" || lower == "date_of_birth" => 0.9,
        PiiClassification::IPAddress if lower == "ip_address" || lower == "ip" => 0.8,
        PiiClassification::Address if lower == "address" || lower == "zip_code" => 0.8,
        PiiClassification::Health if lower == "diagnosis" || lower == "medical_record" => 0.9,
        _ => 0.75, // suffix/prefix patterns get lower confidence
    }
}

/// Classify by SQL data type patterns.
fn classify_by_type(_data_type: &str) -> PiiClassification {
    // Type-based classification is intentionally conservative.
    // We only flag types that are very strongly correlated with PII.
    // Most types (VARCHAR, INT, etc.) are too generic to be useful.
    PiiClassification::None
}

/// Generate a dialect-specific masking expression.
fn masking_expression(
    column_name: &str,
    classification: &PiiClassification,
    dialect: SqlDialect,
) -> String {
    match dialect {
        SqlDialect::Snowflake => masking_snowflake(column_name, classification),
        SqlDialect::BigQuery => masking_bigquery(column_name, classification),
        SqlDialect::Databricks | SqlDialect::Spark => {
            masking_databricks(column_name, classification)
        }
        SqlDialect::Postgres => masking_postgres(column_name, classification),
        SqlDialect::MySQL => masking_mysql(column_name, classification),
        SqlDialect::Redshift => masking_redshift(column_name, classification),
        _ => masking_generic(column_name, classification),
    }
}

fn masking_snowflake(col: &str, class: &PiiClassification) -> String {
    match class {
        PiiClassification::Email => format!("SHA2({col})"),
        PiiClassification::Phone => format!("REGEXP_REPLACE({col}, '.', '*')"),
        PiiClassification::SSN => format!("'***-**-' || RIGHT({col}, 4)"),
        PiiClassification::Name => format!("LEFT({col}, 1) || '***'"),
        PiiClassification::Financial => format!("'****' || RIGHT({col}, 4)"),
        PiiClassification::DateOfBirth => format!("DATE_TRUNC('YEAR', {col})"),
        PiiClassification::IPAddress => format!("REGEXP_REPLACE({col}, '\\\\d', '*')"),
        PiiClassification::Address => "'***REDACTED***'".to_string(),
        PiiClassification::Health => "'***REDACTED***'".to_string(),
        _ => format!("SHA2({col})"),
    }
}

fn masking_bigquery(col: &str, class: &PiiClassification) -> String {
    match class {
        PiiClassification::Email => format!("SHA256({col})"),
        PiiClassification::Phone => format!("REGEXP_REPLACE({col}, r'.', '*')"),
        PiiClassification::SSN => format!("CONCAT('***-**-', RIGHT({col}, 4))"),
        PiiClassification::Name => format!("CONCAT(LEFT({col}, 1), '***')"),
        PiiClassification::Financial => format!("CONCAT('****', RIGHT({col}, 4))"),
        PiiClassification::DateOfBirth => format!("DATE_TRUNC({col}, YEAR)"),
        PiiClassification::IPAddress => format!("REGEXP_REPLACE({col}, r'\\d', '*')"),
        PiiClassification::Address => "'***REDACTED***'".to_string(),
        PiiClassification::Health => "'***REDACTED***'".to_string(),
        _ => format!("SHA256({col})"),
    }
}

fn masking_databricks(col: &str, class: &PiiClassification) -> String {
    match class {
        PiiClassification::Email => format!("sha2({col}, 256)"),
        PiiClassification::Phone => format!("regexp_replace({col}, '.', '*')"),
        PiiClassification::SSN => format!("concat('***-**-', right({col}, 4))"),
        PiiClassification::Name => format!("concat(left({col}, 1), '***')"),
        PiiClassification::Financial => format!("concat('****', right({col}, 4))"),
        PiiClassification::DateOfBirth => format!("date_trunc('YEAR', {col})"),
        PiiClassification::IPAddress => format!("regexp_replace({col}, '\\\\d', '*')"),
        PiiClassification::Address => "'***REDACTED***'".to_string(),
        PiiClassification::Health => "'***REDACTED***'".to_string(),
        _ => format!("sha2({col}, 256)"),
    }
}

fn masking_postgres(col: &str, class: &PiiClassification) -> String {
    match class {
        PiiClassification::Email => format!("encode(sha256({col}::bytea), 'hex')"),
        PiiClassification::Phone => format!("regexp_replace({col}, '.', '*', 'g')"),
        PiiClassification::SSN => format!("'***-**-' || right({col}, 4)"),
        PiiClassification::Name => format!("left({col}, 1) || '***'"),
        PiiClassification::Financial => format!("'****' || right({col}, 4)"),
        PiiClassification::DateOfBirth => format!("date_trunc('year', {col})"),
        PiiClassification::IPAddress => format!("regexp_replace({col}, '\\d', '*', 'g')"),
        PiiClassification::Address => "'***REDACTED***'".to_string(),
        PiiClassification::Health => "'***REDACTED***'".to_string(),
        _ => format!("encode(sha256({col}::bytea), 'hex')"),
    }
}

fn masking_mysql(col: &str, class: &PiiClassification) -> String {
    match class {
        PiiClassification::Email => format!("SHA2({col}, 256)"),
        PiiClassification::Phone => {
            format!("REPLACE({col}, SUBSTR({col}, 1, LENGTH({col})-4), '****')")
        }
        PiiClassification::SSN => format!("CONCAT('***-**-', RIGHT({col}, 4))"),
        PiiClassification::Name => format!("CONCAT(LEFT({col}, 1), '***')"),
        PiiClassification::Financial => format!("CONCAT('****', RIGHT({col}, 4))"),
        PiiClassification::DateOfBirth => format!("DATE_FORMAT({col}, '%Y-01-01')"),
        PiiClassification::IPAddress => "'***.***.***.***'".to_string(),
        PiiClassification::Address => "'***REDACTED***'".to_string(),
        PiiClassification::Health => "'***REDACTED***'".to_string(),
        _ => format!("SHA2({col}, 256)"),
    }
}

fn masking_redshift(col: &str, class: &PiiClassification) -> String {
    match class {
        PiiClassification::Email => format!("SHA2({col}, 256)"),
        PiiClassification::Phone => format!("REGEXP_REPLACE({col}, '.', '*')"),
        PiiClassification::SSN => format!("'***-**-' || RIGHT({col}, 4)"),
        PiiClassification::Name => format!("LEFT({col}, 1) || '***'"),
        PiiClassification::Financial => format!("'****' || RIGHT({col}, 4)"),
        PiiClassification::DateOfBirth => format!("DATE_TRUNC('year', {col})"),
        PiiClassification::IPAddress => format!("REGEXP_REPLACE({col}, '\\\\d', '*')"),
        PiiClassification::Address => "'***REDACTED***'".to_string(),
        PiiClassification::Health => "'***REDACTED***'".to_string(),
        _ => format!("SHA2({col}, 256)"),
    }
}

fn masking_generic(_col: &str, _class: &PiiClassification) -> String {
    "'***MASKED***'".to_string()
}

/// Compute the overall risk level for a schema based on PII columns.
fn compute_schema_risk(columns: &[PiiColumnResult]) -> PiiRiskLevel {
    let has_high = columns
        .iter()
        .any(|c| c.classification.risk_level() == PiiRiskLevel::High);
    let has_medium = columns
        .iter()
        .any(|c| c.classification.risk_level() == PiiRiskLevel::Medium);
    let has_low = columns
        .iter()
        .any(|c| c.classification.risk_level() == PiiRiskLevel::Low);

    if has_high {
        PiiRiskLevel::High
    } else if has_medium {
        PiiRiskLevel::Medium
    } else if has_low {
        PiiRiskLevel::Low
    } else {
        PiiRiskLevel::None
    }
}

/// Ordering function for PiiRiskLevel.
fn risk_ord(level: &PiiRiskLevel) -> u8 {
    match level {
        PiiRiskLevel::None => 0,
        PiiRiskLevel::Low => 1,
        PiiRiskLevel::Medium => 2,
        PiiRiskLevel::High => 3,
    }
}

/// Check if a SQL string contains an identifier as a whole word.
fn sql_contains_identifier(sql_lower: &str, ident_lower: &str) -> bool {
    if ident_lower.is_empty() {
        return false;
    }
    let mut start = 0;
    while let Some(pos) = sql_lower[start..].find(ident_lower) {
        let abs_pos = start + pos;
        let end_pos = abs_pos + ident_lower.len();

        let before_ok = abs_pos == 0
            || (!sql_lower.as_bytes()[abs_pos - 1].is_ascii_alphanumeric()
                && sql_lower.as_bytes()[abs_pos - 1] != b'_');
        let after_ok = end_pos >= sql_lower.len()
            || (!sql_lower.as_bytes()[end_pos].is_ascii_alphanumeric()
                && sql_lower.as_bytes()[end_pos] != b'_');

        if before_ok && after_ok {
            return true;
        }

        start = abs_pos + 1;
        if start >= sql_lower.len() {
            break;
        }
    }
    false
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::schema::types::{ColumnDef, SchemaDefinition, SqlDialect, TableDef};
    use std::collections::HashMap;

    fn make_col(name: &str, data_type: &str, tags: Vec<&str>) -> ColumnDef {
        ColumnDef {
            name: name.to_string(),
            data_type: data_type.to_string(),
            nullable: true,
            description: None,
            tags: tags.into_iter().map(|s| s.to_string()).collect(),
            synonyms: vec![],

            statistics: None,
        }
    }

    fn make_pii_schema() -> SchemaDefinition {
        let mut tables = HashMap::new();
        tables.insert(
            "users".to_string(),
            TableDef {
                description: None,
                database: None,
                schema: None,
                columns: vec![
                    make_col("id", "INT", vec![]),
                    make_col("email", "VARCHAR", vec!["pii"]),
                    make_col("phone_number", "VARCHAR", vec![]),
                    make_col("ssn", "VARCHAR", vec![]),
                    make_col("first_name", "VARCHAR", vec![]),
                    make_col("last_name", "VARCHAR", vec![]),
                    make_col("address", "VARCHAR", vec![]),
                    make_col("salary", "DECIMAL(10,2)", vec![]),
                    make_col("dob", "DATE", vec![]),
                    make_col("ip_address", "VARCHAR", vec![]),
                    make_col("created_at", "TIMESTAMP", vec![]),
                    make_col("status", "VARCHAR", vec![]),
                ],
                primary_key: Some(vec!["id".to_string()]),
                foreign_keys: vec![],
                statistics: None,

                clustering_keys: vec![],
            },
        );
        tables.insert(
            "orders".to_string(),
            TableDef {
                description: None,
                database: None,
                schema: None,
                columns: vec![
                    make_col("id", "INT", vec![]),
                    make_col("user_id", "INT", vec![]),
                    make_col("total", "DECIMAL(10,2)", vec![]),
                    make_col("credit_card", "VARCHAR", vec![]),
                ],
                primary_key: Some(vec!["id".to_string()]),
                foreign_keys: vec![],
                statistics: None,

                clustering_keys: vec![],
            },
        );
        SchemaDefinition {
            version: "1".to_string(),
            dialect: SqlDialect::Generic,
            tables,
            database: None,
            schema_name: None,
            glossary: None,
        }
    }

    // === classify_by_name tests ===

    #[test]
    fn test_classify_email() {
        assert_eq!(classify_by_name("email"), PiiClassification::Email);
        assert_eq!(classify_by_name("email_address"), PiiClassification::Email);
        assert_eq!(classify_by_name("user_email"), PiiClassification::Email);
        assert_eq!(classify_by_name("customer_email"), PiiClassification::Email);
    }

    #[test]
    fn test_classify_phone() {
        assert_eq!(classify_by_name("phone"), PiiClassification::Phone);
        assert_eq!(classify_by_name("phone_number"), PiiClassification::Phone);
        assert_eq!(classify_by_name("mobile"), PiiClassification::Phone);
        assert_eq!(classify_by_name("cell_phone"), PiiClassification::Phone);
        assert_eq!(classify_by_name("telephone"), PiiClassification::Phone);
    }

    #[test]
    fn test_classify_ssn() {
        assert_eq!(classify_by_name("ssn"), PiiClassification::SSN);
        assert_eq!(classify_by_name("social_security"), PiiClassification::SSN);
        assert_eq!(
            classify_by_name("social_security_number"),
            PiiClassification::SSN
        );
        assert_eq!(classify_by_name("tax_id"), PiiClassification::SSN);
        assert_eq!(classify_by_name("national_id"), PiiClassification::SSN);
    }

    #[test]
    fn test_classify_name() {
        assert_eq!(classify_by_name("first_name"), PiiClassification::Name);
        assert_eq!(classify_by_name("last_name"), PiiClassification::Name);
        assert_eq!(classify_by_name("full_name"), PiiClassification::Name);
        assert_eq!(classify_by_name("username"), PiiClassification::Name);
        assert_eq!(classify_by_name("customer_name"), PiiClassification::Name);
    }

    #[test]
    fn test_classify_address() {
        assert_eq!(classify_by_name("address"), PiiClassification::Address);
        assert_eq!(
            classify_by_name("street_address"),
            PiiClassification::Address
        );
        assert_eq!(classify_by_name("zip_code"), PiiClassification::Address);
        assert_eq!(classify_by_name("postal_code"), PiiClassification::Address);
    }

    #[test]
    fn test_classify_financial() {
        assert_eq!(
            classify_by_name("credit_card"),
            PiiClassification::Financial
        );
        assert_eq!(
            classify_by_name("card_number"),
            PiiClassification::Financial
        );
        assert_eq!(classify_by_name("salary"), PiiClassification::Financial);
        assert_eq!(
            classify_by_name("bank_account"),
            PiiClassification::Financial
        );
        assert_eq!(classify_by_name("cvv"), PiiClassification::Financial);
    }

    #[test]
    fn test_classify_dob() {
        assert_eq!(classify_by_name("dob"), PiiClassification::DateOfBirth);
        assert_eq!(
            classify_by_name("date_of_birth"),
            PiiClassification::DateOfBirth
        );
        assert_eq!(classify_by_name("birthday"), PiiClassification::DateOfBirth);
        assert_eq!(
            classify_by_name("birth_date"),
            PiiClassification::DateOfBirth
        );
    }

    #[test]
    fn test_classify_ip() {
        assert_eq!(classify_by_name("ip_address"), PiiClassification::IPAddress);
        assert_eq!(classify_by_name("ip"), PiiClassification::IPAddress);
        assert_eq!(classify_by_name("client_ip"), PiiClassification::IPAddress);
        assert_eq!(classify_by_name("source_ip"), PiiClassification::IPAddress);
    }

    #[test]
    fn test_classify_health() {
        assert_eq!(classify_by_name("diagnosis"), PiiClassification::Health);
        assert_eq!(
            classify_by_name("medical_record"),
            PiiClassification::Health
        );
        assert_eq!(classify_by_name("patient_id"), PiiClassification::Health);
    }

    #[test]
    fn test_classify_passport() {
        assert_eq!(
            classify_by_name("passport"),
            PiiClassification::Custom("passport".to_string())
        );
        assert_eq!(
            classify_by_name("passport_number"),
            PiiClassification::Custom("passport".to_string())
        );
    }

    #[test]
    fn test_classify_non_pii() {
        assert_eq!(classify_by_name("id"), PiiClassification::None);
        assert_eq!(classify_by_name("created_at"), PiiClassification::None);
        assert_eq!(classify_by_name("status"), PiiClassification::None);
        assert_eq!(classify_by_name("total"), PiiClassification::None);
        assert_eq!(classify_by_name("user_id"), PiiClassification::None);
    }

    #[test]
    fn test_classify_case_insensitive() {
        assert_eq!(classify_by_name("Email"), PiiClassification::Email);
        assert_eq!(classify_by_name("PHONE_NUMBER"), PiiClassification::Phone);
        assert_eq!(classify_by_name("SSN"), PiiClassification::SSN);
    }

    // === classify_by_tags tests ===

    #[test]
    fn test_tag_pii_generic() {
        let class = classify_by_tags(&["pii".to_string()]);
        assert_eq!(class, PiiClassification::Custom("unspecified".to_string()));
    }

    #[test]
    fn test_tag_pii_email() {
        let class = classify_by_tags(&["pii".to_string(), "email".to_string()]);
        assert_eq!(class, PiiClassification::Email);
    }

    #[test]
    fn test_tag_pii_phone() {
        let class = classify_by_tags(&["pii_phone".to_string()]);
        assert_eq!(class, PiiClassification::Phone);
    }

    #[test]
    fn test_tag_pii_ssn() {
        let class = classify_by_tags(&["pii_ssn".to_string()]);
        assert_eq!(class, PiiClassification::SSN);
    }

    #[test]
    fn test_tag_financial() {
        let class = classify_by_tags(&["financial".to_string()]);
        assert_eq!(class, PiiClassification::Financial);
    }

    #[test]
    fn test_tag_hipaa() {
        let class = classify_by_tags(&["hipaa".to_string()]);
        assert_eq!(class, PiiClassification::Health);
    }

    #[test]
    fn test_tag_none() {
        let class = classify_by_tags(&["fiscal".to_string()]);
        assert_eq!(class, PiiClassification::None);
    }

    // === classify_schema tests ===

    #[test]
    fn test_classify_schema_counts() {
        let schema = make_pii_schema();
        let report = classify_schema(&schema);
        assert!(report.pii_count > 0);
        assert!(report.total_columns > report.pii_count);
    }

    #[test]
    fn test_classify_schema_finds_email() {
        let schema = make_pii_schema();
        let report = classify_schema(&schema);
        let email_col = report.columns.iter().find(|c| c.column == "email").unwrap();
        // email column has "pii" tag, so it's classified by tag
        assert_ne!(email_col.classification, PiiClassification::None);
    }

    #[test]
    fn test_classify_schema_finds_ssn() {
        let schema = make_pii_schema();
        let report = classify_schema(&schema);
        let ssn_col = report.columns.iter().find(|c| c.column == "ssn").unwrap();
        assert_eq!(ssn_col.classification, PiiClassification::SSN);
    }

    #[test]
    fn test_classify_schema_finds_credit_card() {
        let schema = make_pii_schema();
        let report = classify_schema(&schema);
        let cc_col = report
            .columns
            .iter()
            .find(|c| c.column == "credit_card")
            .unwrap();
        assert_eq!(cc_col.classification, PiiClassification::Financial);
    }

    #[test]
    fn test_classify_schema_non_pii_columns() {
        let schema = make_pii_schema();
        let report = classify_schema(&schema);
        let id_col = report
            .columns
            .iter()
            .find(|c| c.column == "id" && c.table == "users")
            .unwrap();
        assert_eq!(id_col.classification, PiiClassification::None);
    }

    #[test]
    fn test_classify_schema_risk_level() {
        let schema = make_pii_schema();
        let report = classify_schema(&schema);
        // SSN and Financial are high risk
        assert_eq!(report.risk_level, PiiRiskLevel::High);
    }

    #[test]
    fn test_classify_schema_no_pii() {
        let mut tables = HashMap::new();
        tables.insert(
            "stats".to_string(),
            TableDef {
                description: None,
                database: None,
                schema: None,
                columns: vec![
                    make_col("id", "INT", vec![]),
                    make_col("count", "INT", vec![]),
                    make_col("created_at", "TIMESTAMP", vec![]),
                ],
                primary_key: None,
                foreign_keys: vec![],
                statistics: None,

                clustering_keys: vec![],
            },
        );
        let schema = SchemaDefinition {
            version: "1".to_string(),
            dialect: SqlDialect::Generic,
            tables,
            database: None,
            schema_name: None,
            glossary: None,
        };
        let report = classify_schema(&schema);
        assert_eq!(report.pii_count, 0);
        assert_eq!(report.risk_level, PiiRiskLevel::None);
    }

    // === check_query_pii tests ===

    #[test]
    fn test_check_query_accesses_pii() {
        let schema = make_pii_schema();
        let result = check_query_pii("SELECT email, first_name FROM users", &schema);
        assert!(result.accesses_pii);
        assert!(result.pii_columns.len() >= 2);
    }

    #[test]
    fn test_check_query_no_pii() {
        let schema = make_pii_schema();
        let result = check_query_pii("SELECT id, status FROM users", &schema);
        assert!(!result.accesses_pii);
        assert!(result.pii_columns.is_empty());
    }

    #[test]
    fn test_check_query_risk_level() {
        let schema = make_pii_schema();
        let result = check_query_pii("SELECT ssn FROM users", &schema);
        assert_eq!(result.risk_level, PiiRiskLevel::High);
    }

    #[test]
    fn test_check_query_suggests_masking() {
        let schema = make_pii_schema();
        let result = check_query_pii("SELECT email FROM users", &schema);
        assert!(!result.suggested_alternatives.is_empty());
    }

    #[test]
    fn test_check_query_cross_table() {
        let schema = make_pii_schema();
        let result = check_query_pii(
            "SELECT u.email, o.credit_card FROM users u JOIN orders o ON u.id = o.user_id",
            &schema,
        );
        assert!(result.accesses_pii);
        // Should find PII from both tables
        let has_email = result.pii_columns.iter().any(|c| c.column == "email");
        let has_cc = result.pii_columns.iter().any(|c| c.column == "credit_card");
        assert!(has_email);
        assert!(has_cc);
    }

    #[test]
    fn test_check_query_unrelated_table() {
        let mut tables = HashMap::new();
        tables.insert(
            "logs".to_string(),
            TableDef {
                description: None,
                database: None,
                schema: None,
                columns: vec![make_col("message", "VARCHAR", vec![])],
                primary_key: None,
                foreign_keys: vec![],
                statistics: None,

                clustering_keys: vec![],
            },
        );
        tables.insert(
            "users".to_string(),
            TableDef {
                description: None,
                database: None,
                schema: None,
                columns: vec![make_col("email", "VARCHAR", vec![])],
                primary_key: None,
                foreign_keys: vec![],
                statistics: None,

                clustering_keys: vec![],
            },
        );
        let schema = SchemaDefinition {
            version: "1".to_string(),
            dialect: SqlDialect::Generic,
            tables,
            database: None,
            schema_name: None,
            glossary: None,
        };
        // Query only references "logs" table, not "users"
        let result = check_query_pii("SELECT message FROM logs", &schema);
        assert!(!result.accesses_pii);
    }

    // === suggest_masking tests ===

    #[test]
    fn test_suggest_masking_email_snowflake() {
        let col = make_col("email", "VARCHAR", vec![]);
        let masking = suggest_masking(&col, SqlDialect::Snowflake);
        assert!(masking.is_some());
        assert!(masking.unwrap().contains("SHA2"));
    }

    #[test]
    fn test_suggest_masking_email_bigquery() {
        let col = make_col("email", "VARCHAR", vec![]);
        let masking = suggest_masking(&col, SqlDialect::BigQuery);
        assert!(masking.is_some());
        assert!(masking.unwrap().contains("SHA256"));
    }

    #[test]
    fn test_suggest_masking_email_databricks() {
        let col = make_col("email", "VARCHAR", vec![]);
        let masking = suggest_masking(&col, SqlDialect::Databricks);
        assert!(masking.is_some());
        assert!(masking.unwrap().contains("sha2"));
    }

    #[test]
    fn test_suggest_masking_email_postgres() {
        let col = make_col("email", "VARCHAR", vec![]);
        let masking = suggest_masking(&col, SqlDialect::Postgres);
        assert!(masking.is_some());
        assert!(masking.unwrap().contains("sha256"));
    }

    #[test]
    fn test_suggest_masking_email_mysql() {
        let col = make_col("email", "VARCHAR", vec![]);
        let masking = suggest_masking(&col, SqlDialect::MySQL);
        assert!(masking.is_some());
        assert!(masking.unwrap().contains("SHA2"));
    }

    #[test]
    fn test_suggest_masking_email_generic() {
        let col = make_col("email", "VARCHAR", vec![]);
        let masking = suggest_masking(&col, SqlDialect::Generic);
        assert!(masking.is_some());
        assert!(masking.unwrap().contains("MASKED"));
    }

    #[test]
    fn test_suggest_masking_phone_snowflake() {
        let col = make_col("phone", "VARCHAR", vec![]);
        let masking = suggest_masking(&col, SqlDialect::Snowflake);
        assert!(masking.is_some());
        assert!(masking.unwrap().contains("REGEXP_REPLACE"));
    }

    #[test]
    fn test_suggest_masking_ssn_bigquery() {
        let col = make_col("ssn", "VARCHAR", vec![]);
        let masking = suggest_masking(&col, SqlDialect::BigQuery);
        assert!(masking.is_some());
        assert!(masking.unwrap().contains("***-**-"));
    }

    #[test]
    fn test_suggest_masking_non_pii() {
        let col = make_col("id", "INT", vec![]);
        let masking = suggest_masking(&col, SqlDialect::Snowflake);
        assert!(masking.is_none());
    }

    #[test]
    fn test_suggest_masking_credit_card_snowflake() {
        let col = make_col("credit_card", "VARCHAR", vec![]);
        let masking = suggest_masking(&col, SqlDialect::Snowflake);
        assert!(masking.is_some());
        assert!(masking.unwrap().contains("RIGHT"));
    }

    #[test]
    fn test_suggest_masking_dob_bigquery() {
        let col = make_col("dob", "DATE", vec![]);
        let masking = suggest_masking(&col, SqlDialect::BigQuery);
        assert!(masking.is_some());
        assert!(masking.unwrap().contains("DATE_TRUNC"));
    }

    #[test]
    fn test_suggest_masking_redshift() {
        let col = make_col("email", "VARCHAR", vec![]);
        let masking = suggest_masking(&col, SqlDialect::Redshift);
        assert!(masking.is_some());
        assert!(masking.unwrap().contains("SHA2"));
    }

    #[test]
    fn test_suggest_masking_name_postgres() {
        let col = make_col("first_name", "VARCHAR", vec![]);
        let masking = suggest_masking(&col, SqlDialect::Postgres);
        assert!(masking.is_some());
        assert!(masking.unwrap().contains("***"));
    }

    #[test]
    fn test_suggest_masking_address_generic() {
        let col = make_col("address", "VARCHAR", vec![]);
        let masking = suggest_masking(&col, SqlDialect::Generic);
        assert!(masking.is_some());
        assert!(masking.unwrap().contains("MASKED"));
    }

    // === PiiClassification risk_level tests ===

    #[test]
    fn test_risk_level_ssn() {
        assert_eq!(PiiClassification::SSN.risk_level(), PiiRiskLevel::High);
    }

    #[test]
    fn test_risk_level_financial() {
        assert_eq!(
            PiiClassification::Financial.risk_level(),
            PiiRiskLevel::High
        );
    }

    #[test]
    fn test_risk_level_email() {
        assert_eq!(PiiClassification::Email.risk_level(), PiiRiskLevel::Medium);
    }

    #[test]
    fn test_risk_level_name() {
        assert_eq!(PiiClassification::Name.risk_level(), PiiRiskLevel::Low);
    }

    #[test]
    fn test_risk_level_none() {
        assert_eq!(PiiClassification::None.risk_level(), PiiRiskLevel::None);
    }

    // === Display trait tests ===

    #[test]
    fn test_pii_classification_display() {
        assert_eq!(PiiClassification::Email.to_string(), "email");
        assert_eq!(PiiClassification::Phone.to_string(), "phone");
        assert_eq!(PiiClassification::SSN.to_string(), "ssn");
        assert_eq!(PiiClassification::Name.to_string(), "name");
        assert_eq!(PiiClassification::Address.to_string(), "address");
        assert_eq!(PiiClassification::Financial.to_string(), "financial");
        assert_eq!(PiiClassification::Health.to_string(), "health");
        assert_eq!(PiiClassification::DateOfBirth.to_string(), "date_of_birth");
        assert_eq!(PiiClassification::IPAddress.to_string(), "ip_address");
        assert_eq!(
            PiiClassification::Custom("passport".to_string()).to_string(),
            "custom:passport"
        );
        assert_eq!(PiiClassification::None.to_string(), "none");
    }

    #[test]
    fn test_pii_risk_level_display() {
        assert_eq!(PiiRiskLevel::None.to_string(), "none");
        assert_eq!(PiiRiskLevel::Low.to_string(), "low");
        assert_eq!(PiiRiskLevel::Medium.to_string(), "medium");
        assert_eq!(PiiRiskLevel::High.to_string(), "high");
    }

    // === Serialization tests ===

    #[test]
    fn test_pii_report_serialize() {
        let report = PiiReport {
            columns: vec![],
            pii_count: 0,
            total_columns: 0,
            risk_level: PiiRiskLevel::None,
        };
        let json = serde_json::to_string(&report).expect("serialize");
        assert!(json.contains("pii_count"));
    }

    #[test]
    fn test_pii_column_result_serialize() {
        let result = PiiColumnResult {
            table: "users".to_string(),
            column: "email".to_string(),
            classification: PiiClassification::Email,
            confidence: 0.95,
            detection_method: "name_pattern".to_string(),
            suggested_masking: Some("SHA2(email)".to_string()),
        };
        let json = serde_json::to_string(&result).expect("serialize");
        assert!(json.contains("Email"));
        assert!(json.contains("SHA2"));
    }

    #[test]
    fn test_pii_query_result_serialize() {
        let result = PiiQueryResult {
            accesses_pii: true,
            pii_columns: vec![],
            suggested_alternatives: vec![],
            risk_level: PiiRiskLevel::High,
        };
        let json = serde_json::to_string(&result).expect("serialize");
        assert!(json.contains("accesses_pii"));
    }

    // === classify_column public API ===

    #[test]
    fn test_classify_column_public() {
        let col = make_col("email", "VARCHAR", vec!["pii"]);
        let result = classify_column("users", &col, SqlDialect::Snowflake);
        assert_eq!(result.table, "users");
        assert_eq!(result.column, "email");
        assert_ne!(result.classification, PiiClassification::None);
        assert!(result.suggested_masking.is_some());
    }

    #[test]
    fn test_classify_column_tag_priority() {
        // Tag-based should take priority over name-pattern
        let col = make_col("email", "VARCHAR", vec!["pii_ssn"]);
        let result = classify_column("users", &col, SqlDialect::Generic);
        assert_eq!(result.classification, PiiClassification::SSN);
        assert_eq!(result.detection_method, "tag");
    }

    #[test]
    fn test_classify_column_name_pattern_fallback() {
        let col = make_col("email", "VARCHAR", vec![]);
        let result = classify_column("users", &col, SqlDialect::Generic);
        assert_eq!(result.classification, PiiClassification::Email);
        assert_eq!(result.detection_method, "name_pattern");
    }

    // === sql_contains_identifier tests ===

    #[test]
    fn test_sql_identifier_basic() {
        assert!(sql_contains_identifier("select email from users", "email"));
        assert!(sql_contains_identifier("select email from users", "users"));
    }

    #[test]
    fn test_sql_identifier_no_partial() {
        assert!(!sql_contains_identifier(
            "select user_id from orders",
            "user"
        ));
    }

    #[test]
    fn test_sql_identifier_empty() {
        assert!(!sql_contains_identifier("select 1", ""));
        assert!(!sql_contains_identifier("", "email"));
    }

    // === name_pattern_confidence tests ===

    #[test]
    fn test_confidence_exact_email() {
        let c = name_pattern_confidence("email", &PiiClassification::Email);
        assert!(c >= 0.9);
    }

    #[test]
    fn test_confidence_suffix_email() {
        let c = name_pattern_confidence("customer_email", &PiiClassification::Email);
        assert!(c >= 0.7);
    }

    #[test]
    fn test_confidence_ssn_highest() {
        let c = name_pattern_confidence("ssn", &PiiClassification::SSN);
        assert!(c >= 0.95);
    }
}
