//! Per-dialect pricing profiles for cloud data warehouses.
//!
//! Built-in pricing for BigQuery, Snowflake, Redshift, Athena, Databricks,
//! and a generic fallback. Supports credit-based (Snowflake), DBU-based
//! (Databricks), provisioned (Redshift), and per-TB (BigQuery, Athena) models.

use crate::schema::types::SqlDialect;
use serde::{Deserialize, Serialize};

/// Pricing model used by a cloud data warehouse.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum PricingModel {
    /// BigQuery, Athena -- simple $/TB scanned.
    PerTbScanned {
        /// Cost per terabyte scanned in USD.
        cost_per_tb: f64,
    },
    /// Snowflake -- credits consumed per TB, priced per credit.
    CreditBased {
        /// Credits consumed per TB scanned.
        credits_per_tb: f64,
        /// Default price per credit in USD.
        default_credit_price: f64,
    },
    /// Databricks -- DBU consumed per TB, priced per DBU.
    DbuBased {
        /// DBU consumed per TB scanned.
        dbu_per_tb: f64,
        /// Default price per DBU in USD.
        default_dbu_price: f64,
    },
    /// Redshift -- provisioned $/TB with a minimum charge.
    Provisioned {
        /// Cost per terabyte scanned in USD.
        cost_per_tb: f64,
        /// Minimum charge in USD for any query.
        minimum_charge: f64,
    },
}

impl PricingModel {
    /// Compute the effective cost per TB in USD using defaults.
    pub fn effective_cost_per_tb(&self) -> f64 {
        match self {
            PricingModel::PerTbScanned { cost_per_tb } => *cost_per_tb,
            PricingModel::CreditBased {
                credits_per_tb,
                default_credit_price,
            } => credits_per_tb * default_credit_price,
            PricingModel::DbuBased {
                dbu_per_tb,
                default_dbu_price,
            } => dbu_per_tb * default_dbu_price,
            PricingModel::Provisioned { cost_per_tb, .. } => *cost_per_tb,
        }
    }
}

/// Options to override default credit/DBU pricing.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct PricingOptions {
    /// Override for Snowflake credit price ($/credit).
    pub credit_price: Option<f64>,
    /// Override for Databricks DBU price ($/DBU).
    pub dbu_price: Option<f64>,
}

/// A pricing profile for a cloud data warehouse dialect.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CostProfile {
    /// Dialect name.
    pub dialect: String,

    /// Effective cost per terabyte scanned in USD (computed from pricing model).
    pub cost_per_tb: f64,

    /// Human-readable pricing description.
    pub description: String,

    /// Whether window functions have reduced cost on this dialect.
    pub cheap_window_functions: bool,

    /// Whether CTEs are materialized (more expensive) or inlined.
    pub ctes_materialized: bool,

    /// The underlying pricing model.
    pub pricing_model: PricingModel,

    /// Whether this dialect uses columnar storage (micro-partition scanning).
    ///
    /// Columnar stores (Snowflake, BigQuery, Redshift, etc.) scan entire
    /// micro-partitions regardless of filter selectivity. Filter selectivity
    /// reduces the number of *rows returned* but not the *bytes scanned*
    /// from storage. Only partition pruning reduces storage I/O.
    ///
    /// When `true`, the cost estimator applies a floor: bytes scanned for a
    /// table can never be less than `row_count * avg_row_bytes`.
    pub is_columnar: bool,
}

