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

use std::env::current_dir;
use std::path::PathBuf;
use std::sync::{Arc, LazyLock};
use std::time::Duration;
use tokio::sync::Mutex;

use arrow::pyarrow::{FromPyArrow, ToPyArrow};
use dora_download::download_file;
use dora_node_api::dora_core::config::NodeId;
use dora_node_api::dora_core::descriptor::source_is_url;
use dora_node_api::merged::{MergeExternalSend, MergedEvent};
use dora_node_api::{DataflowId, DoraNode, EventStream, TryRecvError, init_tracing};
use dora_operator_api_python::{DelayedCleanup, NodeCleanupHandle, PyEvent, pydict_to_metadata};
use dora_ros2_bridge_python::Ros2Subscription;
use eyre::{Context, ContextCompat};

use futures::{Stream, StreamExt};
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict};
use pyo3_special_method_derive::{Dict, Dir, Repr, Str};
use tokio::runtime::{Builder, Runtime};
use tracing::{Level, span};
static RUNTIME: LazyLock<Runtime> = LazyLock::new(|| {
    Builder::new_multi_thread()
        .worker_threads(4)
        .enable_all()
        .build()
        .context("Failed to create Tokio runtime")
        .unwrap()
});

/// Consume a Python `logging.LogRecord` and emit a Rust `tracing::Event` instead.
#[pyfunction]
fn host_log<'py>(record: Bound<'py, PyAny>) -> PyResult<()> {
    let level = record.getattr("levelno")?.extract::<u8>()?;
    let message = record.getattr("getMessage")?.call0()?.to_string();
    let pathname = record.getattr("pathname")?.to_string();
    let lineno = record.getattr("lineno")?.to_string();
    let target = record.getattr("name")?.to_string();

    RUNTIME.spawn(async move {
    if level.ge(&40u8) {
        let span = span!(Level::ERROR, "dora.python.log.error", file=pathname, line=lineno, %target, %message);
        let _enter = span.enter();
        tracing::event!(tracing::Level::ERROR, file=pathname, line=lineno, %target, %message);
    } else if level.ge(&30u8) {
        let span = span!(Level::ERROR, "dora.python.log.warn", file=pathname, line=lineno, %target, %message);
        let _enter = span.enter();
        tracing::event!(tracing::Level::WARN, file=pathname, line=lineno, %target, %message);
    } else if level.ge(&20u8){
        let span = span!(Level::INFO, "dora.python.log.info", file=pathname, line=lineno, %target, %message);
        let _enter = span.enter();
        tracing::event!(tracing::Level::INFO, file=pathname, line=lineno, %target, %message);
    } else if level.ge(&10u8) {
        let span = span!(Level::DEBUG, "dora.python.log.debug", file=pathname, line=lineno, %target, %message);
        let _enter = span.enter();
        tracing::event!(tracing::Level::DEBUG, file=pathname, line=lineno, %target, %message);
    } else {
        let span = span!(Level::TRACE, "dora.python.log.trace", file=pathname, line=lineno, %target, %message);
        let _enter = span.enter();
        tracing::event!(tracing::Level::TRACE, file=pathname, line=lineno, %target, %message);
    }
    });
    Ok(())
}

/// Modifies the Python `logging` module to deliver its log messages to the host `tracing::Subscriber` by default.
/// To achieve this goal, the following changes are made to the module:
/// - A new builtin function `logging.host_log` transcodes `logging.LogRecord`s to `tracing::Event`s. This function
///   is not exported in `logging.__all__`, as it is not intended to be called directly.
/// - A new class `logging.HostHandler` provides a `logging.Handler` that delivers all records to `host_log`.
/// - `logging.basicConfig` is changed to use `logging.HostHandler` by default.
///
/// Since any call like `logging.warn(...)` sets up logging via `logging.basicConfig`, all log messages are now
/// delivered to `crate::host_log`, which will send them to `tracing::event!`.
pub fn setup_logging(py: Python) -> PyResult<()> {
    let logging = py.import("logging")?;
    logging.setattr("host_log", wrap_pyfunction!(host_log, &logging)?)?;
    py.run(
        cr#"
class HostHandler(Handler):
	def __init__(self, level=0):
		super().__init__(level=level)
	
	def emit(self, record):
		host_log(record)

oldBasicConfig = basicConfig

def basicConfig(*pargs, **kwargs):
	if "handlers" not in kwargs:
		kwargs["handlers"] = [HostHandler()]
	else:
		kwargs["handlers"].append(HostHandler())
	if "level" not in kwargs:
		kwargs["level"] = 0

	return oldBasicConfig(*pargs, **kwargs)
"#,
        Some(&logging.dict()),
        None,
    )?;

    let all = logging.index()?;
    all.append("HostHandler")?;

    Ok(())
}
/// The custom node API lets you integrate `dora` into your application.
/// It allows you to retrieve input and send output in any fashion you want.
///
/// Use with:
///
/// ```python
/// from dora import Node
///
/// node = Node()
/// ```
///
/// :type node_id: str, optional
#[pyclass]
#[derive(Dir, Dict, Str, Repr)]
pub struct Node {
    events: Events,
    node: DelayedCleanup<DoraNode>,

