use std::collections::HashMap;
use std::sync::Arc;

use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::PyList;

use ewal_core::eval::expr::{eval as rust_eval, eval_bool as rust_eval_bool, Expr};
use ewal_core::eval::find::{find as rust_find, parallel_find as rust_parallel_find, parallel_count as rust_parallel_count};
use ewal_core::eval::scope;
use ewal_core::signal::bitpack::{bit_slice, sign_extend};
use ewal_core::signal::SignalValue;
use ewal_core::trace::vcd::VcdTrace;
use ewal_core::trace::Trace;

/// Convert a SignalValue to a Python object.
fn signal_value_to_py(py: Python, val: &SignalValue) -> Py<PyAny> {
    match val {
        SignalValue::Scalar(v) => v.into_pyobject(py).unwrap().into_any().unbind(),
        SignalValue::WideScalar(v) => {
            v.first().copied().unwrap_or(0).into_pyobject(py).unwrap().into_any().unbind()
        }
        SignalValue::FourState(_) => py.None(),
        SignalValue::Real(f) => f.into_pyobject(py).unwrap().into_any().unbind(),
        SignalValue::Str(s) => s.into_pyobject(py).unwrap().into_any().unbind(),
        SignalValue::Unknown => py.None(),
    }
}

#[pyclass]
struct PyTrace {
    trace: Arc<dyn Trace>,
    name_to_id: HashMap<String, usize>,
}

#[pymethods]
impl PyTrace {
    #[staticmethod]
    fn from_vcd(path: &str) -> PyResult<Self> {
        let trace = VcdTrace::from_file(path)
            .map_err(|e| PyValueError::new_err(e))?;
        let mut name_to_id = HashMap::new();
        for name in trace.signal_names() {
            if let Some(id) = trace.signal_id(name) {
                name_to_id.insert(name.clone(), id);
            }
        }
        Ok(PyTrace {
            trace: Arc::new(trace),
            name_to_id,
        })
    }

    #[staticmethod]
    fn from_vcd_string(data: &str) -> PyResult<Self> {
        let trace = VcdTrace::parse(data.as_bytes())
            .map_err(|e| PyValueError::new_err(e))?;
        let mut name_to_id = HashMap::new();
        for name in trace.signal_names() {
            if let Some(id) = trace.signal_id(name) {
                name_to_id.insert(name.clone(), id);
            }
        }
        Ok(PyTrace {
            trace: Arc::new(trace),
            name_to_id,
        })
    }

    fn signal_names(&self) -> Vec<String> {
        self.trace.signal_names().to_vec()
    }

    fn scopes(&self) -> Vec<String> {
        self.trace.scopes().to_vec()
    }

    fn max_index(&self) -> usize {
        self.trace.max_index()
    }

    fn timestamp(&self, index: usize) -> u64 {
        self.trace.timestamp(index)
    }

    fn signal_value(&self, py: Python, name: &str, index: usize) -> Py<PyAny> {
        let val = self.trace.signal_value(name, index);
        signal_value_to_py(py, &val)
    }

    fn signal_value_raw(&self, name: &str, index: usize) -> u64 {
        self.trace.signal_value(name, index).to_u64().unwrap_or(0)
    }

    fn signal_width(&self, name: &str) -> u32 {
        self.trace.signal_width(name)
    }

    fn signal_id(&self, name: &str) -> Option<usize> {
        self.trace.signal_id(name)
    }

    fn signal_values_range(&self, py: Python, name: &str, start: usize, end: usize) -> Py<PyAny> {
        let vals = self.trace.signal_values_range(name, start, end);
        let list = PyList::new(py, vals.iter().map(|v| signal_value_to_py(py, v))).unwrap();
        list.into_any().unbind()
    }

    fn resolve(&self, scope_name: &str, short_name: &str) -> Option<String> {
        scope::resolve(self.trace.as_ref(), scope_name, short_name)
    }

    fn local_signals(&self, scope_name: &str) -> Vec<String> {
        scope::local_signals(self.trace.as_ref(), scope_name)
    }

    fn local_scopes(&self, scope_name: &str) -> Vec<String> {
        scope::local_scopes(self.trace.as_ref(), scope_name)
    }

    fn find(&self, expr: &PyExpr) -> Vec<usize> {
        rust_find(self.trace.as_ref(), &expr.inner)
    }

    /// Find matching indices using rayon parallel iteration across all CPU cores.
    fn parallel_find(&self, expr: &PyExpr) -> Vec<usize> {
        rust_parallel_find(self.trace.as_ref(), &expr.inner)
    }

    /// Count matching indices using rayon parallel iteration.
    fn parallel_count(&self, expr: &PyExpr) -> usize {
        rust_parallel_count(self.trace.as_ref(), &expr.inner)
    }

