use crate::base::value::KeyValue;
use crate::prelude::*;

use bytes::Bytes;
use numpy::{PyArray1, PyArrayDyn, PyArrayMethods};
use pyo3::IntoPyObjectExt;
use pyo3::exceptions::PyTypeError;
use pyo3::types::PyAny;
use pyo3::types::{PyList, PyTuple};
use pyo3::{exceptions::PyException, prelude::*};
use pythonize::{depythonize, pythonize};

fn basic_value_to_py_object<'py>(
    py: Python<'py>,
    v: &value::BasicValue,
) -> PyResult<Bound<'py, PyAny>> {
    let result = match v {
        value::BasicValue::Bytes(v) => v.into_bound_py_any(py)?,
        value::BasicValue::Str(v) => v.into_bound_py_any(py)?,
        value::BasicValue::Bool(v) => v.into_bound_py_any(py)?,
        value::BasicValue::Int64(v) => v.into_bound_py_any(py)?,
        value::BasicValue::Float32(v) => v.into_bound_py_any(py)?,
        value::BasicValue::Float64(v) => v.into_bound_py_any(py)?,
        value::BasicValue::Range(v) => pythonize(py, v)?,
        value::BasicValue::Uuid(uuid_val) => uuid_val.into_bound_py_any(py)?,
        value::BasicValue::Date(v) => v.into_bound_py_any(py)?,
        value::BasicValue::Time(v) => v.into_bound_py_any(py)?,
        value::BasicValue::LocalDateTime(v) => v.into_bound_py_any(py)?,
        value::BasicValue::OffsetDateTime(v) => v.into_bound_py_any(py)?,
        value::BasicValue::TimeDelta(v) => v.into_bound_py_any(py)?,
        value::BasicValue::Json(v) => pythonize(py, v)?,
        value::BasicValue::Vector(v) => handle_vector_to_py(py, v)?,
        value::BasicValue::UnionVariant { tag_id, value } => {
            (*tag_id, basic_value_to_py_object(py, value)?).into_bound_py_any(py)?
        }
    };
    Ok(result)
}

pub fn field_values_to_py_object<'py, 'a>(
    py: Python<'py>,
    values: impl Iterator<Item = &'a value::Value>,
) -> PyResult<Bound<'py, PyAny>> {
    let fields = values
        .map(|v| value_to_py_object(py, v))
        .collect::<PyResult<Vec<_>>>()?;
    Ok(PyTuple::new(py, fields)?.into_any())
}

pub fn key_to_py_object<'py, 'a>(
    py: Python<'py>,
    key: impl IntoIterator<Item = &'a value::KeyPart>,
) -> PyResult<Bound<'py, PyAny>> {
    fn key_part_to_py_object<'py>(
        py: Python<'py>,
        part: &value::KeyPart,
    ) -> PyResult<Bound<'py, PyAny>> {
        let result = match part {
            value::KeyPart::Bytes(v) => v.into_bound_py_any(py)?,
            value::KeyPart::Str(v) => v.into_bound_py_any(py)?,
            value::KeyPart::Bool(v) => v.into_bound_py_any(py)?,
            value::KeyPart::Int64(v) => v.into_bound_py_any(py)?,
            value::KeyPart::Range(v) => pythonize(py, v)?,
            value::KeyPart::Uuid(v) => v.into_bound_py_any(py)?,
            value::KeyPart::Date(v) => v.into_bound_py_any(py)?,
            value::KeyPart::Struct(v) => key_to_py_object(py, v)?,
        };
        Ok(result)
    }
    let fields = key
        .into_iter()
        .map(|part| key_part_to_py_object(py, part))
        .collect::<PyResult<Vec<_>>>()?;
    Ok(PyTuple::new(py, fields)?.into_any())
}