impl CostProfile {
    /// Get the built-in cost profile for a dialect.
    pub fn for_dialect(dialect: SqlDialect) -> Self {
        match dialect {
            SqlDialect::BigQuery => {
                let model = PricingModel::PerTbScanned { cost_per_tb: 6.25 };
                Self {
                    dialect: "bigquery".to_string(),
                    cost_per_tb: model.effective_cost_per_tb(),
                    description: "BigQuery on-demand: $6.25/TB scanned".to_string(),
                    cheap_window_functions: true,
                    ctes_materialized: false,
                    pricing_model: model,
                    is_columnar: true,
                }
            }
            SqlDialect::Snowflake => {
                let model = PricingModel::CreditBased {
                    credits_per_tb: 0.6,
                    default_credit_price: 3.00,
                };
                Self {
                    dialect: "snowflake".to_string(),
                    cost_per_tb: model.effective_cost_per_tb(),
                    description: "Snowflake: 0.6 credits/TB, $3.00/credit (effective $1.80/TB)"
                        .to_string(),
                    cheap_window_functions: true,
                    ctes_materialized: true,
                    pricing_model: model,
                    is_columnar: true,
                }
            }
            SqlDialect::Redshift => {
                let model = PricingModel::Provisioned {
                    cost_per_tb: 1.086,
                    minimum_charge: 0.01,
                };
                Self {
                    dialect: "redshift".to_string(),
                    cost_per_tb: model.effective_cost_per_tb(),
                    description: "Redshift Serverless: $1.086/TB, $0.01 minimum".to_string(),
                    cheap_window_functions: false,
                    ctes_materialized: false,
                    pricing_model: model,
                    is_columnar: true,
                }
            }
            SqlDialect::Athena => {
                let model = PricingModel::PerTbScanned { cost_per_tb: 5.00 };
                Self {
                    dialect: "athena".to_string(),
                    cost_per_tb: model.effective_cost_per_tb(),
                    description: "Athena: $5.00/TB scanned".to_string(),
                    cheap_window_functions: false,
                    ctes_materialized: false,
                    pricing_model: model,
                    is_columnar: true,
                }
            }
            SqlDialect::Databricks => {
                let model = PricingModel::DbuBased {
                    dbu_per_tb: 5.5,
                    default_dbu_price: 0.55,
                };
                Self {
                    dialect: "databricks".to_string(),
                    cost_per_tb: model.effective_cost_per_tb(),
                    description: "Databricks SQL Pro: 5.5 DBU/TB, $0.55/DBU (effective $3.025/TB)"
                        .to_string(),
                    cheap_window_functions: true,
                    ctes_materialized: false,
                    pricing_model: model,
                    is_columnar: true,
                }
            }
            SqlDialect::Trino | SqlDialect::Presto => {
                let model = PricingModel::PerTbScanned { cost_per_tb: 5.00 };
                Self {
                    dialect: dialect.to_string(),
                    cost_per_tb: model.effective_cost_per_tb(),
                    description: format!("{dialect}: ~$5.00/TB (Athena-equivalent)"),
                    cheap_window_functions: false,
                    ctes_materialized: false,
                    pricing_model: model,
                    is_columnar: true,
                }
            }
            _ => Self::generic(),
        }
    }

    /// Generic/fallback cost profile.
    pub fn generic() -> Self {
        let model = PricingModel::PerTbScanned { cost_per_tb: 5.00 };
        Self {
            dialect: "generic".to_string(),
            cost_per_tb: model.effective_cost_per_tb(),
            description: "Generic estimate: $5.00/TB scanned".to_string(),
            cheap_window_functions: false,
            ctes_materialized: false,
            pricing_model: model,
            is_columnar: false,
        }
    }

    /// Create a Snowflake profile with a custom credit price.
    pub fn snowflake_custom(credit_price: f64) -> Self {
        let model = PricingModel::CreditBased {
            credits_per_tb: 0.6,
            default_credit_price: credit_price,
        };
        Self {
            dialect: "snowflake".to_string(),
            cost_per_tb: model.effective_cost_per_tb(),
            description: format!(
                "Snowflake: 0.6 credits/TB, ${:.2}/credit (effective ${:.3}/TB)",
                credit_price,
                0.6 * credit_price
            ),
            cheap_window_functions: true,
            ctes_materialized: true,
            pricing_model: model,
            is_columnar: true,
        }
    }