    dataflow_id: DataflowId,
    node_id: NodeId,
}

#[pymethods]
impl Node {
    #[new]
    #[pyo3(signature = (node_id=None))]
    pub fn new(node_id: Option<String>) -> eyre::Result<Self> {
        let (node, events) = if let Some(node_id) = node_id {
            DoraNode::init_flexible(NodeId::from(node_id))
                .context("Could not setup node from node id. Make sure to have a running dataflow with this dynamic node")?
        } else {
            DoraNode::init_from_env().context("Could not initiate node from environment variable. For dynamic node, please add a node id in the initialization function.")?
        };
        let id = node.id().clone();
        let dataflow_id = node.dataflow_id().clone();
        RUNTIME.spawn(async move {
            let _guard = init_tracing(&id, &dataflow_id).unwrap();
            loop {
                tokio::time::sleep(std::time::Duration::from_secs(1)).await;
            }
        });

        let dataflow_id = *node.dataflow_id();
        let node_id = node.id().clone();
        let node = DelayedCleanup::new(node);
        let events = events;
        let cleanup_handle = NodeCleanupHandle {
            _handles: Arc::new(node.handle()),
        };

        Python::with_gil(|py| {
            // Extend the `logging` module to interact with tracing
            setup_logging(py)
        })?;

        Ok(Node {
            events: Events {
                inner: Arc::new(Mutex::new(EventsInner::Dora(events))),
                _cleanup_handle: cleanup_handle,
            },
            dataflow_id,
            node_id,
            node,
        })
    }

    /// `.next()` gives you the next input that the node has received.
    /// It blocks until the next event becomes available.
    /// You can use timeout in seconds to return if no input is available.
    /// It will return `None` when all senders has been dropped.
    ///
    /// ```python
    /// event = node.next()
    /// ```
    ///
    /// You can also iterate over the event stream with a loop
    ///
    /// ```python
    /// for event in node:
    ///    match event["type"]:
    ///        case "INPUT":
    ///            match event["id"]:
    ///                 case "image":
    /// ```
    ///
    /// :type timeout: float, optional
    /// :rtype: dict
    #[pyo3(signature = (timeout=None))]
    #[allow(clippy::should_implement_trait)]
    pub fn next(&self, py: Python, timeout: Option<f32>) -> PyResult<Option<Py<PyDict>>> {
        let event = py.allow_threads(|| self.events.recv(timeout.map(Duration::from_secs_f32)));
        if let Some(event) = event {
            let dict = event
                .to_py_dict(py)
                .context("Could not convert event into a dict")?;
            Ok(Some(dict))
        } else {
            Ok(None)
        }
    }

    /// `.drain()` gives you all available inputs that the node has received.
    /// It does not block until the next event becomes available.
    ///
    /// ```python
    /// events = node.drain()
    /// for event in events:
    ///     print(event)
    /// ```
    ///
    /// :rtype: list[dict]
    #[allow(clippy::should_implement_trait)]
    pub fn drain(&self, py: Python) -> PyResult<Vec<Py<PyDict>>> {
        let events = self
            .events
            .drain()
            .context("Could not drain events. Channel is closed")?
            .into_iter()
            .map(|event| {
                event
                    .to_py_dict(py)
                    .context("Could not convert event into a dict")
                    .unwrap_or_else(|_| PyDict::new(py).into())
            })
            .collect();
        Ok(events)
    }

