//! Lineage bindings: column_lineage, track_lineage.

use crate::schema::Schema;
use pyo3::prelude::*;
use std::collections::HashMap;

/// Compute column-level lineage for a SQL query.
///
/// Returns a dict with `column_dict` (output->sources mapping) and
/// `column_lineage` (detailed edge list).
///
/// Pass `default_database` and `default_schema` to qualify unqualified table
/// references (important for Snowflake lineage matching).
#[pyfunction]
#[pyo3(signature = (sql, *, dialect="generic", schema=None, default_database=None, default_schema=None))]
pub fn column_lineage(
    py: Python<'_>,
    sql: &str,
    dialect: &str,
    schema: Option<&Schema>,
    default_database: Option<String>,
    default_schema: Option<String>,
) -> PyResult<PyObject> {
    crate::sdk::require_init()?;
    // Release the GIL for the HTTP call so Python threads can run.
    py.allow_threads(|| crate::sdk::check_credits("column_lineage"))?;

    let d = crate::convert::parse_dialect(dialect)?;

    // The lineage engine expects schema as { "table": { "col": "type", ... }, ... }
    let schema_map = match schema {
        Some(s) => schema_def_to_lineage_map(&s.inner),
        None => HashMap::new(),
    };

    let options = if default_database.is_some() || default_schema.is_some() {
        Some(altimate_core::lineage::ColumnLineageOptions {
            default_database,
            default_schema,
            inject_shared: false,
        })
    } else {
        None
    };

    let start = std::time::Instant::now();
    let result = altimate_core::lineage::column_lineage(sql, d, &schema_map, options.as_ref())
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e))?;
    let duration_ms = start.elapsed().as_millis() as u64;

    // Deduct credit only if lineage was actually generated (non-empty result).
    // All background calls (credit, failure, telemetry) are fire-and-forget.
    let has_lineage = !result.column_dict.is_empty() || !result.column_lineage.is_empty();
    if has_lineage {
        crate::sdk::count_credit("column_lineage", 1);
    } else if sql.trim().len() > 5 {
        let reason = if result.errors.is_empty() {
            "EmptyLineage".to_owned()
        } else {
            format!("ParseError: {}", result.errors.join("; "))
        };
        crate::sdk::report_failure(sql, dialect, duration_ms, &reason);
    }

    // Telemetry: emit event to Azure App Insights.
    let mut props = serde_json::Map::new();
    props.insert("operation".into(), serde_json::json!("column_lineage"));
    props.insert("dialect".into(), serde_json::json!(dialect));
    props.insert("success".into(), serde_json::json!(has_lineage));
    props.insert("error_count".into(), serde_json::json!(result.errors.len()));
    let mut measurements = serde_json::Map::new();
    measurements.insert("duration_ms".into(), serde_json::json!(duration_ms as f64));
    crate::sdk::emit_event("sdk/columnLineage", props, measurements);

    crate::convert::to_py_result(py, &result)
}

/// Track column-level lineage across multiple queries.
#[pyfunction]
pub fn track_lineage(py: Python<'_>, queries: Vec<String>, schema: &Schema) -> PyResult<PyObject> {
    crate::sdk::require_init()?;
    py.allow_threads(|| crate::sdk::check_credits("track_lineage"))?;

    let refs: Vec<&str> = queries.iter().map(|s| s.as_str()).collect();
    let result = altimate_core::lineage::track_lineage(&refs, &schema.inner)
        .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e))?;

    // Deduct one credit per query in the batch (fire-and-forget).
    crate::sdk::count_credit("track_lineage", queries.len());

    // Telemetry: emit event to Azure App Insights.
    let mut props = serde_json::Map::new();
    props.insert("operation".into(), serde_json::json!("track_lineage"));
    props.insert("query_count".into(), serde_json::json!(queries.len()));
    crate::sdk::emit_event("sdk/trackLineage", props, serde_json::Map::new());

    crate::convert::to_py_result(py, &result)
}

/// Convert SchemaDefinition to the flat `{ table: { col: type } }` format
/// that the polyglot-sql lineage engine expects.
fn schema_def_to_lineage_map(
    schema: &altimate_core::SchemaDefinition,
) -> HashMap<String, serde_json::Value> {
    schema
        .tables
        .iter()
        .map(|(table_name, table_def)| {
            let cols: serde_json::Map<String, serde_json::Value> = table_def
                .columns
                .iter()
                .map(|col| {
                    (
                        col.name.clone(),
                        serde_json::Value::String(col.data_type.clone()),
                    )
                })
                .collect();
            (table_name.clone(), serde_json::Value::Object(cols))
        })
        .collect()
}
