//! GPT SQL UDFs — 通过 SQL 接口操作 GPT 模型 (Phase-69: T69.3-T69.5)
//!
//! 提供 4 个 SQL 标量函数：
//!
//! - `gpt_create(name[, config_json])` — 创建新模型
//! - `gpt_train(name, text, steps[, lr, batch_size])` — 训练模型
//! - `gpt_generate(name, prompt[, max_tokens, temperature])` — 生成文本
//! - `gpt_status([name])` — 查询模型状态
//!
//! ## SQL 示例
//!
//! ```sql
//! -- 创建一个微型 GPT 模型
//! SELECT gpt_create('my_gpt');
//!
//! -- 用自定义配置创建
//! SELECT gpt_create('big_gpt', '{"n_embd":128,"n_head":4,"n_layer":4,"block_size":256}');
//!
//! -- 训练（text, steps, learning_rate）
//! SELECT gpt_train('my_gpt', 'Hello World! This is training data...', 100, 0.001);
//!
//! -- 生成文本
//! SELECT gpt_generate('my_gpt', 'Hello', 50, 0.8);
//!
//! -- 查看所有模型
//! SELECT gpt_status();
//!
//! -- 查看特定模型
//! SELECT gpt_status('my_gpt');
//! ```

use std::sync::Arc;

use arrow_array::{Array, StringArray};
use arrow_schema::DataType;
use datafusion::error::DataFusionError;
use datafusion::execution::context::SessionContext;
use datafusion::logical_expr::{
    ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
};

use crate::gpt::{GenerateConfig, GptConfig, TrainConfig};
use crate::gpt::model_registry::global_model_registry;

/// 注册所有 GPT 相关的 SQL UDF
pub fn register_gpt_udfs(ctx: &SessionContext) {
    ctx.register_udf(ScalarUDF::from(GptCreateUdf::new()));
    ctx.register_udf(ScalarUDF::from(GptTrainUdf::new()));
    ctx.register_udf(ScalarUDF::from(GptGenerateUdf::new()));
    ctx.register_udf(ScalarUDF::from(GptStatusUdf::new()));
}

// ============================================================================
// Helper functions
// ============================================================================

/// 从 ColumnarValue 中提取字符串
fn extract_string(val: &ColumnarValue) -> Result<String, DataFusionError> {
    match val {
        ColumnarValue::Scalar(s) => {
            let s = s.to_string();
            // DataFusion scalar to_string may include quotes
            Ok(s.trim_matches('\'').trim_matches('"').to_string())
        }
        ColumnarValue::Array(arr) => {
            let str_arr = arr.as_any().downcast_ref::<StringArray>().ok_or_else(|| {
                DataFusionError::Execution("Expected Utf8 argument".to_string())
            })?;
            if str_arr.is_empty() {
                return Err(DataFusionError::Execution(
                    "Empty string array".to_string(),
                ));
            }
            Ok(str_arr.value(0).to_string())
        }
    }
}

/// 从 ColumnarValue 中提取 f64
fn extract_f64(val: &ColumnarValue) -> Result<f64, DataFusionError> {
    match val {
        ColumnarValue::Scalar(s) => {
            let s = s.to_string();
            s.trim()
                .parse::<f64>()
                .map_err(|e| DataFusionError::Execution(format!("Expected number: {}", e)))
        }
        _ => Err(DataFusionError::Execution(
            "Expected scalar number".to_string(),
        )),
    }
}

/// 从 ColumnarValue 中提取 usize
fn extract_usize(val: &ColumnarValue) -> Result<usize, DataFusionError> {
    extract_f64(val).map(|v| v as usize)
}

// ============================================================================
// T69.3 — gpt_create(name[, config_json])
// ============================================================================

/// SQL UDF: `gpt_create(name[, config_json])`
///
/// 创建一个新的 GPT 模型并注册到 ModelRegistry。
///
/// # Arguments (in SQL)
/// - `name`: Utf8 — 模型名称（唯一标识符）
/// - `config_json`: Utf8 (optional) — GptConfig 的 JSON 字符串
///
/// # Returns
/// Utf8 — 创建结果描述
///
/// # Example
/// ```sql
/// SELECT gpt_create('my_model');
/// SELECT gpt_create('my_model', '{"n_embd":128,"n_head":4,"n_layer":4}');
/// ```
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct GptCreateUdf {
    signature: Signature,
}

impl GptCreateUdf {
    pub fn new() -> Self {
        Self {
            signature: Signature::variadic_any(Volatility::Volatile),
        }
    }
}

