// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Lance Authors

use std::sync::Arc;

use lance::{Error, datatypes::Schema};
use lance_table::format::{DeletionFile, Fragment};
use pyo3::{exceptions::PyIOError, prelude::*};

use crate::{Dataset, rt, utils::PyLance};

/// Format the Lance schema of a dataset as a string.
///
/// This can be used to view the field ids and types in the schema.
#[pyfunction]
pub fn format_schema(dataset: &Bound<'_, PyAny>) -> PyResult<String> {
    let py = dataset.py();
    let dataset = dataset.getattr("_ds")?.extract::<Py<Dataset>>()?;
    let dataset_ref = &dataset.bind(py).borrow().ds;
    let schema = dataset_ref.schema();
    Ok(format!("{:#?}", schema))
}

/// Print the full Lance manifest of the dataset.
#[pyfunction]
pub fn format_manifest(dataset: &Bound<'_, PyAny>) -> PyResult<String> {
    let py = dataset.py();
    let dataset = dataset.getattr("_ds")?.extract::<Py<Dataset>>()?;
    let dataset_ref = &dataset.bind(py).borrow().ds;
    let manifest = dataset_ref.manifest();
    Ok(format!("{:#?}", manifest))
}

// These are dead code because they just exist for the debug impl.
#[derive(Debug)]
#[allow(dead_code)]
struct PrettyPrintableFragment {
    id: u64,
    files: Vec<PrettyPrintableDataFile>,
    deletion_file: Option<DeletionFile>,
    physical_rows: Option<usize>,
}

#[derive(Debug)]
#[allow(dead_code)]
struct PrettyPrintableDataFile {
    path: String,
    fields: Vec<i32>,
    column_indices: Vec<i32>,
    schema: Schema,
    major_version: u32,
    minor_version: u32,
}

impl PrettyPrintableFragment {
    fn new(fragment: &Fragment, schema: &Schema) -> Self {
        let files = fragment
            .files
            .iter()
            .map(|file| {
                let schema = schema.project_by_ids(&file.fields, false);
                PrettyPrintableDataFile {
                    path: file.path.clone(),
                    fields: file.fields.to_vec(),
                    column_indices: file.column_indices.to_vec(),
                    schema,
                    major_version: file.file_major_version,
                    minor_version: file.file_minor_version,
                }
            })
            .collect();

        Self {
            id: fragment.id,
            files,
            deletion_file: fragment.deletion_file.clone(),
            physical_rows: fragment.physical_rows,
        }
    }
}

/// Debug print a LanceFragment.
#[pyfunction]
pub fn format_fragment(
    fragment: PyLance<Fragment>,
    dataset: &Bound<'_, PyAny>,
) -> PyResult<String> {
    let py = dataset.py();
    let fragment = fragment.0;

    let dataset = dataset.getattr("_ds")?.extract::<Py<Dataset>>()?;
    let dataset_ref = &dataset.bind(py).borrow().ds;
    let schema = dataset_ref.schema();

    let pp_meta = PrettyPrintableFragment::new(&fragment, schema);
    Ok(format!("{:#?}", pp_meta))
}

/// Return a string representation of each transaction in the dataset, in
/// reverse chronological order.
///
/// If `max_transactions` is provided, only the most recent `max_transactions`
/// transactions will be returned. Defaults to 10.
#[pyfunction]
#[pyo3(signature = (dataset, /, max_transactions = 10))]
pub fn list_transactions(
    dataset: &Bound<'_, PyAny>,
    max_transactions: usize,
) -> PyResult<Vec<Option<String>>> {
    let py = dataset.py();
    let dataset = dataset.getattr("_ds")?.extract::<Py<Dataset>>()?;
    let mut dataset = dataset.bind(py).borrow().ds.clone();

    rt().block_on(Some(py), async move {
        let mut transactions = vec![];

        loop {
            let transaction = dataset.read_transaction().await.map_err(|err| {
                PyIOError::new_err(format!("Failed to read transaction file: {:?}", err))
            })?;
            if let Some(transaction) = transaction {
                transactions.push(Some(format!("{:#?}", transaction)));
            } else {
                transactions.push(None);
            }

            if transactions.len() >= max_transactions {
                break;
            } else {
                match dataset
                    .checkout_version(dataset.version().version - 1)
                    .await
                {
                    Ok(ds) => dataset = Arc::new(ds),
                    Err(Error::DatasetNotFound { .. }) => break,
                    Err(err) => {
                        return Err(PyIOError::new_err(format!(
                            "Failed to checkout version: {:?}",
                            err
                        )));
                    }
                }
            }
        }

        Ok(transactions)
    })?
}