    fn eval_expr(&self, py: Python, expr: &PyExpr, index: usize) -> Py<PyAny> {
        let val = rust_eval(self.trace.as_ref(), &expr.inner, index);
        signal_value_to_py(py, &val)
    }

    fn eval_expr_bool(&self, expr: &PyExpr, index: usize) -> bool {
        rust_eval_bool(self.trace.as_ref(), &expr.inner, index)
    }

    fn bit_slice_val(&self, val: u64, high: u32, low: u32) -> u64 {
        bit_slice(val, high, low)
    }

    fn sign_extend_val(&self, val: u64, width: u32) -> i64 {
        sign_extend(val, width)
    }
}

#[pyclass(skip_from_py_object)]
#[derive(Clone)]
struct PyExpr {
    inner: Expr,
}

#[pymethods]
impl PyExpr {
    #[staticmethod]
    fn signal(id: usize) -> Self {
        PyExpr { inner: Expr::Signal(id) }
    }

    #[staticmethod]
    fn const_int(val: u64) -> Self {
        PyExpr { inner: Expr::Const(SignalValue::Scalar(val)) }
    }

    #[staticmethod]
    fn const_float(val: f64) -> Self {
        PyExpr { inner: Expr::Const(SignalValue::Real(val)) }
    }

    #[staticmethod]
    fn const_str(val: String) -> Self {
        PyExpr { inner: Expr::Const(SignalValue::Str(val)) }
    }

    #[staticmethod]
    fn always_true() -> Self {
        PyExpr { inner: Expr::True }
    }

    fn rel_eval(&self, offset: i64) -> Self {
        PyExpr { inner: Expr::RelEval(Box::new(self.inner.clone()), offset) }
    }

    fn not_(&self) -> Self {
        PyExpr { inner: Expr::Not(Box::new(self.inner.clone())) }
    }

    fn and_(&self, other: &PyExpr) -> Self {
        PyExpr { inner: Expr::And(Box::new(self.inner.clone()), Box::new(other.inner.clone())) }
    }

    fn or_(&self, other: &PyExpr) -> Self {
        PyExpr { inner: Expr::Or(Box::new(self.inner.clone()), Box::new(other.inner.clone())) }
    }

    fn eq_(&self, other: &PyExpr) -> Self {
        PyExpr { inner: Expr::Eq(Box::new(self.inner.clone()), Box::new(other.inner.clone())) }
    }

    fn neq(&self, other: &PyExpr) -> Self {
        PyExpr { inner: Expr::Neq(Box::new(self.inner.clone()), Box::new(other.inner.clone())) }
    }

    fn gt_(&self, other: &PyExpr) -> Self {
        PyExpr { inner: Expr::Gt(Box::new(self.inner.clone()), Box::new(other.inner.clone())) }
    }

    fn lt_(&self, other: &PyExpr) -> Self {
        PyExpr { inner: Expr::Lt(Box::new(self.inner.clone()), Box::new(other.inner.clone())) }
    }

    fn gte(&self, other: &PyExpr) -> Self {
        PyExpr { inner: Expr::Gte(Box::new(self.inner.clone()), Box::new(other.inner.clone())) }
    }

    fn lte(&self, other: &PyExpr) -> Self {
        PyExpr { inner: Expr::Lte(Box::new(self.inner.clone()), Box::new(other.inner.clone())) }
    }

    fn add(&self, other: &PyExpr) -> Self {
        PyExpr { inner: Expr::Add(Box::new(self.inner.clone()), Box::new(other.inner.clone())) }
    }

    fn sub(&self, other: &PyExpr) -> Self {
        PyExpr { inner: Expr::Sub(Box::new(self.inner.clone()), Box::new(other.inner.clone())) }
    }

    fn bit_and(&self, other: &PyExpr) -> Self {
        PyExpr { inner: Expr::BitAnd(Box::new(self.inner.clone()), Box::new(other.inner.clone())) }
    }

    fn bit_or(&self, other: &PyExpr) -> Self {
        PyExpr { inner: Expr::BitOr(Box::new(self.inner.clone()), Box::new(other.inner.clone())) }
    }

    fn bit_xor(&self, other: &PyExpr) -> Self {
        PyExpr { inner: Expr::BitXor(Box::new(self.inner.clone()), Box::new(other.inner.clone())) }
    }

    fn slice(&self, high: u32, low: u32) -> Self {
        PyExpr { inner: Expr::Slice(Box::new(self.inner.clone()), high, low) }
    }

    fn signed(&self) -> Self {
        PyExpr { inner: Expr::Signed(Box::new(self.inner.clone())) }
    }
}

#[pymodule]
fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyTrace>()?;
    m.add_class::<PyExpr>()?;
    Ok(())
}
