use pyo3::prelude::*;
use pyo3::types::{PyByteArray, PyBytes, PyDict, PySlice, PyString};

use crate::ops::{CommandCategory, CommandType, MoveCmd, OpCategory};
use crate::types::Point3D;

use super::dict::{create_and_append_command, py_to_axis_map_helper};
use crate::python::ops::axis::PyAxis;

pub fn ops_to_numpy_arrays(
    py: Python<'_>,
    ops: &crate::ops::Ops,
) -> PyResult<Py<PyDict>> {
    let num_cmds = ops.len();

    let num_arcs: usize = (0..num_cmds)
        .filter(|&i| ops.commands[i].command_type() == CommandType::ArcTo)
        .count();
    let num_beziers: usize = (0..num_cmds)
        .filter(|&i| {
            ops.commands[i].command_type() == CommandType::BezierTo
                || ops.commands[i].command_type()
                    == CommandType::QuadraticBezierTo
        })
        .count();

    let scanline_lengths: Vec<usize> = ops
        .commands
        .iter()
        .filter_map(|node| {
            if let OpCategory::Moving {
                cmd: MoveCmd::ScanLine { power_values },
                ..
            } = &node.category
            {
                Some(power_values.len())
            } else {
                None
            }
        })
        .collect();
    let total_scanline_bytes: usize = scanline_lengths.iter().sum();
    let num_scanlines = scanline_lengths.len();

    let numpy = py.import("numpy")?;

    let types = numpy
        .call_method1("zeros", (num_cmds,))?
        .call_method1("astype", ("int32",))?;
    let endpoints = numpy
        .call_method1("zeros", ((num_cmds, 3),))?
        .call_method1("astype", ("float32",))?;

    let arc_data = numpy
        .call_method1("zeros", ((num_arcs, 3),))?
        .call_method1("astype", ("float32",))?;
    let arc_map = numpy
        .call_method1("full", (num_cmds, -1))?
        .call_method1("astype", ("int32",))?;

    let bezier_data = numpy
        .call_method1("zeros", ((num_beziers, 6),))?
        .call_method1("astype", ("float32",))?;
    let bezier_map = numpy
        .call_method1("full", (num_cmds, -1))?
        .call_method1("astype", ("int32",))?;

    let scanline_data_arr = numpy
        .call_method1("zeros", (total_scanline_bytes,))?
        .call_method1("astype", ("uint8",))?;
    let scanline_map = numpy
        .call_method1("full", (num_cmds, -1))?
        .call_method1("astype", ("int32",))?;
    let scanline_indices = numpy
        .call_method1("zeros", ((num_scanlines, 2),))?
        .call_method1("astype", ("int32",))?;

    let mut state_marker_cmds_data: Vec<(usize, Py<PyDict>)> = Vec::new();
    let mut extra_axes_map: Vec<(usize, Py<PyDict>)> = Vec::new();
    let mut arc_idx: usize = 0;
    let mut bezier_idx: usize = 0;
    let mut scanline_idx: usize = 0;
    let mut scanline_offset: usize = 0;

    for i in 0..num_cmds {
        let node = &ops.commands[i];
        let ct = node.command_type();

        types.call_method1("__setitem__", (i, ct as i32))?;

        if let OpCategory::Moving { end, cmd } = &node.category {
            endpoints
                .call_method1("__setitem__", (i, vec![end.x, end.y, end.z]))?;

            match cmd {
                MoveCmd::BezierTo { control1, control2 } => {
                    bezier_data.call_method1(
                        "__setitem__",
                        (
                            bezier_idx,
                            vec![
                                control1.x, control1.y, control1.z, control2.x,
                                control2.y, control2.z,
                            ],
                        ),
                    )?;
                    bezier_map
                        .call_method1("__setitem__", (i, bezier_idx as i32))?;
                    bezier_idx += 1;
                }
                MoveCmd::QuadraticBezierTo { control } => {
                    bezier_data.call_method1(
                        "__setitem__",
                        (
                            bezier_idx,
                            vec![
                                control.x, control.y, control.z, 0.0, 0.0, 0.0,
                            ],
                        ),
                    )?;
                    bezier_map
                        .call_method1("__setitem__", (i, bezier_idx as i32))?;
                    bezier_idx += 1;
                }
                MoveCmd::ArcTo { center, cw } => {
                    arc_data.call_method1(
                        "__setitem__",
                        (
                            arc_idx,
                            vec![
                                center.x,
                                center.y,
                                if *cw { 1.0 } else { 0.0 },
                            ],
                        ),
                    )?;
                    arc_map.call_method1("__setitem__", (i, arc_idx as i32))?;
                    arc_idx += 1;
                }
                MoveCmd::ScanLine { power_values } => {
                    let length = power_values.len();
                    let py_bytes = PyByteArray::new(py, power_values.as_ref());
                    let slice = PySlice::new(
                        py,
                        scanline_offset as isize,
                        (scanline_offset + length) as isize,
                        1,
                    );
                    scanline_data_arr
                        .call_method1("__setitem__", (slice, py_bytes))?;
                    scanline_indices.call_method1(
                        "__setitem__",
                        (
                            scanline_idx,
                            (scanline_offset, scanline_offset + length),
                        ),
                    )?;
                    scanline_map.call_method1(
                        "__setitem__",
                        (i, scanline_idx as i32),
                    )?;
                    scanline_offset += length;
                    scanline_idx += 1;
                }
                _ => {}
            }

            if let Some(ea) = node.extra_axes() {
                let ea_dict = PyDict::new(py);
                for &(axis, val) in ea {
                    let py_axis = Py::new(py, PyAxis(axis))?;
                    let label: String =
                        py_axis.bind(py).getattr("label")?.extract()?;
                    ea_dict.set_item(label, val)?;
                }
                extra_axes_map.push((i, ea_dict.unbind()));
            }
        } else {
            let d = super::dict::cmd_to_dict(py, ops, i)?;
            state_marker_cmds_data.push((i, d.into()));
        }
    }

    let json_mod = py.import("json")?;
    let sm_dict = PyDict::new(py);
    for (idx, d) in &state_marker_cmds_data {
        sm_dict.set_item(idx.to_string(), d)?;
    }
    let json_str: String =
        json_mod.call_method1("dumps", (sm_dict,))?.extract()?;
    let json_bytes = numpy.call_method1(
        "frombuffer",
        (PyBytes::new(py, json_str.as_bytes()), "uint8"),
    )?;

    let result = PyDict::new(py);
    result.set_item("types", types)?;
    result.set_item("endpoints", endpoints)?;
    result.set_item("arc_data", arc_data)?;
    result.set_item("arc_map", arc_map)?;
    result.set_item("bezier_data", bezier_data)?;
    result.set_item("bezier_map", bezier_map)?;
    result.set_item("scanline_data", scanline_data_arr)?;
    result.set_item("scanline_indices", scanline_indices)?;
    result.set_item("scanline_map", scanline_map)?;
    result.set_item("state_marker_json_bytes", json_bytes)?;

    if !extra_axes_map.is_empty() {
        let ea_dict = PyDict::new(py);
        for (idx, d) in &extra_axes_map {
            ea_dict.set_item(idx.to_string(), d)?;
        }
        let ea_json_str: String =
            json_mod.call_method1("dumps", (ea_dict,))?.extract()?;
        let ea_json_bytes = numpy.call_method1(
            "frombuffer",
            (PyBytes::new(py, ea_json_str.as_bytes()), "uint8"),
        )?;
        result.set_item("extra_axes_json", ea_json_bytes)?;
    }

    Ok(result.unbind())
}

