#![allow(clippy::borrow_deref_ref)] // clippy warns about code generated by #[pymethods]

use dora_node_api::{DoraNode, EventStream};
use dora_operator_api_python::{process_python_output, pydict_to_metadata, PyEvent};
use eyre::Context;
use pyo3::prelude::*;
use pyo3::types::PyDict;

#[pyclass]
pub struct Node {
    events: EventStream,
    node: DoraNode,
}

#[pymethods]
impl Node {
    #[new]
    pub fn new() -> eyre::Result<Self> {
        let (node, events) = DoraNode::init_from_env()?;

        Ok(Node { events, node })
    }

    #[allow(clippy::should_implement_trait)]
    pub fn next(&mut self, py: Python) -> PyResult<Option<PyEvent>> {
        self.__next__(py)
    }

    pub fn __next__(&mut self, py: Python) -> PyResult<Option<PyEvent>> {
        let event = py.allow_threads(|| self.events.recv());
        Ok(event.map(PyEvent::from))
    }

    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    pub fn send_output(
        &mut self,
        output_id: String,
        data: PyObject,
        metadata: Option<&PyDict>,
        py: Python,
    ) -> eyre::Result<()> {
        process_python_output(&data, py, |data| {
            self.send_output_slice(output_id, data.len(), data, metadata)
        })
    }
}

impl Node {
    fn send_output_slice(
        &mut self,
        output_id: String,
        len: usize,
        data: &[u8],
        metadata: Option<&PyDict>,
    ) -> eyre::Result<()> {
        let metadata = pydict_to_metadata(metadata)?;
        self.node
            .send_output(output_id.into(), metadata, len, |out| {
                out.copy_from_slice(data);
            })
            .wrap_err("failed to send output")
    }

    pub fn id(&self) -> String {
        self.node.id().to_string()
    }
}

#[pyfunction]
fn start_runtime() -> eyre::Result<()> {
    dora_runtime::main().wrap_err("Dora Runtime raised an error.")
}

#[pymodule]
fn dora(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(start_runtime, m)?)?;
    m.add_class::<Node>().unwrap();
    Ok(())
}
