use crate::{PyIndexSet, PyModel};
use arco_core::Model;
use arco_core::model::{PrettyBoundGroup, PrettyPrintAdapter, PrettySection, format_ascii_number};
use arco_core::types::Bounds;
use arco_expr::VariableId;
use pyo3::prelude::*;
use std::collections::HashSet;
use std::fmt::Write as _;

const FLOAT_EQ_EPSILON: f64 = 1e-12;

pub(crate) struct ArrayPrintSpec {
    pub(crate) start_var_id: u32,
    pub(crate) len: usize,
    pub(crate) base_name: String,
    pub(crate) shape: Vec<usize>,
    pub(crate) strides: Vec<usize>,
    pub(crate) index_sets: Vec<Py<PyIndexSet>>,
}

pub(crate) struct PythonPrettyAdapter<'a> {
    pub(crate) model: &'a PyModel,
}

impl PrettyPrintAdapter for PythonPrettyAdapter<'_> {
    fn variable_label(&self, model: &Model, var_id: VariableId) -> Option<String> {
        let spec = self.model.find_array_print_spec(var_id)?;
        let array_label = PyModel::array_label_for_var(spec, var_id)?;
        if let Some(name) = model.get_variable_name(var_id) {
            if !PyModel::is_autogenerated_array_name(spec, var_id, name) {
                return Some(name.to_string());
            }
        }
        Some(array_label)
    }

    fn sections(&self, _model: &Model) -> Vec<PrettySection> {
        let lines = self.model.render_index_set_lines();
        if lines.is_empty() {
            Vec::new()
        } else {
            vec![PrettySection {
                heading: "Index sets".to_string(),
                entries: lines,
            }]
        }
    }

    fn grouped_bounds(&self, _model: &Model) -> Vec<PrettyBoundGroup> {
        self.model.render_grouped_array_bounds()
    }
}

impl ArrayPrintSpec {
    fn offset_of(&self, var_id: VariableId) -> Option<usize> {
        let raw = var_id.inner();
        if raw < self.start_var_id {
            return None;
        }
        let offset = (raw - self.start_var_id) as usize;
        (offset < self.len).then_some(offset)
    }
}

impl PyModel {
    pub(crate) fn register_array_print_spec(
        &mut self,
        py: Python<'_>,
        start_var_id: u32,
        total: usize,
        index_sets: &[Py<PyIndexSet>],
        shape: &[usize],
        base_name: Option<&str>,
    ) {
        if total == 0 || index_sets.is_empty() {
            return;
        }
        let strides = (0..shape.len())
            .map(|axis| shape[axis + 1..].iter().product::<usize>().max(1))
            .collect::<Vec<_>>();

        self.array_print_specs.push(ArrayPrintSpec {
            start_var_id,
            len: total,
            base_name: base_name.unwrap_or("x").to_string(),
            shape: shape.to_vec(),
            strides,
            index_sets: index_sets.iter().map(|set| set.clone_ref(py)).collect(),
        });
    }

    pub(crate) fn find_array_print_spec(&self, var_id: VariableId) -> Option<&ArrayPrintSpec> {
        self.array_print_specs
            .iter()
            .find(|spec| spec.offset_of(var_id).is_some())
    }

    pub(crate) fn array_label_for_var(spec: &ArrayPrintSpec, var_id: VariableId) -> Option<String> {
        let offset = spec.offset_of(var_id)?;
        if spec.shape.is_empty() {
            return Some(spec.base_name.clone());
        }
        if spec.index_sets.len() != spec.shape.len() {
            return None;
        }

        Python::attach(|py| {
            let mut parts = Vec::with_capacity(spec.shape.len());
            for axis in 0..spec.shape.len() {
                let stride = *spec.strides.get(axis)?;
                let size = *spec.shape.get(axis)?;
                if size == 0 {
                    return None;
                }
                let coord = (offset / stride) % size;
                let set_ref = spec.index_sets.get(axis)?.borrow(py);
                let member = set_ref.members.get(coord)?;
                parts.push(format_index_member(member));
            }
            Some(format!("{}[{}]", spec.base_name, parts.join(",")))
        })
    }

    pub(crate) fn is_autogenerated_array_name(
        spec: &ArrayPrintSpec,
        var_id: VariableId,
        name: &str,
    ) -> bool {
        let Some(offset) = spec.offset_of(var_id) else {
            return false;
        };
        if spec.len == 1 {
            return name == spec.base_name || name == format!("{}[0]", spec.base_name);
        }
        name == format!("{}[{offset}]", spec.base_name)
    }

    fn render_index_set_lines(&self) -> Vec<String> {
        if self.array_print_specs.is_empty() {
            return Vec::new();
        }
        Python::attach(|py| {
            let mut seen = HashSet::new();
            let mut lines = Vec::new();
            for spec in &self.array_print_specs {
                for index_set in &spec.index_sets {
                    let set_ref = index_set.borrow(py);
                    let members = set_ref
                        .members
                        .iter()
                        .map(format_index_member)
                        .collect::<Vec<_>>()
                        .join(", ");
                    let line = format!("{} = [{}]", set_ref.name, members);
                    if seen.insert(line.clone()) {
                        lines.push(line);
                    }
                }
            }
            lines
        })
    }