pub fn ops_from_numpy_arrays(
    arrays: &Bound<'_, PyDict>,
) -> PyResult<crate::ops::Ops> {
    let py = arrays.py();
    let numpy = py.import("numpy")?;

    let types_arr: Vec<i32> = arrays
        .get_item("types")?
        .ok_or_else(|| {
            pyo3::exceptions::PyKeyError::new_err("missing 'types'")
        })?
        .call_method0("tolist")?
        .extract()?;

    let mut ops = crate::ops::Ops::new();

    let json_bytes_bound = arrays
        .get_item("state_marker_json_bytes")?
        .unwrap_or_else(|| {
            numpy.call_method1("array", (Vec::<u8>::new(),)).unwrap()
        });
    let json_bytes_data: Vec<u8> =
        json_bytes_bound.call_method1("tobytes", ())?.extract()?;
    let json_mod = py.import("json")?;
    let state_marker_cmds_data: Py<PyDict> = if json_bytes_data.is_empty() {
        PyDict::new(py).unbind()
    } else {
        let json_str = String::from_utf8(json_bytes_data).map_err(|e| {
            pyo3::exceptions::PyValueError::new_err(e.to_string())
        })?;
        json_mod.call_method1("loads", (json_str,))?.extract()?
    };

    let extra_axes_data_bound: Option<Bound<'_, PyAny>> =
        arrays.get_item("extra_axes_json")?;
    let extra_axes_data: Py<PyDict> = if let Some(ref ea) =
        extra_axes_data_bound
    {
        let ea_bytes: Vec<u8> = ea.call_method1("tobytes", ())?.extract()?;
        if ea_bytes.is_empty() {
            PyDict::new(py).unbind()
        } else {
            let ea_str = String::from_utf8(ea_bytes).map_err(|e| {
                pyo3::exceptions::PyValueError::new_err(e.to_string())
            })?;
            json_mod.call_method1("loads", (ea_str,))?.extract()?
        }
    } else {
        PyDict::new(py).unbind()
    };

    let endpoints_bound = arrays.get_item("endpoints")?.ok_or_else(|| {
        pyo3::exceptions::PyKeyError::new_err("missing 'endpoints'")
    })?;
    let arc_data_bound = arrays.get_item("arc_data")?.ok_or_else(|| {
        pyo3::exceptions::PyKeyError::new_err("missing 'arc_data'")
    })?;
    let arc_map_bound = arrays.get_item("arc_map")?.ok_or_else(|| {
        pyo3::exceptions::PyKeyError::new_err("missing 'arc_map'")
    })?;
    let bezier_data_bound =
        arrays.get_item("bezier_data")?.unwrap_or_else(|| {
            numpy
                .call_method1("zeros", ((0, 6),))
                .unwrap()
                .call_method1("astype", ("float32",))
                .unwrap()
        });
    let bezier_map_bound =
        arrays.get_item("bezier_map")?.unwrap_or_else(|| {
            numpy
                .call_method1("full", (0, -1))
                .unwrap()
                .call_method1("astype", ("int32",))
                .unwrap()
        });
    let scanline_data_bound =
        arrays.get_item("scanline_data")?.ok_or_else(|| {
            pyo3::exceptions::PyKeyError::new_err("missing 'scanline_data'")
        })?;
    let scanline_indices_bound =
        arrays.get_item("scanline_indices")?.ok_or_else(|| {
            pyo3::exceptions::PyKeyError::new_err("missing 'scanline_indices'")
        })?;
    let scanline_map_bound =
        arrays.get_item("scanline_map")?.ok_or_else(|| {
            pyo3::exceptions::PyKeyError::new_err("missing 'scanline_map'")
        })?;

    for (i, &cmd_type_val) in types_arr.iter().enumerate() {
        let i_str = i.to_string();
        let py_key = PyString::new(py, &i_str);

        let sm_bound = state_marker_cmds_data.bind(py);
        if let Some(cmd_data_any) = sm_bound.get_item(&py_key)? {
            let cmd_data = cmd_data_any.cast::<PyDict>()?;
            create_and_append_command(cmd_data, &mut ops)?;
            continue;
        }

        let ct = CommandType::try_from(cmd_type_val as u8).map_err(|e| {
            pyo3::exceptions::PyValueError::new_err(format!("{e}"))
        })?;

        if ct.category() != CommandCategory::Moving {
            continue;
        }

        let end_list: Vec<f32> = endpoints_bound
            .call_method1("__getitem__", (i,))?
            .call_method0("tolist")?
            .extract()?;
        let end_tuple = Point3D::new(
            end_list[0] as f64,
            end_list[1] as f64,
            end_list[2] as f64,
        );

        match ct {
            CommandType::MoveTo => {
                ops.move_to(end_tuple.x, end_tuple.y, end_tuple.z, None);
            }
            CommandType::LineTo => {
                ops.line_to(end_tuple.x, end_tuple.y, end_tuple.z, None);
            }
            CommandType::ArcTo => {
                let arc_idx_val: i32 = arc_map_bound
                    .call_method1("__getitem__", (i,))?
                    .extract()?;
                let arc_vals: Vec<f32> = arc_data_bound
                    .call_method1("__getitem__", (arc_idx_val as isize,))?
                    .call_method0("tolist")?
                    .extract()?;
                ops.arc_to(
                    end_tuple.x,
                    end_tuple.y,
                    arc_vals[0] as f64,
                    arc_vals[1] as f64,
                    arc_vals[2] != 0.0,
                    end_tuple.z,
                    None,
                );
            }
            CommandType::BezierTo => {
                let bez_idx_val: i32 = bezier_map_bound
                    .call_method1("__getitem__", (i,))?
                    .extract()?;
                let bez_vals: Vec<f32> = bezier_data_bound
                    .call_method1("__getitem__", (bez_idx_val as isize,))?
                    .call_method0("tolist")?
                    .extract()?;
                ops.bezier_to(
                    Point3D::new(
                        bez_vals[0] as f64,
                        bez_vals[1] as f64,
                        bez_vals[2] as f64,
                    ),
                    Point3D::new(
                        bez_vals[3] as f64,
                        bez_vals[4] as f64,
                        bez_vals[5] as f64,
                    ),
                    end_tuple,
                    None,
                );
            }
            CommandType::QuadraticBezierTo => {
                let bez_idx_val: i32 = bezier_map_bound
                    .call_method1("__getitem__", (i,))?
                    .extract()?;
                let bez_vals: Vec<f32> = bezier_data_bound
                    .call_method1("__getitem__", (bez_idx_val as isize,))?
                    .call_method0("tolist")?
                    .extract()?;
                ops.quadratic_bezier_to(
                    Point3D::new(
                        bez_vals[0] as f64,
                        bez_vals[1] as f64,
                        bez_vals[2] as f64,
                    ),
                    end_tuple,
                    None,
                );
            }
            CommandType::ScanLine => {
                let scan_idx_val: i32 = scanline_map_bound
                    .call_method1("__getitem__", (i,))?
                    .extract()?;
                let si_list: Vec<i32> = scanline_indices_bound
                    .call_method1("__getitem__", (scan_idx_val as isize,))?
                    .call_method0("tolist")?
                    .extract()?;
                let start = si_list[0] as usize;
                let end = si_list[1] as usize;
                let pv_bytes: Vec<u8> = scanline_data_bound
                    .call_method1(
                        "__getitem__",
                        ((start..end).collect::<Vec<_>>(),),
                    )?
                    .call_method1("tobytes", ())?
                    .extract()?;
                ops.scan_to(
                    end_tuple.x,
                    end_tuple.y,
                    end_tuple.z,
                    pv_bytes,
                    None,
                );
            }
            _ => continue,
        }

        let ea_bound = extra_axes_data.bind(py);
        if let Some(ea_item) = ea_bound.get_item(&py_key)? {
            let ea_dict = ea_item.cast::<PyDict>()?;
            let ea_vec = py_to_axis_map_helper(ea_dict)?;
            let last_idx = ops.len() - 1;
            ops.commands[last_idx].set_extra_axes(std::sync::Arc::from(ea_vec));
        }
    }

    Ok(ops)
}
