#![allow(non_local_definitions)]
use ::pgqrs as rust_pgqrs;
use gethostname::gethostname;
#[cfg(feature = "s3")]
use pyo3::basic::CompareOp;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList};
#[cfg(feature = "s3")]
use rust_pgqrs::store::s3::S3Store as RustS3Store;
use rust_pgqrs::store::{AnyStore, Store};
use rust_pgqrs::{BackoffStrategy as RustBackoffStrategy, StepRetryPolicy as RustStepRetryPolicy};

use std::future::Future;
use std::sync::{Arc, LazyLock};
use tokio::runtime::Runtime;

mod builder;
mod r#enum;
mod errors;
mod tables;
mod workers;
mod workflow;

use builder::*;
use errors::*;
use r#enum::*;
use tables::*;
use workers::*;
use workflow::*;

pub(crate) struct RuntimeManager {
    pub(crate) rt: Runtime,
}

impl RuntimeManager {
    pub(crate) fn new() -> Self {
        Self {
            rt: Runtime::new().expect("Failed to initialize Tokio runtime"),
        }
    }

    pub fn block_on<F: Future>(&self, future: F) -> F::Output {
        self.rt.block_on(future)
    }
}

pub(crate) static RUNTIME: LazyLock<RuntimeManager> = LazyLock::new(RuntimeManager::new);

pub(crate) fn get_runtime() -> &'static RuntimeManager {
    &RUNTIME
}

// Helper to convert serde-json value to PyObject
pub(crate) fn json_to_py(py: Python, value: &serde_json::Value) -> PyResult<PyObject> {
    match value {
        serde_json::Value::Null => Ok(py.None()),
        serde_json::Value::Bool(b) => Ok(b.to_object(py)),
        serde_json::Value::Number(n) => {
            if let Some(i) = n.as_i64() {
                Ok(i.to_object(py))
            } else if let Some(f) = n.as_f64() {
                Ok(f.to_object(py))
            } else {
                Ok(n.to_string().to_object(py)) // Fallback
            }
        }
        serde_json::Value::String(s) => Ok(s.to_object(py)),
        serde_json::Value::Array(arr) => {
            let list = PyList::empty(py);
            for item in arr {
                list.append(json_to_py(py, item)?)?;
            }
            Ok(list.to_object(py))
        }
        serde_json::Value::Object(map) => {
            let dict = PyDict::new(py);
            for (k, v) in map {
                dict.set_item(k, json_to_py(py, v)?)?;
            }
            Ok(dict.to_object(py))
        }
    }
}

pub(crate) fn py_to_json(py: Python, val: &PyAny) -> PyResult<serde_json::Value> {
    let json_module = py.import("json")?;
    let json_str = json_module
        .call_method1("dumps", (val,))?
        .extract::<String>()?;
    serde_json::from_str(&json_str)
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))
}

#[pyclass(name = "Config")]
#[derive(Clone)]
pub struct PyConfig {
    #[allow(dead_code)]
    inner: rust_pgqrs::Config,
}

#[pymethods]
impl PyConfig {
    #[new]
    #[pyo3(signature = (dsn, schema=None, max_connections=None))]
    fn new(dsn: String, schema: Option<String>, max_connections: Option<u32>) -> PyResult<Self> {
        let mut inner = rust_pgqrs::Config::from_dsn(&dsn);
        if let Some(s) = schema {
            inner = inner.with_schema(&s);
        }
        if let Some(m) = max_connections {
            inner = inner.with_max_connections(m);
        }
        Ok(PyConfig { inner })
    }

    #[staticmethod]
    fn from_dsn(dsn: String) -> Self {
        PyConfig {
            inner: rust_pgqrs::Config::from_dsn(dsn),
        }
    }

    #[getter]
    fn get_dsn(&self) -> String {
        self.inner.dsn.clone()
    }

    #[getter]
    fn get_schema(&self) -> String {
        self.inner.schema.clone()
    }

    #[setter]
    fn set_schema(&mut self, schema: String) {
        self.inner.schema = schema;
    }

    #[getter]
    fn get_max_connections(&self) -> u32 {
        self.inner.max_connections
    }

    #[setter]
    fn set_max_connections(&mut self, max: u32) {
        self.inner.max_connections = max;
    }

    #[getter]
    fn get_connection_timeout_seconds(&self) -> u64 {
        self.inner.connection_timeout_seconds
    }

    #[setter]
    fn set_connection_timeout_seconds(&mut self, timeout: u64) {
        self.inner.connection_timeout_seconds = timeout;
    }