    fn render_grouped_array_bounds(&self) -> Vec<PrettyBoundGroup> {
        let mut groups = Vec::new();

        for spec in &self.array_print_specs {
            if spec.len <= 1 || spec.index_sets.len() != spec.shape.len() {
                continue;
            }

            let first_id = VariableId::new(spec.start_var_id);
            let Ok(first_var) = self.inner.get_variable(first_id) else {
                continue;
            };
            if is_binary_variable(first_var.bounds, first_var.is_integer) {
                continue;
            }

            let uniform = (1..spec.len).all(|offset| {
                let var_id = VariableId::new(spec.start_var_id + offset as u32);
                if let Ok(var) = self.inner.get_variable(var_id) {
                    var.is_integer == first_var.is_integer
                        && float_approx_equal(var.bounds.lower, first_var.bounds.lower)
                        && float_approx_equal(var.bounds.upper, first_var.bounds.upper)
                } else {
                    false
                }
            });
            if !uniform {
                continue;
            }

            let Some(bounds_label) = Self::array_bounds_label(spec) else {
                continue;
            };
            let Some(mut line) = format_variable_bounds_line(&bounds_label, first_var.bounds)
            else {
                continue;
            };
            if let Some(quantifier) = Self::array_bounds_quantifier(spec) {
                let _ = write!(line, "  {quantifier}");
            }
            let vars = (0..spec.len)
                .map(|offset| VariableId::new(spec.start_var_id + offset as u32))
                .collect();
            groups.push(PrettyBoundGroup { text: line, vars });
        }

        groups
    }

    fn array_bounds_label(spec: &ArrayPrintSpec) -> Option<String> {
        if spec.shape.is_empty() {
            return Some(spec.base_name.clone());
        }
        if spec.index_sets.len() != spec.shape.len() {
            return None;
        }
        Python::attach(|py| {
            let mut axis_symbols = Vec::with_capacity(spec.index_sets.len());
            for (axis, index_set) in spec.index_sets.iter().enumerate() {
                let set_ref = index_set.borrow(py);
                axis_symbols.push(axis_symbol_from_name(&set_ref.name, axis));
            }
            Some(format!("{}[{}]", spec.base_name, axis_symbols.join(",")))
        })
    }

    fn array_bounds_quantifier(spec: &ArrayPrintSpec) -> Option<String> {
        if spec.shape.is_empty() || spec.index_sets.len() != spec.shape.len() {
            return None;
        }
        Python::attach(|py| {
            let mut parts = Vec::with_capacity(spec.index_sets.len());
            for (axis, index_set) in spec.index_sets.iter().enumerate() {
                let set_ref = index_set.borrow(py);
                let symbol = axis_symbol_from_name(&set_ref.name, axis);
                parts.push(format!("{symbol} in {}", set_ref.name));
            }
            Some(format!("for {}", parts.join(", ")))
        })
    }
}

fn float_approx_equal(lhs: f64, rhs: f64) -> bool {
    if lhs.to_bits() == rhs.to_bits() {
        return true;
    }
    if !lhs.is_finite() || !rhs.is_finite() {
        return false;
    }
    let scale = lhs.abs().max(rhs.abs()).max(1.0);
    (lhs - rhs).abs() <= FLOAT_EQ_EPSILON * scale
}

fn is_binary_variable(bounds: Bounds, is_integer: bool) -> bool {
    is_integer && float_approx_equal(bounds.lower, 0.0) && float_approx_equal(bounds.upper, 1.0)
}

fn format_variable_bounds_line(label: &str, bounds: Bounds) -> Option<String> {
    let lower_finite = bounds.lower.is_finite();
    let upper_finite = bounds.upper.is_finite();
    if !lower_finite && !upper_finite {
        return None;
    }

    if lower_finite && upper_finite {
        return Some(format!(
            "{} <= {label} <= {}",
            format_ascii_number(bounds.lower),
            format_ascii_number(bounds.upper)
        ));
    }
    if lower_finite {
        return Some(format!("{} <= {label}", format_ascii_number(bounds.lower)));
    }
    Some(format!("{label} <= {}", format_ascii_number(bounds.upper)))
}

fn format_index_member(member: &crate::index_set::IndexMember) -> String {
    match member {
        crate::index_set::IndexMember::Int(value) => value.to_string(),
        crate::index_set::IndexMember::Float(value) => format_ascii_number(*value),
        crate::index_set::IndexMember::Str(value) => value.clone(),
    }
}

fn axis_symbol_from_name(name: &str, axis: usize) -> String {
    if let Some(ch) = name.chars().find(|ch| ch.is_ascii_alphabetic()) {
        return ch.to_ascii_lowercase().to_string();
    }
    format!("i{}", axis + 1)
}
