use std::sync::{Arc, Mutex};

use arrow::{array::RecordBatch, datatypes::Schema};
use autark_dataframe::DataFrame;
use tokio::runtime::Builder;
use transcend_sdk::{client::Client, datasets::DatasetIpcReadOptions};

use crate::readers::{OnceReader, dataframe_from_stream, log_stream_batch};
use crate::{Error, Result};

pub struct TranscendReader {
    base_url: String,
    api_key: String,
    dataset: String,
    max_batches: Option<usize>,
    latest: Option<bool>,
    batches: Mutex<Option<Vec<RecordBatch>>>,
}

impl TranscendReader {
    pub fn new(base_url: &str, api_key: &str, dataset: &str) -> Result<Self> {
        if dataset.trim().is_empty() {
            return Err(Error::EmptyReader);
        }

        Ok(Self {
            base_url: base_url.to_string(),
            api_key: api_key.to_string(),
            dataset: dataset.to_string(),
            max_batches: None,
            latest: None,
            batches: Mutex::new(None),
        })
    }

    pub fn with_max_batches(mut self, max_batches: usize) -> Self {
        self.max_batches = Some(max_batches);
        self
    }

    pub fn with_latest(mut self, latest: bool) -> Self {
        self.latest = Some(latest);
        self
    }

    fn batches_cached(&self) -> Result<Vec<RecordBatch>> {
        let mut guard = self.batches.lock().map_err(|_| Error::PoisonedLock)?;
        if let Some(batches) = guard.as_ref() {
            return Ok(batches.clone());
        }

        let runtime = Builder::new_current_thread().enable_all().build()?;
        let batches = runtime.block_on(async {
            let client = Client::new(&self.base_url, &self.api_key)
                .map_err(|err| Error::Reader(format!("transcend error: {err}")))?;
            client
                .dataset_ipc(
                    &self.dataset,
                    DatasetIpcReadOptions::new(self.max_batches, self.latest),
                )
                .await
                .map_err(|err| Error::Reader(format!("transcend error: {err}")))
        })?;

        if batches.is_empty() {
            return Err(Error::EmptyReader);
        }

        *guard = Some(batches.clone());
        Ok(batches)
    }
}

impl OnceReader for TranscendReader {
    fn read(&mut self) -> Result<DataFrame> {
        dataframe_from_stream(self)
    }

    fn schema(&self) -> Result<Arc<Schema>> {
        self.batches_cached()?
            .first()
            .map(|batch| batch.schema())
            .ok_or(Error::EmptyReader)
    }

    fn read_stream(&mut self, sink: &mut dyn FnMut(RecordBatch) -> Result<()>) -> Result<()> {
        self.batches_cached()?
            .into_iter()
            .enumerate()
            .try_for_each(|(batch_ix, batch)| {
                log_stream_batch("transcend", batch_ix, &batch);
                sink(batch)
            })
    }
}