    /// `.try_recv()` gives you the next input in the queue that the node has received.
    /// It does not block until the next event becomes available.
    ///
    /// ```python
    /// event = events.try_recv()
    ///     print(event)
    /// ```
    ///
    /// :rtype: dict
    #[allow(clippy::should_implement_trait)]
    pub fn try_recv(&mut self, py: Python) -> Option<Py<PyDict>> {
        match self.events.try_recv() {
            Ok(event) => match event.to_py_dict(py) {
                Ok(dict) => Some(dict),
                Err(_) => None,
            },
            Err(_) => None,
        }
    }

    /// Check if there are any buffered events in the event stream.
    ///
    /// :rtype: bool
    #[allow(clippy::should_implement_trait)]
    pub fn is_empty(&self) -> bool {
        self.events.is_empty()
    }

    /// `.recv_async()` gives you the next input that the node has received asynchronously.
    /// It does not blocks until the next event becomes available.
    /// You can use timeout in seconds to return if no input is available.
    /// It will return an Error if the timeout is reached.
    /// It will return `None` when all senders has been dropped.
    ///
    /// warning::
    ///     This feature is experimental as pyo3 async (rust-python FFI) is still in development.
    ///
    /// ```python
    /// event = await node.recv_async()
    /// ```
    ///
    /// You can also iterate over the event stream with a loop
    ///
    /// :type timeout: float, optional
    /// :rtype: dict
    #[pyo3(signature = (timeout=None))]
    #[allow(clippy::should_implement_trait)]
    pub async fn recv_async(&self, timeout: Option<f32>) -> PyResult<Option<Py<PyDict>>> {
        let event = self
            .events
            .recv_async_timeout(timeout.map(Duration::from_secs_f32))
            .await;
        if let Some(event) = event {
            // Get python
            Python::with_gil(|py| {
                let dict = event
                    .to_py_dict(py)
                    .context("Could not convert event into a dict")?;

                Ok(Some(dict))
            })
        } else {
            Ok(None)
        }
    }

    /// You can iterate over the event stream with a loop
    ///
    /// ```python
    /// for event in node:
    ///    match event["type"]:
    ///        case "INPUT":
    ///            match event["id"]:
    ///                 case "image":
    /// ```
    ///
    /// Default behaviour is to timeout after 2 seconds.
    ///
    /// :rtype: dict
    pub fn __next__(&self, py: Python) -> PyResult<Option<Py<PyDict>>> {
        self.next(py, None)
    }

    /// You can iterate over the event stream with a loop
    ///
    /// ```python
    /// for event in node:
    ///    match event["type"]:
    ///        case "INPUT":
    ///            match event["id"]:
    ///                 case "image":
    /// ```
    ///
    /// :rtype: dict
    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    /// `send_output` send data from the node.
    ///
    /// ```python
    /// Args:
    ///    output_id: str,
    ///    data: pyarrow.Array,
    ///    metadata: Option[Dict],
    /// ```
    ///
    /// ex:
    ///
    /// ```python
    /// node.send_output("string", b"string", {"open_telemetry_context": "7632e76"})
    /// ```
    ///
    /// :type output_id: str
    /// :type data: pyarrow.Array
    /// :type metadata: dict, optional
    /// :rtype: None
    #[pyo3(signature = (output_id, data, metadata=None))]
    pub fn send_output(
        &self,
        output_id: String,
        data: PyObject,
        metadata: Option<Bound<'_, PyDict>>,
        py: Python,
    ) -> eyre::Result<()> {
        let parameters = pydict_to_metadata(metadata)?;

        if let Ok(py_bytes) = data.downcast_bound::<PyBytes>(py) {
            let data = py_bytes.as_bytes();
            self.node
                .get_mut()
                .send_output_bytes(output_id.into(), parameters, data.len(), data)
                .wrap_err("failed to send output")?;
        } else if let Ok(arrow_array) = arrow::array::ArrayData::from_pyarrow_bound(data.bind(py)) {
            self.node.get_mut().send_output(
                output_id.into(),
                parameters,
                arrow::array::make_array(arrow_array),
            )?;
        } else {
            eyre::bail!("invalid `data` type, must by `PyBytes` or arrow array")
        }

        Ok(())
    }