pub fn value_to_py_object<'py>(py: Python<'py>, v: &value::Value) -> PyResult<Bound<'py, PyAny>> {
    let result = match v {
        value::Value::Null => py.None().into_bound(py),
        value::Value::Basic(v) => basic_value_to_py_object(py, v)?,
        value::Value::Struct(v) => field_values_to_py_object(py, v.fields.iter())?,
        value::Value::UTable(v) | value::Value::LTable(v) => {
            let rows = v
                .iter()
                .map(|v| field_values_to_py_object(py, v.0.fields.iter()))
                .collect::<PyResult<Vec<_>>>()?;
            PyList::new(py, rows)?.into_any()
        }
        value::Value::KTable(v) => {
            let rows = v
                .iter()
                .map(|(k, v)| {
                    let k: Box<[value::Value]> =
                        k.into_iter().map(|k| value::Value::from(k)).collect();
                    field_values_to_py_object(py, k.iter().chain(v.0.fields.iter()))
                })
                .collect::<PyResult<Vec<_>>>()?;
            PyList::new(py, rows)?.into_any()
        }
    };
    Ok(result)
}

fn basic_value_from_py_object<'py>(
    typ: &schema::BasicValueType,
    v: &Bound<'py, PyAny>,
) -> PyResult<value::BasicValue> {
    let result = match typ {
        schema::BasicValueType::Bytes => {
            value::BasicValue::Bytes(Bytes::from(v.extract::<Vec<u8>>()?))
        }
        schema::BasicValueType::Str => value::BasicValue::Str(Arc::from(v.extract::<String>()?)),
        schema::BasicValueType::Bool => value::BasicValue::Bool(v.extract::<bool>()?),
        schema::BasicValueType::Int64 => value::BasicValue::Int64(v.extract::<i64>()?),
        schema::BasicValueType::Float32 => value::BasicValue::Float32(v.extract::<f32>()?),
        schema::BasicValueType::Float64 => value::BasicValue::Float64(v.extract::<f64>()?),
        schema::BasicValueType::Range => value::BasicValue::Range(depythonize(v)?),
        schema::BasicValueType::Uuid => value::BasicValue::Uuid(v.extract::<uuid::Uuid>()?),
        schema::BasicValueType::Date => value::BasicValue::Date(v.extract::<chrono::NaiveDate>()?),
        schema::BasicValueType::Time => value::BasicValue::Time(v.extract::<chrono::NaiveTime>()?),
        schema::BasicValueType::LocalDateTime => {
            value::BasicValue::LocalDateTime(v.extract::<chrono::NaiveDateTime>()?)
        }
        schema::BasicValueType::OffsetDateTime => {
            if v.getattr_opt("tzinfo")?
                .ok_or_else(|| {
                    PyErr::new::<PyTypeError, _>(format!(
                        "expecting a datetime.datetime value, got {}",
                        v.get_type()
                    ))
                })?
                .is_none()
            {
                value::BasicValue::OffsetDateTime(
                    v.extract::<chrono::NaiveDateTime>()?.and_utc().into(),
                )
            } else {
                value::BasicValue::OffsetDateTime(
                    v.extract::<chrono::DateTime<chrono::FixedOffset>>()?,
                )
            }
        }
        schema::BasicValueType::TimeDelta => {
            value::BasicValue::TimeDelta(v.extract::<chrono::TimeDelta>()?)
        }
        schema::BasicValueType::Json => {
            value::BasicValue::Json(Arc::from(depythonize::<serde_json::Value>(v)?))
        }
        schema::BasicValueType::Vector(elem) => {
            if let Some(vector) = handle_ndarray_from_py(&elem.element_type, v)? {
                vector
            } else {
                // Fallback to list
                value::BasicValue::Vector(Arc::from(
                    v.extract::<Vec<Bound<'py, PyAny>>>()?
                        .into_iter()
                        .map(|v| basic_value_from_py_object(&elem.element_type, &v))
                        .collect::<PyResult<Vec<_>>>()?,
                ))
            }
        }
        schema::BasicValueType::Union(s) => {
            let mut valid_value = None;

            // Try parsing the value
            for (i, typ) in s.types.iter().enumerate() {
                if let Ok(value) = basic_value_from_py_object(typ, v) {
                    valid_value = Some(value::BasicValue::UnionVariant {
                        tag_id: i,
                        value: Box::new(value),
                    });
                    break;
                }
            }

            valid_value.ok_or_else(|| {
                PyErr::new::<PyTypeError, _>(format!(
                    "invalid union value: {}, available types: {:?}",
                    v, s.types
                ))
            })?
        }
    };
    Ok(result)
}

