use std::sync::Arc;

use arrow::{
    array::{
        Array, Date32Array, Int32Array, Int64Array, RecordBatch, StringArray,
        TimestampMicrosecondArray,
    },
    datatypes::{DataType, Field, Schema, TimeUnit},
};
use futures::TryStreamExt;
use mysql_async::{Conn, Pool, Row, params::Params, prelude::*};
use parquet::arrow::AsyncArrowWriter;
use rayon::prelude::*;
use tokio::fs::File;

#[cfg(test)]
mod test;

type Error = Box<dyn std::error::Error + Send + Sync>;

pub async fn load_data(pool: &Pool) -> Result<(), Error> {
    let data = String::from_utf8_lossy(include_bytes!("../data.sql"));
    // Obviously do something better if you don't control the input data.
    for stmt in data.trim().split(";") {
        if !stmt.trim().is_empty() {
            let mut conn = pool.get_conn().await?;
            conn.exec_drop(stmt, Params::Empty).await?;
        }
    }
    Ok(())
}

pub async fn convert_data(pool: &Pool, table: &str) -> Result<(), Error> {
    let mut conn = pool.get_conn().await?;
    let schema = discover_schema(&mut conn, table).await?;
    let file = File::create(format!("{table}.parquet")).await?;
    let mut writer = AsyncArrowWriter::try_new(file, schema.clone(), Default::default())?;
    const BATCH_SIZE: usize = 50000;
    let mut rows = Vec::with_capacity(BATCH_SIZE);
    let mut stream = conn
        .exec_stream(format!("select * from {table}"), Params::Empty)
        .await?;
    while let Some(row) = stream.try_next().await? {
        rows.push(row);
        if rows.len() >= BATCH_SIZE {
            process_parquet(&mut rows, schema.clone(), &mut writer).await?;
        }
    }
    if !rows.is_empty() {
        process_parquet(&mut rows, schema.clone(), &mut writer).await?;
    }
    writer.close().await?;
    Ok(())
}

async fn discover_schema(conn: &mut Conn, table: &str) -> Result<Arc<Schema>, Error> {
    let query = format!("describe {}", table);
    let rows: Vec<Row> = conn.exec(query, ()).await?;
    let mut fields = Vec::new();
    for row in rows {
        let name: String = row.get("Field").unwrap();
        let mysql_type: String = row.get("Type").unwrap();
        let is_nullable: String = row.get("Null").unwrap();
        let arrow_type = match mysql_type.as_str() {
            t if t.starts_with("int") => DataType::Int32,
            t if t.starts_with("bigint") => DataType::Int64,
            t if t.starts_with("varchar") => DataType::Utf8,
            t if t.starts_with("datetime") => DataType::Timestamp(TimeUnit::Microsecond, None),
            t if t.starts_with("date") => DataType::Date32,
            _ => DataType::Utf8,
        };
        let nullable = is_nullable == "YES";
        fields.push(Field::new(name, arrow_type, nullable));
    }
    Ok(Arc::new(Schema::new(fields)))
}

async fn process_parquet(
    rows: &mut Vec<Row>,
    schema: Arc<Schema>,
    writer: &mut AsyncArrowWriter<File>,
) -> Result<(), Error> {
    if rows.is_empty() {
        return Ok(());
    }
    let epoch = mysql_common::chrono::NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
    let mut data: Vec<Box<dyn Array>> = Vec::new();
    for field in schema.fields() {
        let name = field.name().as_str();
        match field.data_type() {
            DataType::Int32 => {
                let values: Vec<Option<i32>> =
                    rows.par_iter().map(|row| row.get::<i32, _>(name)).collect();
                data.push(Box::new(Int32Array::from(values)));
            }
            DataType::Int64 => {
                let values: Vec<Option<i64>> =
                    rows.par_iter().map(|row| row.get::<i64, _>(name)).collect();
                data.push(Box::new(Int64Array::from(values)));
            }
            DataType::Utf8 => {
                let values: Vec<Option<String>> = rows
                    .par_iter()
                    .map(|row| row.get::<String, _>(name))
                    .collect();
                data.push(Box::new(StringArray::from(values)));
            }
            DataType::Timestamp(TimeUnit::Microsecond, None) => {
                let values: Vec<Option<i64>> = rows
                    .par_iter()
                    .map(|row| {
                        row.get::<mysql_common::chrono::NaiveDateTime, _>(name)
                            .map(|dt| dt.and_utc().timestamp_micros())
                    })
                    .collect();
                data.push(Box::new(TimestampMicrosecondArray::from(values)));
            }
            DataType::Date32 => {
                let values: Vec<Option<i32>> = rows
                    .par_iter()
                    .map(|row| {
                        let date = row.get::<mysql_common::chrono::NaiveDate, _>(name);
                        Some((date.unwrap() - epoch).num_days() as i32)
                    })
                    .collect();
                data.push(Box::new(Date32Array::from(values)));
            }
            _ => {
                let values: Vec<Option<String>> = rows
                    .par_iter()
                    .map(|row| row.get::<String, _>(name))
                    .collect();
                data.push(Box::new(StringArray::from(values)));
            }
        }
    }
    let columns: Vec<Arc<dyn Array>> = data.into_iter().map(|arr| arr.into()).collect();
    let batch = RecordBatch::try_new(schema.clone(), columns)?;
    writer.write(&batch).await?;
    rows.clear();
    Ok(())
}