    /// Returns the full dataflow descriptor that this node is part of.
    ///
    /// This method returns the parsed dataflow YAML file.
    ///
    /// :rtype: dict
    pub fn dataflow_descriptor(&self, py: Python) -> eyre::Result<PyObject> {
        Ok(
            pythonize::pythonize(py, &self.node.get_mut().dataflow_descriptor()?)
                .map(|x| x.unbind())?,
        )
    }

    /// Returns the node configuration.
    ///
    /// :rtype: dict
    pub fn node_config(&self, py: Python) -> eyre::Result<PyObject> {
        Ok(pythonize::pythonize(py, &self.node.get_mut().node_config()).map(|x| x.unbind())?)
    }

    /// Returns the dataflow id.
    ///
    /// :rtype: str
    pub fn dataflow_id(&self) -> String {
        self.dataflow_id.to_string()
    }

    /// Merge an external event stream with dora main loop.
    /// This currently only work with ROS2.
    ///
    /// :type subscription: dora.Ros2Subscription
    /// :rtype: None
    pub fn merge_external_events(&self, subscription: &mut Ros2Subscription) -> eyre::Result<()> {
        let subscription = subscription.into_stream()?;
        let stream = futures::stream::poll_fn(move |cx| {
            let s = subscription.as_stream().map(|item| {
                match item.context("failed to read ROS2 message") {
                    Ok((value, _info)) => Python::with_gil(|py| {
                        value
                            .to_pyarrow(py)
                            .context("failed to convert value to pyarrow")
                            .unwrap_or_else(|err| err_to_pyany(err, py))
                    }),
                    Err(err) => Python::with_gil(|py| err_to_pyany(err, py)),
                }
            });
            futures::pin_mut!(s);
            s.poll_next_unpin(cx)
        });

        // take out the event stream and temporarily replace it with a dummy
        let mut inner = self.events.inner.blocking_lock();
        let events = std::mem::replace(
            &mut *inner,
            EventsInner::Merged(Box::new(futures::stream::empty())),
        );
        // update self.events with the merged stream
        *inner = EventsInner::Merged(events.merge_external_send(Box::pin(stream)));

        Ok(())
    }
}

fn err_to_pyany(err: eyre::Report, gil: Python<'_>) -> Py<PyAny> {
    PyErr::from(err)
        .into_pyobject(gil)
        .unwrap_or_else(|infallible| match infallible {})
        .into_any()
        .unbind()
}

struct Events {
    inner: Arc<Mutex<EventsInner>>,
    _cleanup_handle: NodeCleanupHandle,
}

impl Events {
    fn recv(&self, timeout: Option<Duration>) -> Option<PyEvent> {
        let mut inner = self.inner.blocking_lock();
        let event = match &mut *inner {
            EventsInner::Dora(events) => match timeout {
                Some(timeout) => events.recv_timeout(timeout).map(MergedEvent::Dora),
                None => events.recv().map(MergedEvent::Dora),
            },
            EventsInner::Merged(events) => futures::executor::block_on(events.next()),
        };
        event.map(|event| PyEvent { event })
    }

    fn try_recv(&self) -> Result<PyEvent, TryRecvError> {
        let mut inner = self.inner.blocking_lock();
        let event = match &mut *inner {
            EventsInner::Dora(events) => events.try_recv().map(MergedEvent::Dora),
            EventsInner::Merged(_events) => {
                todo!("try_recv on external event stream is not yet implemented!")
            }
        };
        event.map(|event| PyEvent { event })
    }

    async fn recv_async_timeout(&self, timeout: Option<Duration>) -> Option<PyEvent> {
        let mut inner = self.inner.lock().await;
        let event = match &mut *inner {
            EventsInner::Dora(events) => match timeout {
                Some(timeout) => events
                    .recv_async_timeout(timeout)
                    .await
                    .map(MergedEvent::Dora),
                None => events.recv_async().await.map(MergedEvent::Dora),
            },
            EventsInner::Merged(events) => events.next().await,
        };
        event.map(|event| PyEvent { event })
    }

    fn drain(&self) -> Option<Vec<PyEvent>> {
        let mut inner = self.inner.blocking_lock();
        match &mut *inner {
            EventsInner::Dora(events) => match events.drain() {
                Some(items) => {
                    return Some(
                        items
                            .into_iter()
                            .map(MergedEvent::Dora)
                            .map(|event| PyEvent { event })
                            .collect(),
                    );
                }
                None => return None,
            },
            EventsInner::Merged(_events) => {
                todo!("Draining external event is not yet implemented!")
            }
        };
    }

