use std::sync::atomic::{AtomicU64, Ordering};

use autark::mpera::op::{BinaryOpKind, Literal, UnaryOpKind};
use autark::mpera::program::Program;
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::PyNone;

use crate::error::py_err;
use crate::kinds::{join_kind, reduce_kind, reduce_kind_with_x};
use crate::schema::PySchema;

static NEXT_PROGRAM_ID: AtomicU64 = AtomicU64::new(1);

#[pyclass(name = "Program", unsendable)]
#[derive(Clone)]
pub(crate) struct PyProgram {
    graph_id: u64,
    inner: Program,
}

impl PyProgram {
    pub(crate) fn new(graph_id: u64, inner: Program) -> PyProgram {
        PyProgram { graph_id, inner }
    }

    pub(crate) fn graph_id(&self) -> u64 {
        self.graph_id
    }

    pub(crate) fn inner_clone(&self) -> Program {
        self.inner.clone()
    }

    fn ensure_same(&self, other: &PyProgram) -> PyResult<()> {
        if self.graph_id != other.graph_id {
            return Err(PyValueError::new_err(
                "programs must share the same Program graph",
            ));
        }
        Ok(())
    }

    fn map<F>(&self, f: F) -> PyResult<PyProgram>
    where
        F: FnOnce(&Program) -> PyResult<Program>,
    {
        Ok(PyProgram::new(self.graph_id, f(&self.inner)?))
    }

    fn rhs(&self, rhs: &Bound<'_, PyAny>) -> PyResult<Program> {
        if let Ok(other) = rhs.extract::<PyRef<'_, PyProgram>>() {
            self.ensure_same(&other)?;
            return Ok(other.inner.clone());
        }
        if let Ok(value) = rhs.extract::<f64>() {
            return self.inner.lit(Literal::F64(value)).map_err(py_err);
        }
        if let Ok(value) = rhs.extract::<String>() {
            return self.inner.lit(Literal::String(value)).map_err(py_err);
        }
        Err(PyValueError::new_err("expected Program, number, or string"))
    }

    fn binary(&self, rhs: &Bound<'_, PyAny>, kind: BinaryOpKind) -> PyResult<PyProgram> {
        let rhs = self.rhs(rhs)?;
        self.map(|p| p.binaryop(rhs, kind).map_err(py_err))
    }

    fn reverse_scalar(&self, lhs: f64, kind: BinaryOpKind) -> PyResult<PyProgram> {
        let lhs = self.inner.lit(Literal::F64(lhs)).map_err(py_err)?;
        Ok(PyProgram::new(
            self.graph_id,
            lhs.binaryop(self.inner.clone(), kind).map_err(py_err)?,
        ))
    }
}

#[pymethods]
impl PyProgram {
    #[new]
    fn new_py() -> PyProgram {
        PyProgram::new(
            NEXT_PROGRAM_ID.fetch_add(1, Ordering::Relaxed),
            Program::new(),
        )
    }

    fn ir_json(&self) -> PyResult<String> {
        self.inner.ir().map_err(py_err)
    }

    fn lit(&self, value: &Bound<'_, PyAny>) -> PyResult<PyProgram> {
        if let Ok(v) = value.extract::<f64>() {
            return self.map(|p| p.lit(Literal::F64(v)).map_err(py_err));
        }
        if let Ok(v) = value.extract::<String>() {
            return self.map(|p| p.lit(Literal::String(v)).map_err(py_err));
        }
        Err(PyValueError::new_err("lit expects number or string"))
    }

    #[pyo3(signature = (names, index = None))]
    fn columns(&self, names: &Bound<'_, PyAny>, index: Option<usize>) -> PyResult<PyProgram> {
        if names.downcast::<PyNone>().is_ok() {
            return self.map(|p| {
                p.columns(
                    Some(index.unwrap_or(0)),
                    autark::mpera::op::ColumnSelector::All,
                )
                .map_err(py_err)
            });
        }
        if let Ok(column_index) = names.extract::<usize>() {
            return self.map(|p| {
                p.columns(
                    index,
                    autark::mpera::op::ColumnSelector::Indices(vec![column_index]),
                )
                .map_err(py_err)
            });
        }
        if let Ok(indices) = names.extract::<Vec<usize>>() {
            return self.map(|p| {
                p.columns(index, autark::mpera::op::ColumnSelector::Indices(indices))
                    .map_err(py_err)
            });
        }
        if let Ok(name) = names.extract::<String>() {
            return self.map(|p| {
                p.columns(index, autark::mpera::op::ColumnSelector::Names(vec![name]))
                    .map_err(py_err)
            });
        }
        let names = names.extract::<Vec<String>>()?;
        self.map(|p| {
            p.columns(index, autark::mpera::op::ColumnSelector::Names(names))
                .map_err(py_err)
        })
    }

