use crate::__private::S3SinkConfig;
use anyhow::Context;
use bytes::Bytes;
use futures::TryStreamExt;
use object_store::aws::AmazonS3Builder;
use object_store::path::Path as ObjectPath;
use object_store::{ObjectStore, ObjectStoreExt};
use parquet::arrow::ArrowWriter;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use parquet::basic::Compression;
use parquet::file::properties::WriterProperties;
use std::collections::BTreeMap;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::runtime::Builder as RuntimeBuilder;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use tokio::task::JoinSet;

const DEFAULT_MAX_CONCURRENT_DOWNLOADS: usize = 128;

#[derive(Debug, Clone)]
pub struct S3ParquetCompactConfig {
    pub s3: S3SinkConfig,
    pub output_path: PathBuf,
    pub max_concurrent_downloads: usize,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ParquetCompactResult {
    pub source_objects: usize,
    pub rows_written: usize,
    pub output_path: PathBuf,
}

pub fn compact_s3_parquet_to_local(
    config: S3ParquetCompactConfig,
) -> anyhow::Result<ParquetCompactResult> {
    let rt = RuntimeBuilder::new_current_thread()
        .enable_all()
        .build()
        .context("failed to build compaction runtime")?;
    rt.block_on(compact_s3_parquet_to_local_async(config))
}

pub async fn compact_s3_parquet_to_local_async(
    config: S3ParquetCompactConfig,
) -> anyhow::Result<ParquetCompactResult> {
    let store = build_s3_store(&config.s3)?;
    let max_concurrent_downloads = match config.max_concurrent_downloads {
        0 => DEFAULT_MAX_CONCURRENT_DOWNLOADS,
        value => value,
    };
    compact_store_parquet_prefix_to_local(
        store,
        config.s3.prefix.as_str(),
        config.output_path.as_path(),
        max_concurrent_downloads,
    )
    .await
}

async fn compact_store_parquet_prefix_to_local(
    store: Arc<dyn ObjectStore>,
    prefix: &str,
    output_path: &Path,
    max_concurrent_downloads: usize,
) -> anyhow::Result<ParquetCompactResult> {
    let prefix = prefix.trim_matches('/');
    let prefix_path = if prefix.is_empty() {
        None
    } else {
        Some(ObjectPath::from(prefix))
    };

    let mut objects = store
        .list(prefix_path.as_ref())
        .try_collect::<Vec<_>>()
        .await
        .context("failed to list parquet objects")?;
    objects.sort_by(|a, b| a.location.cmp(&b.location));
    objects.retain(|meta| meta.location.as_ref().ends_with(".parquet"));

    if objects.is_empty() {
        anyhow::bail!("no parquet objects found under prefix '{prefix}'");
    }

    if let Some(parent) = output_path.parent()
        && !parent.as_os_str().is_empty()
    {
        std::fs::create_dir_all(parent)
            .with_context(|| format!("failed to create output directory {}", parent.display()))?;
    }

    let semaphore = Arc::new(Semaphore::new(max_concurrent_downloads));
    let mut join_set = JoinSet::new();
    let mut ready = BTreeMap::new();
    let mut next_spawn_index = 0usize;
    let mut next_write_index = 0usize;
    let mut in_flight = 0usize;

    let mut writer = None;
    let mut expected_schema = None;
    let mut rows_written = 0usize;

    while next_write_index < objects.len() {
        while next_spawn_index < objects.len() && in_flight < max_concurrent_downloads {
            let permit = semaphore
                .clone()
                .acquire_owned()
                .await
                .context("compaction download semaphore closed unexpectedly")?;
            let object = objects[next_spawn_index].clone();
            let store = Arc::clone(&store);
            join_set.spawn(async move {
                fetch_object_payload(store, next_spawn_index, object.location, permit).await
            });
            next_spawn_index += 1;
            in_flight += 1;
        }

        let Some(result) = join_set.join_next().await else {
            anyhow::bail!("download queue terminated before all parquet objects were processed");
        };
        in_flight = in_flight.saturating_sub(1);

        let fetched = match result {
            Ok(Ok(fetched)) => fetched,
            Ok(Err(err)) => {
                join_set.abort_all();
                return Err(err);
            }
            Err(err) => {
                join_set.abort_all();
                return Err(anyhow::anyhow!("compaction download task failed: {err}"));
            }
        };
        ready.insert(fetched.index, fetched);

        while let Some(fetched) = ready.remove(&next_write_index) {
            append_parquet_payload(
                fetched.location.as_ref(),
                fetched.payload,
                output_path,
                &mut writer,
                &mut expected_schema,
                &mut rows_written,
            )?;
            next_write_index += 1;
        }
    }

    writer
        .context("writer should be initialized when objects are present")?
        .close()
        .with_context(|| format!("failed to finalize {}", output_path.display()))?;

    Ok(ParquetCompactResult {
        source_objects: objects.len(),
        rows_written,
        output_path: output_path.to_path_buf(),
    })
}

struct FetchedObject {
    index: usize,
    location: ObjectPath,
    payload: Bytes,
}

async fn fetch_object_payload(
    store: Arc<dyn ObjectStore>,
    index: usize,
    location: ObjectPath,
    _permit: OwnedSemaphorePermit,
) -> anyhow::Result<FetchedObject> {
    let payload = store
        .get(&location)
        .await
        .with_context(|| format!("failed to fetch {location}"))?
        .bytes()
        .await
        .with_context(|| format!("failed to read {location}"))?;

    Ok(FetchedObject {
        index,
        location,
        payload,
    })
}

fn append_parquet_payload(
    location: &str,
    payload: Bytes,
    output_path: &Path,
    writer: &mut Option<ArrowWriter<File>>,
    expected_schema: &mut Option<arrow::datatypes::SchemaRef>,
    rows_written: &mut usize,
) -> anyhow::Result<()> {
    let builder = ParquetRecordBatchReaderBuilder::try_new(payload)
        .with_context(|| format!("failed to open parquet {location}"))?;
    let schema = builder.schema().clone();

    if let Some(existing) = expected_schema.as_ref() {
        if existing.as_ref() != schema.as_ref() {
            anyhow::bail!(
                "schema mismatch for {} while compacting prefix '{}'",
                location,
                output_path.display()
            );
        }
    } else {
        let file = File::create(output_path)
            .with_context(|| format!("failed to create {}", output_path.display()))?;
        let props = WriterProperties::builder()
            .set_compression(Compression::SNAPPY)
            .build();
        *writer = Some(
            ArrowWriter::try_new(file, schema.clone(), Some(props))
                .with_context(|| format!("failed to open {}", output_path.display()))?,
        );
        *expected_schema = Some(schema.clone());
    }

    let record_reader = builder
        .build()
        .with_context(|| format!("failed to build batch reader for {location}"))?;
    for batch in record_reader {
        let batch = batch.with_context(|| format!("failed to read batch from {location}"))?;
        *rows_written += batch.num_rows();
        writer
            .as_mut()
            .expect("writer should be initialized from first schema")
            .write(&batch)
            .with_context(|| format!("failed to write batch from {location}"))?;
    }
    Ok(())
}

fn build_s3_store(config: &S3SinkConfig) -> anyhow::Result<Arc<dyn ObjectStore>> {
    let mut builder = AmazonS3Builder::new()
        .with_bucket_name(config.bucket.clone())
        .with_region(config.region.clone());
    if let Some(endpoint) = &config.endpoint {
        builder = builder.with_endpoint(endpoint.clone());
    }
    if let Some(access_key_id) = &config.access_key_id {
        builder = builder.with_access_key_id(resolve_credential(access_key_id));
    }
    if let Some(secret_access_key) = &config.secret_access_key {
        builder = builder.with_secret_access_key(resolve_credential(secret_access_key));
    }
    let store = builder
        .build()
        .map_err(|e| anyhow::anyhow!("failed to build s3 store: {e}"))?;
    Ok(Arc::new(store))
}

fn read_credential_from_path(path: &str) -> anyhow::Result<String> {
    let content = std::fs::read_to_string(path)
        .map_err(|e| anyhow::anyhow!("failed to read credential from {path}: {e}"))?;
    Ok(content.trim().to_string())
}

fn resolve_credential(value: &str) -> String {
    if value.starts_with('/') && Path::new(value).exists() {
        match read_credential_from_path(value) {
            Ok(credential) => credential,
            Err(_) => value.to_string(),
        }
    } else {
        value.to_string()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow::array::{Array, Int64Array, StringArray};
    use arrow::datatypes::{DataType, Field, Schema};
    use arrow::record_batch::RecordBatch;
    use bytes::Bytes;
    use object_store::memory::InMemory;
    use object_store::{PutOptions, PutPayload};
    use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
    use std::sync::Arc;
    use std::time::{SystemTime, UNIX_EPOCH};

    fn make_test_parquet(rows: &[(&str, i64)]) -> Vec<u8> {
        let schema = Arc::new(Schema::new(vec![
            Field::new("name", DataType::Utf8, false),
            Field::new("value", DataType::Int64, false),
        ]));
        let names = rows
            .iter()
            .map(|(name, _)| (*name).to_string())
            .collect::<Vec<_>>();
        let values = rows.iter().map(|(_, value)| *value).collect::<Vec<_>>();
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![
                Arc::new(StringArray::from(names)),
                Arc::new(Int64Array::from(values)),
            ],
        )
        .unwrap();

        let mut buf = Vec::new();
        let props = WriterProperties::builder()
            .set_compression(Compression::SNAPPY)
            .build();
        let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap();
        writer.write(&batch).unwrap();
        writer.close().unwrap();
        buf
    }

    fn temp_output_path(name: &str) -> PathBuf {
        let suffix = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        std::env::temp_dir().join(format!("pm-compact-{name}-{suffix}.parquet"))
    }

    #[tokio::test]
    async fn compact_store_parquet_prefix_to_local_merges_multiple_files() {
        let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
        let payload_a = make_test_parquet(&[("a", 1), ("b", 2)]);
        let payload_b = make_test_parquet(&[("c", 3)]);

        store
            .put_opts(
                &ObjectPath::from("prefix/part-0001.parquet"),
                PutPayload::from(Bytes::from(payload_a)),
                PutOptions::default(),
            )
            .await
            .unwrap();
        store
            .put_opts(
                &ObjectPath::from("prefix/part-0002.parquet"),
                PutPayload::from(Bytes::from(payload_b)),
                PutOptions::default(),
            )
            .await
            .unwrap();

        let output_path = temp_output_path("merge");
        let result =
            compact_store_parquet_prefix_to_local(store, "prefix", output_path.as_path(), 2)
                .await
                .unwrap();

        assert_eq!(result.source_objects, 2);
        assert_eq!(result.rows_written, 3);
        assert_eq!(result.output_path, output_path);

        let output_bytes = std::fs::read(output_path.as_path()).unwrap();
        let record_reader = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(output_bytes))
            .unwrap()
            .build()
            .unwrap();
        let mut names = Vec::new();
        let mut values = Vec::new();
        for batch in record_reader {
            let batch = batch.unwrap();
            let name_col = batch
                .column(0)
                .as_any()
                .downcast_ref::<StringArray>()
                .unwrap();
            let value_col = batch
                .column(1)
                .as_any()
                .downcast_ref::<Int64Array>()
                .unwrap();
            for i in 0..batch.num_rows() {
                names.push(name_col.value(i).to_string());
                values.push(value_col.value(i));
            }
        }
        assert_eq!(names, vec!["a", "b", "c"]);
        assert_eq!(values, vec![1, 2, 3]);

        std::fs::remove_file(output_path).unwrap();
    }

    #[tokio::test]
    async fn compact_store_parquet_prefix_to_local_ignores_non_parquet_objects() {
        let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
        let payload = make_test_parquet(&[("a", 1)]);

        store
            .put_opts(
                &ObjectPath::from("prefix/part-0001.parquet"),
                PutPayload::from(Bytes::from(payload)),
                PutOptions::default(),
            )
            .await
            .unwrap();
        store
            .put_opts(
                &ObjectPath::from("prefix/notes.txt"),
                PutPayload::from(Bytes::from_static(b"ignored")),
                PutOptions::default(),
            )
            .await
            .unwrap();

        let output_path = temp_output_path("filter");
        let result =
            compact_store_parquet_prefix_to_local(store, "prefix", output_path.as_path(), 2)
                .await
                .unwrap();
        assert_eq!(result.source_objects, 1);
        assert_eq!(result.rows_written, 1);

        std::fs::remove_file(output_path).unwrap();
    }
}