    #[getter]
    fn get_default_lock_time_seconds(&self) -> u32 {
        self.inner.default_lock_time_seconds
    }

    #[setter]
    fn set_default_lock_time_seconds(&mut self, lock_time: u32) {
        self.inner.default_lock_time_seconds = lock_time;
    }

    #[getter]
    fn get_default_max_batch_size(&self) -> usize {
        self.inner.default_max_batch_size
    }

    #[setter]
    fn set_default_max_batch_size(&mut self, batch_size: usize) {
        self.inner.default_max_batch_size = batch_size;
    }

    #[getter]
    fn get_max_read_ct(&self) -> i32 {
        self.inner.max_read_ct
    }

    #[setter]
    fn set_max_read_ct(&mut self, max_read_ct: i32) {
        self.inner.max_read_ct = max_read_ct;
    }

    #[getter]
    fn get_heartbeat_interval_seconds(&self) -> u64 {
        self.inner.heartbeat_interval
    }

    #[setter]
    fn set_heartbeat_interval_seconds(&mut self, interval: u64) {
        self.inner.heartbeat_interval = interval;
    }

    #[getter]
    fn get_max_payload_size_bytes(&self) -> usize {
        self.inner.validation_config.max_payload_size_bytes
    }

    #[setter]
    fn set_max_payload_size_bytes(&mut self, size: usize) {
        self.inner.validation_config.max_payload_size_bytes = size;
    }

    #[getter]
    fn get_max_string_length(&self) -> usize {
        self.inner.validation_config.max_string_length
    }

    #[setter]
    fn set_max_string_length(&mut self, length: usize) {
        self.inner.validation_config.max_string_length = length;
    }

    #[getter]
    fn get_max_object_depth(&self) -> usize {
        self.inner.validation_config.max_object_depth
    }

    #[setter]
    fn set_max_object_depth(&mut self, depth: usize) {
        self.inner.validation_config.max_object_depth = depth;
    }

    #[getter]
    fn get_forbidden_keys(&self) -> Vec<String> {
        self.inner.validation_config.forbidden_keys.clone()
    }

    #[setter]
    fn set_forbidden_keys(&mut self, keys: Vec<String>) {
        self.inner.validation_config.forbidden_keys = keys;
    }

    #[getter]
    fn get_required_keys(&self) -> Vec<String> {
        self.inner.validation_config.required_keys.clone()
    }

    #[setter]
    fn set_required_keys(&mut self, keys: Vec<String>) {
        self.inner.validation_config.required_keys = keys;
    }

    #[getter]
    fn get_max_enqueue_per_second(&self) -> Option<u32> {
        self.inner.validation_config.max_enqueue_per_second
    }

    #[setter]
    fn set_max_enqueue_per_second(&mut self, rate: Option<u32>) {
        self.inner.validation_config.max_enqueue_per_second = rate;
    }

    #[getter]
    fn get_max_enqueue_burst(&self) -> Option<u32> {
        self.inner.validation_config.max_enqueue_burst
    }

    #[setter]
    fn set_max_enqueue_burst(&mut self, burst: Option<u32>) {
        self.inner.validation_config.max_enqueue_burst = burst;
    }

    #[cfg(feature = "s3")]
    #[getter]
    fn get_s3_mode(&self) -> PyDurabilityMode {
        self.inner.s3.mode.into()
    }

    #[cfg(feature = "s3")]
    #[setter]
    fn set_s3_mode(&mut self, mode: PyDurabilityMode) {
        self.inner.s3.mode = mode.into();
    }
}

/// Backoff strategy for step retries
#[pyclass(name = "BackoffStrategy")]
#[derive(Clone)]
pub struct PyBackoffStrategy {
    inner: RustBackoffStrategy,
}

#[pymethods]
impl PyBackoffStrategy {
    /// Create a fixed delay backoff strategy
    #[staticmethod]
    fn fixed(delay_seconds: u32) -> Self {
        Self {
            inner: RustBackoffStrategy::Fixed { delay_seconds },
        }
    }

    /// Create an exponential backoff strategy
    #[staticmethod]
    fn exponential(base_seconds: u32, max_seconds: u32) -> Self {
        Self {
            inner: RustBackoffStrategy::Exponential {
                base_seconds,
                max_seconds,
            },
        }
    }

    /// Create an exponential backoff with jitter strategy
    #[staticmethod]
    fn exponential_with_jitter(base_seconds: u32, max_seconds: u32) -> Self {
        Self {
            inner: RustBackoffStrategy::ExponentialWithJitter {
                base_seconds,
                max_seconds,
            },
        }
    }