// Helper function to convert PyAny to BasicValue for NDArray
fn handle_ndarray_from_py<'py>(
    elem_type: &schema::BasicValueType,
    v: &Bound<'py, PyAny>,
) -> PyResult<Option<value::BasicValue>> {
    macro_rules! try_convert {
        ($t:ty, $cast:expr) => {
            if let Ok(array) = v.cast::<PyArrayDyn<$t>>() {
                let data = array.readonly().as_slice()?.to_vec();
                let vec = data.into_iter().map($cast).collect::<Vec<_>>();
                return Ok(Some(value::BasicValue::Vector(Arc::from(vec))));
            }
        };
    }

    match *elem_type {
        schema::BasicValueType::Float32 => try_convert!(f32, value::BasicValue::Float32),
        schema::BasicValueType::Float64 => try_convert!(f64, value::BasicValue::Float64),
        schema::BasicValueType::Int64 => try_convert!(i64, value::BasicValue::Int64),
        _ => {}
    }

    Ok(None)
}

// Helper function to convert BasicValue::Vector to PyAny
fn handle_vector_to_py<'py>(
    py: Python<'py>,
    v: &[value::BasicValue],
) -> PyResult<Bound<'py, PyAny>> {
    match v.first() {
        Some(value::BasicValue::Float32(_)) => {
            let data = v
                .iter()
                .map(|x| match x {
                    value::BasicValue::Float32(f) => Ok(*f),
                    _ => Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
                        "Expected all elements to be Float32",
                    )),
                })
                .collect::<PyResult<Vec<_>>>()?;

            Ok(PyArray1::from_vec(py, data).into_any())
        }
        Some(value::BasicValue::Float64(_)) => {
            let data = v
                .iter()
                .map(|x| match x {
                    value::BasicValue::Float64(f) => Ok(*f),
                    _ => Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
                        "Expected all elements to be Float64",
                    )),
                })
                .collect::<PyResult<Vec<_>>>()?;

            Ok(PyArray1::from_vec(py, data).into_any())
        }
        Some(value::BasicValue::Int64(_)) => {
            let data = v
                .iter()
                .map(|x| match x {
                    value::BasicValue::Int64(i) => Ok(*i),
                    _ => Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
                        "Expected all elements to be Int64",
                    )),
                })
                .collect::<PyResult<Vec<_>>>()?;

            Ok(PyArray1::from_vec(py, data).into_any())
        }
        _ => Ok(v
            .iter()
            .map(|v| basic_value_to_py_object(py, v))
            .collect::<PyResult<Vec<_>>>()?
            .into_bound_py_any(py)?),
    }
}

pub fn field_values_from_py_seq<'py>(
    fields_schema: &[schema::FieldSchema],
    v: &Bound<'py, PyAny>,
) -> PyResult<value::FieldValues> {
    let list = v.extract::<Vec<Bound<'py, PyAny>>>()?;
    if list.len() != fields_schema.len() {
        return Err(PyException::new_err(format!(
            "struct field number mismatch, expected {}, got {}",
            fields_schema.len(),
            list.len()
        )));
    }

    Ok(value::FieldValues {
        fields: std::iter::zip(fields_schema, list.into_iter())
            .map(|(f, v)| value_from_py_object(&f.value_type.typ, &v))
            .collect::<PyResult<Vec<_>>>()?,
    })
}