    /// Create a Databricks profile for a specific pricing tier.
    ///
    /// Supported tiers: "classic" ($0.22/DBU), "pro" ($0.55/DBU), "serverless" ($0.70/DBU).
    /// Falls back to "pro" for unrecognized tiers.
    pub fn databricks_tier(tier: &str) -> Self {
        let dbu_price = match tier {
            "classic" => 0.22,
            "pro" => 0.55,
            "serverless" => 0.70,
            _ => 0.55, // default to pro
        };
        let model = PricingModel::DbuBased {
            dbu_per_tb: 5.5,
            default_dbu_price: dbu_price,
        };
        Self {
            dialect: "databricks".to_string(),
            cost_per_tb: model.effective_cost_per_tb(),
            description: format!(
                "Databricks SQL {}: 5.5 DBU/TB, ${:.2}/DBU (effective ${:.3}/TB)",
                tier,
                dbu_price,
                5.5 * dbu_price
            ),
            cheap_window_functions: true,
            ctes_materialized: false,
            pricing_model: model,
            is_columnar: true,
        }
    }

    /// Calculate USD cost from bytes scanned using default pricing.
    pub fn calculate_cost(&self, bytes_scanned: u64) -> f64 {
        self.calculate_cost_with_options(bytes_scanned, &PricingOptions::default())
    }

    /// Calculate USD cost from bytes scanned with custom pricing overrides.
    pub fn calculate_cost_with_options(&self, bytes_scanned: u64, options: &PricingOptions) -> f64 {
        let tb = bytes_scanned as f64 / (1024.0 * 1024.0 * 1024.0 * 1024.0);
        match &self.pricing_model {
            PricingModel::PerTbScanned { cost_per_tb } => tb * cost_per_tb,
            PricingModel::CreditBased {
                credits_per_tb,
                default_credit_price,
            } => {
                let price = options.credit_price.unwrap_or(*default_credit_price);
                tb * credits_per_tb * price
            }
            PricingModel::DbuBased {
                dbu_per_tb,
                default_dbu_price,
            } => {
                let price = options.dbu_price.unwrap_or(*default_dbu_price);
                tb * dbu_per_tb * price
            }
            PricingModel::Provisioned {
                cost_per_tb,
                minimum_charge,
            } => {
                let raw = tb * cost_per_tb;
                if bytes_scanned == 0 {
                    0.0
                } else {
                    raw.max(*minimum_charge)
                }
            }
        }
    }

    /// Return the effective cost per TB in USD (computed from pricing model defaults).
    pub fn cost_per_tb_effective(&self) -> f64 {
        self.pricing_model.effective_cost_per_tb()
    }

    /// Return a disclaimer string for pricing models that use credits or DBUs.
    ///
    /// Returns `None` for simple per-TB models that need no disclaimer.
    pub fn disclaimer(&self) -> Option<String> {
        match &self.pricing_model {
            PricingModel::CreditBased {
                credits_per_tb,
                default_credit_price,
            } => Some(format!(
                "Cost estimate uses {credits_per_tb} credits/TB at ${default_credit_price:.2}/credit. \
                 Actual cost depends on your Snowflake contract and warehouse size."
            )),
            PricingModel::DbuBased {
                dbu_per_tb,
                default_dbu_price,
            } => Some(format!(
                "Cost estimate uses {dbu_per_tb} DBU/TB at ${default_dbu_price:.2}/DBU. \
                 Actual cost depends on your Databricks tier and cluster configuration."
            )),
            PricingModel::Provisioned { minimum_charge, .. } => Some(format!(
                "Redshift Serverless has a ${minimum_charge:.2} minimum charge per query."
            )),
            PricingModel::PerTbScanned { .. } => None,
        }
    }
}

/// Snowflake warehouse sizes with throughput and credit consumption.
/// Throughput values derived from production ClickHouse data (60.3M rows).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SnowflakeWarehouseSize {
    XSmall,
    Small,
    Medium,
    Large,
    XLarge,
    TwoXLarge,
    ThreeXLarge,
    FourXLarge,
}

impl SnowflakeWarehouseSize {
    /// Throughput in MB/s (median bytes_scanned / median execution_time per size).
    pub fn throughput_mb_per_sec(&self) -> f64 {
        match self {
            Self::XSmall => 126.0,
            Self::Small => 144.0,
            Self::Medium => 183.0,
            Self::Large => 1_059.0,
            Self::XLarge => 2_447.0,
            Self::TwoXLarge => 2_580.0,
            Self::ThreeXLarge => 3_640.0,
            Self::FourXLarge => 5_000.0,
        }
    }