    fn __repr__(&self) -> String {
        format!("{:?}", self.inner)
    }
}

/// Retry policy for workflow steps
#[pyclass(name = "StepRetryPolicy")]
#[derive(Clone)]
pub struct PyStepRetryPolicy {
    inner: RustStepRetryPolicy,
}

#[pymethods]
impl PyStepRetryPolicy {
    #[new]
    #[pyo3(signature = (max_attempts=3, backoff=None))]
    fn new(max_attempts: u32, backoff: Option<PyBackoffStrategy>) -> Self {
        let backoff_inner =
            backoff
                .map(|b| b.inner)
                .unwrap_or(RustBackoffStrategy::ExponentialWithJitter {
                    base_seconds: 1,
                    max_seconds: 60,
                });

        Self {
            inner: RustStepRetryPolicy {
                max_attempts,
                backoff: backoff_inner,
            },
        }
    }

    #[getter]
    fn max_attempts(&self) -> u32 {
        self.inner.max_attempts
    }

    fn __repr__(&self) -> String {
        format!(
            "StepRetryPolicy(max_attempts={}, backoff={:?})",
            self.inner.max_attempts, self.inner.backoff
        )
    }
}

#[pyclass(name = "Store")]
#[derive(Clone)]
pub struct PyStore {
    pub(crate) inner: AnyStore,
}

#[pymethods]
impl PyStore {
    fn producer<'a>(&self, py: Python<'a>, queue: String) -> PyResult<&'a PyAny> {
        let store = self.inner.clone();
        pyo3_asyncio::tokio::future_into_py(py, async move {
            let hostname = format!(
                "{}-{}",
                gethostname().to_string_lossy(),
                uuid::Uuid::new_v4()
            );
            // Use Store trait method directly - returns Box<dyn Producer + 'static>
            let producer = store
                .producer(&queue, &hostname, 0, store.config())
                .await
                .map_err(to_py_err)?;

            Ok(PyProducer {
                inner: Arc::new(producer),
            })
        })
    }

    fn consumer<'a>(&self, py: Python<'a>, queue: String) -> PyResult<&'a PyAny> {
        let store = self.inner.clone();
        pyo3_asyncio::tokio::future_into_py(py, async move {
            let hostname = format!(
                "{}-{}",
                gethostname().to_string_lossy(),
                uuid::Uuid::new_v4()
            );
            // Use Store trait method directly - returns Box<dyn Consumer + 'static>
            let consumer = store
                .consumer(&queue, &hostname, 0, store.config())
                .await
                .map_err(to_py_err)?;

            Ok(PyConsumer {
                inner: Arc::new(consumer),
            })
        })
    }

    fn consume_iter(&self, queue: String, poll_interval_ms: Option<u64>) -> PyConsumerIterator {
        PyConsumerIterator::new(self.inner.clone(), queue, poll_interval_ms.unwrap_or(50))
    }

    fn queue<'a>(&self, py: Python<'a>, name: String) -> PyResult<&'a PyAny> {
        let store = self.inner.clone();
        pyo3_asyncio::tokio::future_into_py(py, async move {
            let q = store.queue(&name).await.map_err(to_py_err)?;
            Ok(PyQueueInfo::from(q))
        })
    }

    fn get_workers<'a>(&self, py: Python<'a>) -> PyResult<&'a PyAny> {
        let store = self.inner.clone();
        pyo3_asyncio::tokio::future_into_py(py, async move { Ok(PyWorkers { store }) })
    }

    fn get_queues<'a>(&self, py: Python<'a>) -> PyResult<&'a PyAny> {
        let store = self.inner.clone();
        pyo3_asyncio::tokio::future_into_py(py, async move { Ok(PyQueues { store }) })
    }

    fn get_messages<'a>(&self, py: Python<'a>) -> PyResult<&'a PyAny> {
        let store = self.inner.clone();
        pyo3_asyncio::tokio::future_into_py(py, async move { Ok(PyMessages { store }) })
    }

    fn get_workflows<'a>(&self, py: Python<'a>) -> PyResult<&'a PyAny> {
        let store = self.inner.clone();
        pyo3_asyncio::tokio::future_into_py(py, async move { Ok(PyWorkflows { store }) })
    }

    fn get_workflow_runs<'a>(&self, py: Python<'a>) -> PyResult<&'a PyAny> {
        let store = self.inner.clone();
        pyo3_asyncio::tokio::future_into_py(py, async move { Ok(PyWorkflowRuns { store }) })
    }

    fn get_workflow_steps<'a>(&self, py: Python<'a>) -> PyResult<&'a PyAny> {
        let store = self.inner.clone();
        pyo3_asyncio::tokio::future_into_py(py, async move { Ok(PyWorkflowSteps { store }) })
    }

    fn bootstrap<'a>(&self, py: Python<'a>) -> PyResult<&'a PyAny> {
        let store = self.inner.clone();
        pyo3_asyncio::tokio::future_into_py(py, async move {
            store.bootstrap().await.map_err(to_py_err)
        })
    }
}