    fn concat(&self, others: Vec<PyRef<'_, PyProgram>>) -> PyResult<PyProgram> {
        if others.is_empty() {
            return Err(PyValueError::new_err(
                "concat requires at least one program",
            ));
        }
        let mut programs = Vec::with_capacity(others.len());
        for other in others {
            self.ensure_same(&other)?;
            programs.push(other.inner.clone());
        }
        self.map(|p| p.concat(&programs).map_err(py_err))
    }

    fn slice(&self, start: Option<isize>, end: Option<isize>) -> PyResult<PyProgram> {
        self.map(|p| p.slice(start, end).map_err(py_err))
    }

    fn rolling(&self, n: usize) -> PyResult<PyProgram> {
        self.map(|p| p.rolling(n).map_err(py_err))
    }

    fn cumsum(&self) -> PyResult<PyProgram> {
        self.map(|p| p.cumsum().map_err(py_err))
    }

    #[pyo3(signature = (by = None, ascending = true))]
    fn order_by(&self, by: Option<PyRef<'_, PyProgram>>, ascending: bool) -> PyResult<PyProgram> {
        if let Some(by_ref) = &by {
            self.ensure_same(by_ref)?;
        }
        let by = by.map(|p| p.inner.clone());
        self.map(move |p| p.order_by(by, ascending).map_err(py_err))
    }

    #[pyo3(signature = (name, schema = None))]
    fn alias(&self, name: &str, schema: Option<PyRef<'_, PySchema>>) -> PyResult<PyProgram> {
        let schema = schema.map(|s| s.inner.clone());
        self.map(|p| p.alias(name, schema).map_err(py_err))
    }

    fn filter(&self, mask: PyRef<'_, PyProgram>) -> PyResult<PyProgram> {
        self.ensure_same(&mask)?;
        self.map(|p| p.filter(mask.inner.clone()).map_err(py_err))
    }