impl ScalarUDFImpl for GptCreateUdf {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }
    fn name(&self) -> &str {
        "gpt_create"
    }
    fn signature(&self) -> &Signature {
        &self.signature
    }
    fn return_type(&self, _: &[DataType]) -> Result<DataType, DataFusionError> {
        Ok(DataType::Utf8)
    }

    fn invoke_with_args(
        &self,
        args: ScalarFunctionArgs,
    ) -> Result<ColumnarValue, DataFusionError> {
        let args = &args.args;
        if args.is_empty() {
            return Err(DataFusionError::Execution(
                "gpt_create(name[, config_json]): requires at least 1 argument".to_string(),
            ));
        }

        let name = extract_string(&args[0])?;

        // 可选：自定义配置 JSON
        let config = if args.len() >= 2 {
            let json_str = extract_string(&args[1])?;
            serde_json::from_str::<GptConfig>(&json_str).map_err(|e| {
                DataFusionError::Execution(format!("Invalid GptConfig JSON: {}", e))
            })?
        } else {
            GptConfig::default()
        };

        let registry = global_model_registry();
        let meta = registry
            .create_model(&name, config)
            .map_err(|e| DataFusionError::Execution(format!("{}", e)))?;

        let result = format!(
            "Model '{}' created (~{}K params, {}L/{}H/{}E, block_size={})",
            meta.name,
            meta.config.estimate_params() / 1000,
            meta.config.n_layer,
            meta.config.n_head,
            meta.config.n_embd,
            meta.config.block_size,
        );

        Ok(ColumnarValue::Scalar(
            datafusion::common::ScalarValue::Utf8(Some(result)),
        ))
    }
}

// ============================================================================
// T69.4 — gpt_train(name, text, steps[, lr, batch_size])
// ============================================================================

/// SQL UDF: `gpt_train(name, text, steps[, lr, batch_size])`
///
/// 训练指定模型。文本数据通过字符级 tokenizer 分词后进行训练。
///
/// # Arguments (in SQL)
/// - `name`: Utf8 — 模型名称
/// - `text`: Utf8 — 训练文本数据
/// - `steps`: Int/Float — 训练步数
/// - `lr`: Float (optional) — 学习率，默认 3e-4
/// - `batch_size`: Int (optional) — 批大小，默认 4
///
/// # Returns
/// Utf8 — JSON 格式的训练报告
///
/// # Example
/// ```sql
/// SELECT gpt_train('my_model', 'Training text data...', 100);
/// SELECT gpt_train('my_model', 'Training text data...', 100, 0.001, 8);
/// ```
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct GptTrainUdf {
    signature: Signature,
}

impl GptTrainUdf {
    pub fn new() -> Self {
        Self {
            signature: Signature::variadic_any(Volatility::Volatile),
        }
    }
}

impl ScalarUDFImpl for GptTrainUdf {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }
    fn name(&self) -> &str {
        "gpt_train"
    }
    fn signature(&self) -> &Signature {
        &self.signature
    }
    fn return_type(&self, _: &[DataType]) -> Result<DataType, DataFusionError> {
        Ok(DataType::Utf8)
    }

    fn invoke_with_args(
        &self,
        args: ScalarFunctionArgs,
    ) -> Result<ColumnarValue, DataFusionError> {
        let args = &args.args;
        if args.len() < 3 {
            return Err(DataFusionError::Execution(
                "gpt_train(name, text, steps[, lr, batch_size]): requires at least 3 arguments"
                    .to_string(),
            ));
        }

        let name = extract_string(&args[0])?;
        let text = extract_string(&args[1])?;
        let steps = extract_usize(&args[2])?;
        let lr = if args.len() >= 4 {
            extract_f64(&args[3])?
        } else {
            3e-4
        };
        let batch_size = if args.len() >= 5 {
            extract_usize(&args[4])?
        } else {
            4
        };

        let train_config = TrainConfig {
            learning_rate: lr,
            ..Default::default()
        };

        let registry = global_model_registry();
        let report = registry
            .train_model(&name, &text, steps, train_config, batch_size, None)
            .map_err(|e| DataFusionError::Execution(format!("{}", e)))?;

        let result = serde_json::to_string(&report).unwrap_or_else(|_| {
            format!(
                "{{\"initial_loss\":{:.4},\"final_loss\":{:.4},\"steps\":{},\"tok_per_sec\":{:.0}}}",
                report.initial_loss, report.final_loss, report.total_steps, report.avg_tokens_per_sec
            )
        });

        Ok(ColumnarValue::Scalar(
            datafusion::common::ScalarValue::Utf8(Some(result)),
        ))
    }
}

// ============================================================================
// T69.5 — gpt_generate(name, prompt[, max_tokens, temperature])
// ============================================================================