#[cfg(feature = "s3")]
#[pyclass(name = "S3StoreHandle")]
#[derive(Clone)]
pub struct PyS3StoreHandle {
    pub(crate) inner: RustS3Store,
}

#[cfg(feature = "s3")]
#[pymethods]
impl PyS3StoreHandle {
    fn snapshot<'a>(&self, py: Python<'a>) -> PyResult<&'a PyAny> {
        let mut store = self.inner.clone();
        pyo3_asyncio::tokio::future_into_py(
            py,
            async move { store.snapshot().await.map_err(to_py_err) },
        )
    }

    fn sync<'a>(&self, py: Python<'a>) -> PyResult<&'a PyAny> {
        let mut store = self.inner.clone();
        pyo3_asyncio::tokio::future_into_py(
            py,
            async move { store.sync().await.map_err(to_py_err) },
        )
    }
}

#[pyfunction]
fn connect<'a>(py: Python<'a>, dsn: String) -> PyResult<&'a PyAny> {
    pyo3_asyncio::tokio::future_into_py(py, async move {
        let config = rust_pgqrs::Config::from_dsn(&dsn);
        let store = rust_pgqrs::connect_with_config(&config)
            .await
            .map_err(to_py_err)?;
        Ok(PyStore { inner: store })
    })
}

#[pyfunction]
fn connect_with<'a>(py: Python<'a>, config: PyConfig) -> PyResult<&'a PyAny> {
    pyo3_asyncio::tokio::future_into_py(py, async move {
        let store = rust_pgqrs::connect_with_config(&config.inner)
            .await
            .map_err(to_py_err)?;
        Ok(PyStore { inner: store })
    })
}

#[cfg(feature = "s3")]
#[pyfunction]
fn as_s3(store: &PyStore) -> PyResult<PyS3StoreHandle> {
    match &store.inner {
        AnyStore::S3(store) => Ok(PyS3StoreHandle {
            inner: store.clone(),
        }),
        _ => Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(format!(
            "store backend is '{}', expected 's3'",
            store.inner.backend_name()
        ))),
    }
}

#[pyclass(name = "Admin")]
pub struct PyAdmin {
    pub(crate) store: AnyStore,
}

#[pymethods]
impl PyAdmin {
    #[new]
    fn new(store: PyStore) -> Self {
        PyAdmin { store: store.inner }
    }

    fn install<'a>(&self, py: Python<'a>) -> PyResult<&'a PyAny> {
        let store = self.store.clone();
        pyo3_asyncio::tokio::future_into_py(py, async move {
            store.bootstrap().await.map_err(to_py_err)
        })
    }

    fn verify<'a>(&self, py: Python<'a>) -> PyResult<&'a PyAny> {
        let store = self.store.clone();
        pyo3_asyncio::tokio::future_into_py(py, async move {
            rust_pgqrs::admin(&store).verify().await.map_err(to_py_err)
        })
    }

    fn delete_queue<'a>(&self, py: Python<'a>, name: String) -> PyResult<&'a PyAny> {
        let store = self.store.clone();
        pyo3_asyncio::tokio::future_into_py(py, async move {
            rust_pgqrs::admin(&store)
                .delete_queue_by_name(&name)
                .await
                .map_err(to_py_err)?;
            Ok(true)
        })
    }

    fn release_worker_messages<'a>(&self, py: Python<'a>, worker_id: i64) -> PyResult<&'a PyAny> {
        let store = self.store.clone();
        pyo3_asyncio::tokio::future_into_py(py, async move {
            let released = rust_pgqrs::admin(&store)
                .release_worker_messages(worker_id)
                .await
                .map_err(to_py_err)?;
            Ok(released)
        })
    }

    #[pyo3(signature = (name, arg=None))]
    fn create_workflow<'a>(
        &self,
        py: Python<'a>,
        name: String,
        arg: Option<PyObject>,
    ) -> PyResult<&'a PyAny> {
        let store = self.store.clone();
        let json_arg = if let Some(a) = arg {
            Some(py_to_json(py, a.as_ref(py))?)
        } else {
            None
        };

        pyo3_asyncio::tokio::future_into_py(py, async move {
            // Ensure workflow exists
            let _wf_rec = store.workflow(&name).await.map_err(to_py_err)?;

            if let Some(input) = json_arg {
                // If arg is provided, trigger and return a run handle
                let producer = store
                    .producer_ephemeral(&name, store.config())
                    .await
                    .map_err(to_py_err)?;

                let payload = serde_json::json!({ "input": input });
                let msgs = producer
                    .batch_enqueue(&[payload])
                    .await
                    .map_err(to_py_err)?;
                let run_msg = msgs[0].clone();

                let workflow = store.run(run_msg).await.map_err(to_py_err)?;
                let workflow_id = workflow.id();

                Ok(PyRun {
                    inner: Arc::new(workflow),
                    store,
                    workflow_id,
                })
            } else {
                // Return dummy result for non-triggering create
                Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(
                    "create_workflow in Python must have an arg to trigger a run",
                ))
            }
        })
    }
}