    #[pyo3(name = "where", signature = (mask, else_))]
    fn where_py(
        &self,
        mask: PyRef<'_, PyProgram>,
        else_: &Bound<'_, PyAny>,
    ) -> PyResult<PyProgram> {
        self.ensure_same(&mask)?;
        let else_ = self.rhs(else_)?;
        self.map(|p| p.r#where(mask.inner.clone(), else_).map_err(py_err))
    }

    #[pyo3(signature = (kind, x = None))]
    fn reduce(&self, kind: &str, x: Option<f64>) -> PyResult<PyProgram> {
        self.map(|p| p.reduce(reduce_kind_with_x(kind, x)?).map_err(py_err))
    }

    #[pyo3(signature = (kind, n = 1))]
    fn unaryop(&self, kind: &str, n: usize) -> PyResult<PyProgram> {
        let kind = match kind {
            "neg" => UnaryOpKind::Neg,
            "pct_change" => UnaryOpKind::PctChange { n },
            "diff" => UnaryOpKind::Diff { n },
            "lead" => UnaryOpKind::Lead { n },
            "lag" => UnaryOpKind::Lag { n },
            _ => {
                return Err(PyValueError::new_err(
                    "unaryop kind must be one of: neg, pct_change, diff, lead, lag",
                ));
            }
        };
        self.map(|p| p.unaryop(kind).map_err(py_err))
    }

    #[pyo3(signature = (keys = None))]
    fn group_by(&self, keys: Option<PyRef<'_, PyProgram>>) -> PyResult<PyProgram> {
        if let Some(keys_ref) = &keys {
            self.ensure_same(keys_ref)?;
        }
        let keys = keys.map(|p| p.inner.clone());
        self.map(move |p| p.group_by(keys).map_err(py_err))
    }

    #[pyo3(signature = (kind, keys = None))]
    fn group_reduce(
        &self,
        kind: &Bound<'_, PyAny>,
        keys: Option<PyRef<'_, PyProgram>>,
    ) -> PyResult<PyProgram> {
        if let Some(keys_ref) = &keys {
            self.ensure_same(keys_ref)?;
        }
        let keys = keys.map(|p| p.inner.clone());
        let kind = kind.as_any();
        let kinds = if let Ok(kind) = kind.extract::<String>() {
            vec![reduce_kind(&kind)?]
        } else if let Ok(kinds) = kind.extract::<Vec<String>>() {
            kinds
                .iter()
                .map(|kind| reduce_kind(kind))
                .collect::<PyResult<Vec<_>>>()?
        } else {
            return Err(PyValueError::new_err(
                "group_reduce kind must be a string or list of strings",
            ));
        };
        self.map(move |p| p.group_reduce(keys, kinds).map_err(py_err))
    }

    #[pyo3(signature = (right, left_on, right_on, kind = "inner"))]
    fn join(
        &self,
        right: PyRef<'_, PyProgram>,
        left_on: PyRef<'_, PyProgram>,
        right_on: PyRef<'_, PyProgram>,
        kind: &str,
    ) -> PyResult<PyProgram> {
        self.ensure_same(&right)?;
        self.ensure_same(&left_on)?;
        self.ensure_same(&right_on)?;
        self.map(|p| {
            p.join(
                right.inner.clone(),
                left_on.inner.clone(),
                right_on.inner.clone(),
                join_kind(kind)?,
            )
            .map_err(py_err)
        })
    }

    fn __add__(&self, rhs: &Bound<'_, PyAny>) -> PyResult<PyProgram> {
        self.binary(rhs, BinaryOpKind::Add)
    }

    fn __sub__(&self, rhs: &Bound<'_, PyAny>) -> PyResult<PyProgram> {
        self.binary(rhs, BinaryOpKind::Sub)
    }

    fn __mul__(&self, rhs: &Bound<'_, PyAny>) -> PyResult<PyProgram> {
        self.binary(rhs, BinaryOpKind::Mul)
    }

    fn __truediv__(&self, rhs: &Bound<'_, PyAny>) -> PyResult<PyProgram> {
        self.binary(rhs, BinaryOpKind::Div)
    }

    fn __lt__(&self, rhs: &Bound<'_, PyAny>) -> PyResult<PyProgram> {
        self.binary(rhs, BinaryOpKind::LesserThan)
    }

    fn __gt__(&self, rhs: &Bound<'_, PyAny>) -> PyResult<PyProgram> {
        self.binary(rhs, BinaryOpKind::GreaterThan)
    }

    fn __le__(&self, rhs: &Bound<'_, PyAny>) -> PyResult<PyProgram> {
        self.binary(rhs, BinaryOpKind::LesserEquals)
    }

    fn __ge__(&self, rhs: &Bound<'_, PyAny>) -> PyResult<PyProgram> {
        self.binary(rhs, BinaryOpKind::GreaterEquals)
    }

    fn __eq__(&self, rhs: &Bound<'_, PyAny>) -> PyResult<PyProgram> {
        self.binary(rhs, BinaryOpKind::Equals)
    }

    fn __ne__(&self, rhs: &Bound<'_, PyAny>) -> PyResult<PyProgram> {
        self.binary(rhs, BinaryOpKind::NotEquals)
    }

    fn __and__(&self, rhs: &Bound<'_, PyAny>) -> PyResult<PyProgram> {
        self.binary(rhs, BinaryOpKind::And)
    }

    fn __or__(&self, rhs: &Bound<'_, PyAny>) -> PyResult<PyProgram> {
        self.binary(rhs, BinaryOpKind::Or)
    }

    fn __radd__(&self, lhs: f64) -> PyResult<PyProgram> {
        self.reverse_scalar(lhs, BinaryOpKind::Add)
    }

    fn __rsub__(&self, lhs: f64) -> PyResult<PyProgram> {
        self.reverse_scalar(lhs, BinaryOpKind::Sub)
    }

    fn __rmul__(&self, lhs: f64) -> PyResult<PyProgram> {
        self.reverse_scalar(lhs, BinaryOpKind::Mul)
    }

    fn __rtruediv__(&self, lhs: f64) -> PyResult<PyProgram> {
        self.reverse_scalar(lhs, BinaryOpKind::Div)
    }
}