    fn is_empty(&self) -> bool {
        let inner = self.inner.blocking_lock();
        match &*inner {
            EventsInner::Dora(events) => events.is_empty(),
            EventsInner::Merged(_events) => {
                todo!("is_empty on external event stream is not yet implemented!")
            }
        }
    }
}

#[allow(clippy::large_enum_variant)]
enum EventsInner {
    Dora(EventStream),
    Merged(Box<dyn Stream<Item = MergedEvent<PyObject>> + Unpin + Send + Sync>),
}

impl<'a> MergeExternalSend<'a, PyObject> for EventsInner {
    type Item = MergedEvent<PyObject>;

    fn merge_external_send(
        self,
        external_events: impl Stream<Item = PyObject> + Unpin + Send + Sync + 'a,
    ) -> Box<dyn Stream<Item = Self::Item> + Unpin + Send + Sync + 'a> {
        match self {
            EventsInner::Dora(events) => events.merge_external_send(external_events),
            EventsInner::Merged(events) => {
                let merged = events.merge_external_send(external_events);
                Box::new(merged.map(|event| match event {
                    MergedEvent::Dora(e) => MergedEvent::Dora(e),
                    MergedEvent::External(e) => MergedEvent::External(e.flatten()),
                }))
            }
        }
    }
}

impl Node {
    pub fn id(&self) -> String {
        self.node_id.to_string()
    }
}

/// Start a runtime for Operators
///
/// :rtype: None
#[pyfunction]
pub fn start_runtime() -> eyre::Result<()> {
    dora_runtime::main().wrap_err("Dora Runtime raised an error.")
}

pub fn resolve_dataflow(dataflow: String) -> eyre::Result<PathBuf> {
    let dataflow = if source_is_url(&dataflow) {
        // try to download the shared library
        let target_path = current_dir().context("Could not access the current dir")?;
        let rt = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .context("tokio runtime failed")?;
        rt.block_on(async { download_file(&dataflow, &target_path).await })
            .wrap_err("failed to download dataflow yaml file")?
    } else {
        PathBuf::from(dataflow)
    };
    Ok(dataflow)
}

/// Build a Dataflow, exactly the same way as `dora build` command line tool.
///
///
/// :type dataflow_path: str
/// :type uv: bool, optional
/// :type coordinator_addr: str, optional
/// :type coordinator_port: int, optional
/// :type force_local: bool, optional
/// :rtype: None
#[pyfunction]
#[pyo3(signature = (dataflow_path, uv=None, coordinator_addr=None, coordinator_port=None, force_local=false))]
pub fn build(
    dataflow_path: String,
    uv: Option<bool>,
    coordinator_addr: Option<String>,
    coordinator_port: Option<u16>,
    force_local: bool,
) -> eyre::Result<()> {
    dora_cli::build(
        dataflow_path,
        coordinator_addr.map(|addr| addr.parse().unwrap()),
        coordinator_port,
        uv.unwrap_or_default(),
        force_local,
    )
}

/// Run a Dataflow, exactly the same way as `dora run` command line tool.
///
/// :type dataflow_path: str
/// :type uv: bool, optional
/// :rtype: None
#[pyfunction]
#[pyo3(signature = (dataflow_path, uv=None))]
pub fn run(dataflow_path: String, uv: Option<bool>) -> eyre::Result<()> {
    dora_cli::run(dataflow_path, uv.unwrap_or_default())
}

#[pymodule]
fn dora(_py: Python, m: Bound<'_, PyModule>) -> PyResult<()> {
    dora_ros2_bridge_python::create_dora_ros2_bridge_module(&m)?;

    m.add_function(wrap_pyfunction!(start_runtime, &m)?)?;
    m.add_function(wrap_pyfunction!(run, &m)?)?;
    m.add_function(wrap_pyfunction!(build, &m)?)?;
    m.add_class::<Node>()?;
    m.setattr("__version__", env!("CARGO_PKG_VERSION"))?;
    m.setattr("__author__", "Dora-rs Authors")?;

    Ok(())
}