/// SQL UDF: `gpt_generate(name, prompt[, max_tokens, temperature])`
///
/// 使用训练好的模型生成文本。
///
/// # Arguments (in SQL)
/// - `name`: Utf8 — 模型名称
/// - `prompt`: Utf8 — 输入 prompt
/// - `max_tokens`: Int (optional) — 最大生成 token 数，默认 100
/// - `temperature`: Float (optional) — 温度参数，默认 0.8
///
/// # Returns
/// Utf8 — 生成的文本
///
/// # Example
/// ```sql
/// SELECT gpt_generate('my_model', 'Hello World');
/// SELECT gpt_generate('my_model', 'Once upon a', 200, 0.5);
/// ```
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct GptGenerateUdf {
    signature: Signature,
}

impl GptGenerateUdf {
    pub fn new() -> Self {
        Self {
            signature: Signature::variadic_any(Volatility::Volatile),
        }
    }
}

impl ScalarUDFImpl for GptGenerateUdf {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }
    fn name(&self) -> &str {
        "gpt_generate"
    }
    fn signature(&self) -> &Signature {
        &self.signature
    }
    fn return_type(&self, _: &[DataType]) -> Result<DataType, DataFusionError> {
        Ok(DataType::Utf8)
    }

    fn invoke_with_args(
        &self,
        args: ScalarFunctionArgs,
    ) -> Result<ColumnarValue, DataFusionError> {
        let args = &args.args;
        if args.len() < 2 {
            return Err(DataFusionError::Execution(
                "gpt_generate(name, prompt[, max_tokens, temperature]): requires at least 2 arguments".to_string(),
            ));
        }

        let name = extract_string(&args[0])?;
        let prompt = extract_string(&args[1])?;
        let max_tokens = if args.len() >= 3 {
            extract_usize(&args[2])?
        } else {
            100
        };
        let temperature = if args.len() >= 4 {
            extract_f64(&args[3])?
        } else {
            0.8
        };

        let gen_config = GenerateConfig {
            max_new_tokens: max_tokens,
            temperature,
            top_k: 40,
        };

        let registry = global_model_registry();
        let text = registry
            .generate(&name, &prompt, &gen_config)
            .map_err(|e| DataFusionError::Execution(format!("{}", e)))?;

        Ok(ColumnarValue::Scalar(
            datafusion::common::ScalarValue::Utf8(Some(text)),
        ))
    }
}

// ============================================================================
// T69.5 — gpt_status([name])
// ============================================================================

/// SQL UDF: `gpt_status([name])`
///
/// 查询模型状态。无参数时列出所有模型，有参数时显示指定模型详情。
///
/// # Arguments (in SQL)
/// - `name`: Utf8 (optional) — 模型名称
///
/// # Returns
/// Utf8 — JSON 格式的模型状态
///
/// # Example
/// ```sql
/// SELECT gpt_status();           -- 列出所有模型
/// SELECT gpt_status('my_model'); -- 查看特定模型
/// ```
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct GptStatusUdf {
    signature: Signature,
}

impl GptStatusUdf {
    pub fn new() -> Self {
        Self {
            signature: Signature::variadic_any(Volatility::Volatile),
        }
    }
}

impl ScalarUDFImpl for GptStatusUdf {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }
    fn name(&self) -> &str {
        "gpt_status"
    }
    fn signature(&self) -> &Signature {
        &self.signature
    }
    fn return_type(&self, _: &[DataType]) -> Result<DataType, DataFusionError> {
        Ok(DataType::Utf8)
    }

    fn invoke_with_args(
        &self,
        args: ScalarFunctionArgs,
    ) -> Result<ColumnarValue, DataFusionError> {
        let args = &args.args;
        let registry = global_model_registry();

        if args.is_empty() {
            // 列出所有模型
            let models = registry.list_models();
            let summary: Vec<serde_json::Value> = models
                .iter()
                .map(|m| {
                    serde_json::json!({
                        "name": m.name,
                        "status": m.status.to_string(),
                        "params": format!("~{}K", m.config.estimate_params() / 1000),
                        "checkpoints": m.checkpoints.len(),
                        "best_loss": m.best_loss,
                    })
                })
                .collect();
            let result =
                serde_json::to_string_pretty(&summary).unwrap_or_else(|_| "[]".to_string());
            return Ok(ColumnarValue::Scalar(
                datafusion::common::ScalarValue::Utf8(Some(result)),
            ));
        }

        // 查看特定模型
        let name = extract_string(&args[0])?;
        let meta = registry
            .get_model(&name)
            .map_err(|e| DataFusionError::Execution(format!("{}", e)))?;

        let result = serde_json::to_string_pretty(&meta).unwrap_or_else(|_| {
            format!(
                "{{\"name\":\"{}\",\"status\":\"{}\"}}",
                meta.name, meta.status
            )
        });

        Ok(ColumnarValue::Scalar(
            datafusion::common::ScalarValue::Utf8(Some(result)),
        ))
    }
}