pub fn value_from_py_object<'py>(
    typ: &schema::ValueType,
    v: &Bound<'py, PyAny>,
) -> PyResult<value::Value> {
    let result = if v.is_none() {
        value::Value::Null
    } else {
        match typ {
            schema::ValueType::Basic(typ) => {
                value::Value::Basic(basic_value_from_py_object(typ, v)?)
            }
            schema::ValueType::Struct(schema) => {
                value::Value::Struct(field_values_from_py_seq(&schema.fields, v)?)
            }
            schema::ValueType::Table(schema) => {
                let list = v.extract::<Vec<Bound<'py, PyAny>>>()?;
                let values = list
                    .into_iter()
                    .map(|v| field_values_from_py_seq(&schema.row.fields, &v))
                    .collect::<PyResult<Vec<_>>>()?;

                match schema.kind {
                    schema::TableKind::UTable => {
                        value::Value::UTable(values.into_iter().map(|v| v.into()).collect())
                    }
                    schema::TableKind::LTable => {
                        value::Value::LTable(values.into_iter().map(|v| v.into()).collect())
                    }

                    schema::TableKind::KTable(info) => {
                        let num_key_parts = info.num_key_parts;
                        let k_table_values = values
                            .into_iter()
                            .map(|v| {
                                let mut iter = v.fields.into_iter();
                                if iter.len() < num_key_parts {
                                    client_bail!(
                                        "Invalid KTable value: expect at least {} fields, got {}",
                                        num_key_parts,
                                        iter.len()
                                    );
                                }
                                let keys: Box<[value::KeyPart]> = (0..num_key_parts)
                                    .map(|_| iter.next().unwrap().into_key())
                                    .collect::<Result<_>>()?;
                                let values = value::FieldValues {
                                    fields: iter.collect::<Vec<_>>(),
                                };
                                Ok((KeyValue(keys), values.into()))
                            })
                            .collect::<Result<BTreeMap<_, _>>>();
                        let k_table_values = k_table_values.into_py_result()?;

                        value::Value::KTable(k_table_values)
                    }
                }
            }
        }
    };
    Ok(result)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::base::schema;
    use crate::base::value;
    use crate::base::value::ScopeValue;
    use pyo3::Python;
    use std::collections::BTreeMap;
    use std::sync::Arc;

    fn assert_roundtrip_conversion(original_value: &value::Value, value_type: &schema::ValueType) {
        Python::attach(|py| {
            // Convert Rust value to Python object using value_to_py_object
            let py_object = value_to_py_object(py, original_value)
                .expect("Failed to convert Rust value to Python object");

            println!("Python object: {py_object:?}");
            let roundtripped_value = value_from_py_object(value_type, &py_object)
                .expect("Failed to convert Python object back to Rust value");

            println!("Roundtripped value: {roundtripped_value:?}");
            assert_eq!(
                original_value, &roundtripped_value,
                "Value mismatch after roundtrip"
            );
        });
    }

    #[test]
    fn test_roundtrip_basic_values() {
        let values_and_types = vec![
            (
                value::Value::Basic(value::BasicValue::Int64(42)),
                schema::ValueType::Basic(schema::BasicValueType::Int64),
            ),
            (
                value::Value::Basic(value::BasicValue::Float64(3.14)),
                schema::ValueType::Basic(schema::BasicValueType::Float64),
            ),
            (
                value::Value::Basic(value::BasicValue::Str(Arc::from("hello"))),
                schema::ValueType::Basic(schema::BasicValueType::Str),
            ),
            (
                value::Value::Basic(value::BasicValue::Bool(true)),
                schema::ValueType::Basic(schema::BasicValueType::Bool),
            ),
        ];

        for (val, typ) in values_and_types {
            assert_roundtrip_conversion(&val, &typ);
        }
    }

    #[test]
    fn test_roundtrip_struct() {
        let struct_schema = schema::StructSchema {
            description: Some(Arc::from("Test struct description")),
            fields: Arc::new(vec![
                schema::FieldSchema {
                    name: "a".to_string(),
                    value_type: schema::EnrichedValueType {
                        typ: schema::ValueType::Basic(schema::BasicValueType::Int64),
                        nullable: false,
                        attrs: Default::default(),
                    },
                    description: None,
                },
                schema::FieldSchema {
                    name: "b".to_string(),
                    value_type: schema::EnrichedValueType {
                        typ: schema::ValueType::Basic(schema::BasicValueType::Str),
                        nullable: false,
                        attrs: Default::default(),
                    },
                    description: None,
                },
            ]),
        };

        let struct_val_data = value::FieldValues {
            fields: vec![
                value::Value::Basic(value::BasicValue::Int64(10)),
                value::Value::Basic(value::BasicValue::Str(Arc::from("world"))),
            ],
        };

        let struct_val = value::Value::Struct(struct_val_data);
        let struct_typ = schema::ValueType::Struct(struct_schema); // No clone needed

        assert_roundtrip_conversion(&struct_val, &struct_typ);
    }

    #[test]
    fn test_roundtrip_table_types() {
        let row_schema_struct = Arc::new(schema::StructSchema {
            description: Some(Arc::from("Test table row description")),
            fields: Arc::new(vec![
                schema::FieldSchema {
                    name: "key_col".to_string(), // Will be used as key for KTable implicitly
                    value_type: schema::EnrichedValueType {
                        typ: schema::ValueType::Basic(schema::BasicValueType::Int64),
                        nullable: false,
                        attrs: Default::default(),
                    },
                    description: None,
                },
                schema::FieldSchema {
                    name: "data_col_1".to_string(),
                    value_type: schema::EnrichedValueType {
                        typ: schema::ValueType::Basic(schema::BasicValueType::Str),
                        nullable: false,
                        attrs: Default::default(),
                    },
                    description: None,
                },
                schema::FieldSchema {
                    name: "data_col_2".to_string(),
                    value_type: schema::EnrichedValueType {
                        typ: schema::ValueType::Basic(schema::BasicValueType::Bool),
                        nullable: false,
                        attrs: Default::default(),
                    },
                    description: None,
                },
            ]),
        });

        let row1_fields = value::FieldValues {
            fields: vec![
                value::Value::Basic(value::BasicValue::Int64(1)),
                value::Value::Basic(value::BasicValue::Str(Arc::from("row1_data"))),
                value::Value::Basic(value::BasicValue::Bool(true)),
            ],
        };
        let row1_scope_val: value::ScopeValue = row1_fields.into();

        let row2_fields = value::FieldValues {
            fields: vec![
                value::Value::Basic(value::BasicValue::Int64(2)),
                value::Value::Basic(value::BasicValue::Str(Arc::from("row2_data"))),
                value::Value::Basic(value::BasicValue::Bool(false)),
            ],
        };
        let row2_scope_val: value::ScopeValue = row2_fields.into();

        // UTable
        let utable_schema = schema::TableSchema {
            kind: schema::TableKind::UTable,
            row: (*row_schema_struct).clone(),
        };
        let utable_val = value::Value::UTable(vec![row1_scope_val.clone(), row2_scope_val.clone()]);
        let utable_typ = schema::ValueType::Table(utable_schema);
        assert_roundtrip_conversion(&utable_val, &utable_typ);

        // LTable
        let ltable_schema = schema::TableSchema {
            kind: schema::TableKind::LTable,
            row: (*row_schema_struct).clone(),
        };
        let ltable_val = value::Value::LTable(vec![row1_scope_val.clone(), row2_scope_val.clone()]);
        let ltable_typ = schema::ValueType::Table(ltable_schema);
        assert_roundtrip_conversion(&ltable_val, &ltable_typ);

        // KTable
        let ktable_schema = schema::TableSchema {
            kind: schema::TableKind::KTable(schema::KTableInfo { num_key_parts: 1 }),
            row: (*row_schema_struct).clone(),
        };
        let mut ktable_data = BTreeMap::new();

        // Create KTable entries where the ScopeValue doesn't include the key field
        // This matches how the Python code will serialize/deserialize
        let row1_fields = value::FieldValues {
            fields: vec![
                value::Value::Basic(value::BasicValue::Str(Arc::from("row1_data"))),
                value::Value::Basic(value::BasicValue::Bool(true)),
            ],
        };
        let row1_scope_val: value::ScopeValue = row1_fields.into();

        let row2_fields = value::FieldValues {
            fields: vec![
                value::Value::Basic(value::BasicValue::Str(Arc::from("row2_data"))),
                value::Value::Basic(value::BasicValue::Bool(false)),
            ],
        };
        let row2_scope_val: value::ScopeValue = row2_fields.into();

        // For KTable, the key is extracted from the first field of ScopeValue based on current serialization
        let key1 = value::Value::<ScopeValue>::Basic(value::BasicValue::Int64(1))
            .into_key()
            .unwrap();
        let key2 = value::Value::<ScopeValue>::Basic(value::BasicValue::Int64(2))
            .into_key()
            .unwrap();

        ktable_data.insert(KeyValue(Box::from([key1])), row1_scope_val.clone());
        ktable_data.insert(KeyValue(Box::from([key2])), row2_scope_val.clone());

        let ktable_val = value::Value::KTable(ktable_data);
        let ktable_typ = schema::ValueType::Table(ktable_schema);
        assert_roundtrip_conversion(&ktable_val, &ktable_typ);
    }
}