    /// Credits consumed per hour.
    pub fn credits_per_hour(&self) -> f64 {
        match self {
            Self::XSmall => 1.0,
            Self::Small => 2.0,
            Self::Medium => 4.0,
            Self::Large => 8.0,
            Self::XLarge => 16.0,
            Self::TwoXLarge => 32.0,
            Self::ThreeXLarge => 64.0,
            Self::FourXLarge => 128.0,
        }
    }
}

impl std::fmt::Display for SnowflakeWarehouseSize {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::XSmall => write!(f, "x_small"),
            Self::Small => write!(f, "small"),
            Self::Medium => write!(f, "medium"),
            Self::Large => write!(f, "large"),
            Self::XLarge => write!(f, "x_large"),
            Self::TwoXLarge => write!(f, "two_x_large"),
            Self::ThreeXLarge => write!(f, "three_x_large"),
            Self::FourXLarge => write!(f, "four_x_large"),
        }
    }
}

/// SQL query type with cost multiplier relative to SELECT.
/// Multipliers derived from ClickHouse production data (60.3M rows, Rakuten Snowflake).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum QueryType {
    Select,
    CreateTableAs,
    Merge,
    Delete,
    Update,
    Insert,
    Copy,
    Other,
}

impl QueryType {
    /// Cost multiplier relative to SELECT (from production $/GB data).
    pub fn cost_multiplier(&self) -> f64 {
        match self {
            Self::Select => 1.0,
            Self::CreateTableAs => 1.7,
            Self::Merge => 15.0,
            Self::Delete => 30.0,
            Self::Update => 30.0,
            Self::Insert => 63.0,
            Self::Copy => 253.0,
            Self::Other => 1.0,
        }
    }

    /// Detect query type from SQL text.
    /// Strips leading comments and whitespace, then checks first keyword(s).
    pub fn from_sql(sql: &str) -> Self {
        let stripped = strip_sql_comments(sql);
        let upper = stripped.trim().to_uppercase();
        let first_word = upper.split_whitespace().next().unwrap_or("");
        match first_word {
            "SELECT" | "WITH" => Self::Select,
            "CREATE" if upper.contains(" AS SELECT") || upper.contains(" AS (") => {
                Self::CreateTableAs
            }
            "MERGE" => Self::Merge,
            "DELETE" => Self::Delete,
            "UPDATE" => Self::Update,
            "INSERT" => Self::Insert,
            "COPY" => Self::Copy,
            _ => Self::Other,
        }
    }
}

impl std::fmt::Display for QueryType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Select => write!(f, "select"),
            Self::CreateTableAs => write!(f, "create_table_as"),
            Self::Merge => write!(f, "merge"),
            Self::Delete => write!(f, "delete"),
            Self::Update => write!(f, "update"),
            Self::Insert => write!(f, "insert"),
            Self::Copy => write!(f, "copy"),
            Self::Other => write!(f, "other"),
        }
    }
}