#[pyfunction]
fn admin(store: PyStore) -> PyAdmin {
    PyAdmin { store: store.inner }
}

#[pymodule]
fn _pgqrs(py: Python, m: &PyModule) -> PyResult<()> {
    #[cfg(feature = "s3")]
    m.add_class::<PyDurabilityMode>()?;
    m.add_class::<PyAdmin>()?;
    m.add_class::<PyProducer>()?;
    m.add_class::<PyConsumer>()?;
    m.add_class::<PyWorkers>()?;
    m.add_class::<PyQueues>()?;
    m.add_class::<PyMessages>()?;
    m.add_class::<PyWorkflows>()?;
    m.add_class::<PyWorkflowRuns>()?;
    m.add_class::<PyWorkflowSteps>()?;
    m.add_class::<PyRunRecord>()?;
    m.add_class::<PyStepRecord>()?;
    m.add_class::<PyConfig>()?;
    m.add_class::<PyWorkerStatus>()?;
    m.add_class::<PyWorkflowStatus>()?;
    m.add_class::<PyStepResultStatus>()?;
    m.add_class::<PyStore>()?;
    #[cfg(feature = "s3")]
    m.add_class::<PyS3StoreHandle>()?;
    m.add_class::<PyQueueMessage>()?;
    m.add_class::<PyQueueInfo>()?;
    m.add_class::<PyRun>()?;
    m.add_class::<PyStepResult>()?;
    m.add_class::<PyStepGuard>()?;
    m.add_class::<PyConsumerIterator>()?;
    m.add_class::<PyBackoffStrategy>()?;
    m.add_class::<PyStepRetryPolicy>()?;
    m.add_class::<WorkflowRecord>()?;
    m.add_class::<PyWorkflowBuilder>()?;
    m.add_class::<PyWorkflowTriggerBuilder>()?;
    m.add_class::<PyRunBuilder>()?;
    m.add_class::<PyStepBuilder>()?;
    m.add_class::<PyDequeueBuilder>()?;

    add_exceptions(py, m)?;

    m.add_function(wrap_pyfunction!(connect, m)?)?;
    m.add_function(wrap_pyfunction!(connect_with, m)?)?;
    #[cfg(feature = "s3")]
    m.add_function(wrap_pyfunction!(as_s3, m)?)?;
    m.add_function(wrap_pyfunction!(admin, m)?)?;
    m.add_function(wrap_pyfunction!(produce, m)?)?;
    m.add_function(wrap_pyfunction!(produce_batch, m)?)?;
    m.add_function(wrap_pyfunction!(consume, m)?)?;
    m.add_function(wrap_pyfunction!(consume_batch, m)?)?;
    m.add_function(wrap_pyfunction!(enqueue, m)?)?;
    m.add_function(wrap_pyfunction!(enqueue_batch, m)?)?;
    m.add_function(wrap_pyfunction!(dequeue, m)?)?;
    m.add_function(wrap_pyfunction!(dequeue_builder, m)?)?;
    m.add_function(wrap_pyfunction!(archive, m)?)?;
    m.add_function(wrap_pyfunction!(archive_batch, m)?)?;
    m.add_function(wrap_pyfunction!(delete, m)?)?;
    m.add_function(wrap_pyfunction!(enqueue_delayed, m)?)?;
    m.add_function(wrap_pyfunction!(extend_vt, m)?)?;
    m.add_function(wrap_pyfunction!(workflow::workflow, m)?)?;
    m.add_function(wrap_pyfunction!(workflow::run, m)?)?;
    m.add_function(wrap_pyfunction!(workflow::step, m)?)?;

    Ok(())
}