/// Strip leading SQL comments (-- and /* */) and whitespace.
fn strip_sql_comments(sql: &str) -> &str {
    let mut s = sql.trim_start();
    loop {
        if s.starts_with("--") {
            s = s.find('\n').map(|i| &s[i + 1..]).unwrap_or("");
            s = s.trim_start();
        } else if s.starts_with("/*") {
            s = s.find("*/").map(|i| &s[i + 2..]).unwrap_or("");
            s = s.trim_start();
        } else {
            break;
        }
    }
    s
}

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

    // ---- Existing tests (preserved with updated expected values) ----

    #[test]
    fn test_bigquery_profile() {
        let profile = CostProfile::for_dialect(SqlDialect::BigQuery);
        assert!((profile.cost_per_tb - 6.25).abs() < 1e-10);
        assert!(profile.cheap_window_functions);
    }

    #[test]
    fn test_snowflake_profile() {
        let profile = CostProfile::for_dialect(SqlDialect::Snowflake);
        // 0.6 credits/TB * $3.00/credit = $1.80/TB
        assert!((profile.cost_per_tb - 1.80).abs() < 1e-10);
        assert!(profile.ctes_materialized);
    }

    #[test]
    fn test_generic_fallback() {
        let profile = CostProfile::for_dialect(SqlDialect::Generic);
        assert!((profile.cost_per_tb - 5.0).abs() < 1e-10);
    }

    #[test]
    fn test_calculate_cost_1tb() {
        let profile = CostProfile::for_dialect(SqlDialect::BigQuery);
        let one_tb = 1024u64 * 1024 * 1024 * 1024;
        let cost = profile.calculate_cost(one_tb);
        assert!((cost - 6.25).abs() < 0.01);
    }

    #[test]
    fn test_calculate_cost_1gb() {
        let profile = CostProfile::for_dialect(SqlDialect::BigQuery);
        let one_gb = 1024u64 * 1024 * 1024;
        let cost = profile.calculate_cost(one_gb);
        // 1 GB / 1024 = 0.000977 TB * 6.25 = ~0.0061
        assert!(cost < 0.01);
    }

    #[test]
    fn test_calculate_cost_zero() {
        let profile = CostProfile::generic();
        assert!((profile.calculate_cost(0) - 0.0).abs() < 1e-10);
    }

    // ---- New tests for PricingModel ----

    #[test]
    fn test_per_tb_scanned_model() {
        let model = PricingModel::PerTbScanned { cost_per_tb: 5.00 };
        assert!((model.effective_cost_per_tb() - 5.00).abs() < 1e-10);
    }

    #[test]
    fn test_credit_based_model() {
        let model = PricingModel::CreditBased {
            credits_per_tb: 0.6,
            default_credit_price: 3.00,
        };
        assert!((model.effective_cost_per_tb() - 1.80).abs() < 1e-10);
    }

    #[test]
    fn test_dbu_based_model() {
        let model = PricingModel::DbuBased {
            dbu_per_tb: 5.5,
            default_dbu_price: 0.55,
        };
        assert!((model.effective_cost_per_tb() - 3.025).abs() < 1e-10);
    }

    #[test]
    fn test_provisioned_model() {
        let model = PricingModel::Provisioned {
            cost_per_tb: 1.086,
            minimum_charge: 0.01,
        };
        assert!((model.effective_cost_per_tb() - 1.086).abs() < 1e-10);
    }

    #[test]
    fn test_snowflake_custom_credit_price() {
        let profile = CostProfile::snowflake_custom(4.00);
        // 0.6 credits/TB * $4.00/credit = $2.40/TB
        assert!((profile.cost_per_tb - 2.40).abs() < 1e-10);
        assert!(profile.ctes_materialized);
        assert!(profile.description.contains("$4.00"));
    }

    #[test]
    fn test_databricks_tier_classic() {
        let profile = CostProfile::databricks_tier("classic");
        // 5.5 DBU/TB * $0.22/DBU = $1.21/TB
        assert!((profile.cost_per_tb - 1.21).abs() < 1e-10);
        assert!(profile.description.contains("classic"));
    }

    #[test]
    fn test_databricks_tier_pro() {
        let profile = CostProfile::databricks_tier("pro");
        // 5.5 DBU/TB * $0.55/DBU = $3.025/TB
        assert!((profile.cost_per_tb - 3.025).abs() < 1e-10);
    }

    #[test]
    fn test_databricks_tier_serverless() {
        let profile = CostProfile::databricks_tier("serverless");
        // 5.5 DBU/TB * $0.70/DBU = $3.85/TB
        assert!((profile.cost_per_tb - 3.85).abs() < 1e-10);
    }

    #[test]
    fn test_databricks_tier_unknown_falls_back_to_pro() {
        let profile = CostProfile::databricks_tier("unknown_tier");
        let pro = CostProfile::databricks_tier("pro");
        assert!((profile.cost_per_tb - pro.cost_per_tb).abs() < 1e-10);
    }

    #[test]
    fn test_calculate_cost_with_options_snowflake_custom() {
        let profile = CostProfile::for_dialect(SqlDialect::Snowflake);
        let one_tb = 1024u64 * 1024 * 1024 * 1024;
        let options = PricingOptions {
            credit_price: Some(4.00),
            dbu_price: None,
        };
        let cost = profile.calculate_cost_with_options(one_tb, &options);
        // 0.6 credits/TB * $4.00/credit = $2.40
        assert!((cost - 2.40).abs() < 0.01);
    }

    #[test]
    fn test_calculate_cost_with_options_databricks_custom() {
        let profile = CostProfile::for_dialect(SqlDialect::Databricks);
        let one_tb = 1024u64 * 1024 * 1024 * 1024;
        let options = PricingOptions {
            credit_price: None,
            dbu_price: Some(0.70),
        };
        let cost = profile.calculate_cost_with_options(one_tb, &options);
        // 5.5 DBU/TB * $0.70/DBU = $3.85
        assert!((cost - 3.85).abs() < 0.01);
    }

    #[test]
    fn test_calculate_cost_with_options_no_overrides() {
        let profile = CostProfile::for_dialect(SqlDialect::Snowflake);
        let one_tb = 1024u64 * 1024 * 1024 * 1024;
        let default_cost = profile.calculate_cost(one_tb);
        let options_cost = profile.calculate_cost_with_options(one_tb, &PricingOptions::default());
        assert!((default_cost - options_cost).abs() < 1e-10);
    }

    #[test]
    fn test_redshift_minimum_charge() {
        let profile = CostProfile::for_dialect(SqlDialect::Redshift);
        // Very small scan: 1 KB
        let cost = profile.calculate_cost(1024);
        // Raw would be ~0.000000001, but minimum is $0.01
        assert!((cost - 0.01).abs() < 1e-10);
    }

    #[test]
    fn test_redshift_zero_bytes_no_minimum() {
        let profile = CostProfile::for_dialect(SqlDialect::Redshift);
        let cost = profile.calculate_cost(0);
        assert!((cost - 0.0).abs() < 1e-10);
    }

    #[test]
    fn test_cost_per_tb_effective() {
        let profile = CostProfile::for_dialect(SqlDialect::Snowflake);
        assert!((profile.cost_per_tb_effective() - 1.80).abs() < 1e-10);
    }

    #[test]
    fn test_disclaimer_snowflake() {
        let profile = CostProfile::for_dialect(SqlDialect::Snowflake);
        let d = profile.disclaimer();
        assert!(d.is_some());
        assert!(d.unwrap().contains("credits/TB"));
    }

    #[test]
    fn test_disclaimer_databricks() {
        let profile = CostProfile::for_dialect(SqlDialect::Databricks);
        let d = profile.disclaimer();
        assert!(d.is_some());
        assert!(d.unwrap().contains("DBU/TB"));
    }

    #[test]
    fn test_disclaimer_redshift() {
        let profile = CostProfile::for_dialect(SqlDialect::Redshift);
        let d = profile.disclaimer();
        assert!(d.is_some());
        assert!(d.unwrap().contains("minimum charge"));
    }

    #[test]
    fn test_disclaimer_bigquery_none() {
        let profile = CostProfile::for_dialect(SqlDialect::BigQuery);
        assert!(profile.disclaimer().is_none());
    }

    #[test]
    fn test_disclaimer_athena_none() {
        let profile = CostProfile::for_dialect(SqlDialect::Athena);
        assert!(profile.disclaimer().is_none());
    }

    #[test]
    fn test_athena_profile() {
        let profile = CostProfile::for_dialect(SqlDialect::Athena);
        assert!((profile.cost_per_tb - 5.00).abs() < 1e-10);
    }

    #[test]
    fn test_databricks_default_profile() {
        let profile = CostProfile::for_dialect(SqlDialect::Databricks);
        // 5.5 DBU/TB * $0.55/DBU = $3.025/TB
        assert!((profile.cost_per_tb - 3.025).abs() < 1e-10);
    }

    #[test]
    fn test_redshift_profile() {
        let profile = CostProfile::for_dialect(SqlDialect::Redshift);
        assert!((profile.cost_per_tb - 1.086).abs() < 1e-10);
    }
}
