use bazaar::textmerge;
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict, PyList, PyTuple};

type ExtractedLines = (Vec<Py<PyBytes>>, Vec<Vec<u8>>);

fn extract_byte_lines(seq: &Bound<PyAny>) -> PyResult<ExtractedLines> {
    let mut items = Vec::new();
    let mut keys = Vec::new();
    for item in seq.try_iter()? {
        let item = item?;
        let bytes = item
            .cast_into::<PyBytes>()
            .map_err(|_| pyo3::exceptions::PyTypeError::new_err("lines must be bytes"))?;
        keys.push(bytes.as_bytes().to_vec());
        items.push(bytes.unbind());
    }
    Ok((items, keys))
}

fn slice_pylist<'py>(
    py: Python<'py>,
    items: &[Py<PyBytes>],
    start: usize,
    end: usize,
) -> PyResult<Bound<'py, PyList>> {
    PyList::new(py, items[start..end].iter().map(|o| o.bind(py).clone()))
}

fn group_to_tuple<'py>(py: Python<'py>, group: &Group) -> PyResult<Bound<'py, PyTuple>> {
    match group {
        Group::Unchanged(lines) => PyTuple::new(py, [lines.bind(py).clone().into_any()]),
        Group::Conflict { a, b } => PyTuple::new(
            py,
            [a.bind(py).clone().into_any(), b.bind(py).clone().into_any()],
        ),
    }
}

/// A two-way merge group, with the inner line lists held as Python objects so
/// the original `bytes` instances round-trip back out unchanged.
enum Group {
    Unchanged(Py<PyList>),
    Conflict { a: Py<PyList>, b: Py<PyList> },
}

impl Group {
    fn is_useful(&self, py: Python) -> bool {
        match self {
            Group::Unchanged(lines) => lines.bind(py).len() > 0,
            Group::Conflict { a, b } => a.bind(py).len() > 0 || b.bind(py).len() > 0,
        }
    }
}

fn run_merge(
    py: Python,
    items_a: &[Py<PyBytes>],
    keys_a: &[Vec<u8>],
    items_b: &[Py<PyBytes>],
    keys_b: &[Vec<u8>],
) -> PyResult<Vec<Group>> {
    let raw = textmerge::merge2(keys_a, keys_b);
    let mut out = Vec::with_capacity(raw.len());
    let mut pa = 0usize;
    let mut pb = 0usize;
    for group in &raw {
        match group {
            textmerge::Group::Unchanged(lines) => {
                let len = lines.len();
                out.push(Group::Unchanged(
                    slice_pylist(py, items_a, pa, pa + len)?.unbind(),
                ));
                pa += len;
                pb += len;
            }
            textmerge::Group::Conflict { a, b } => {
                let la = a.len();
                let lb = b.len();
                out.push(Group::Conflict {
                    a: slice_pylist(py, items_a, pa, pa + la)?.unbind(),
                    b: slice_pylist(py, items_b, pb, pb + lb)?.unbind(),
                });
                pa += la;
                pb += lb;
            }
        }
    }
    Ok(out)
}

/// Two-way text merge.
///
/// Common regions are reported as one-element tuples; conflicts as two-element
/// tuples `(this_lines, other_lines)`.
#[pyclass(frozen, module = "bzrformats._bzr_rs.textmerge")]
struct Merge2 {
    items_a: Vec<Py<PyBytes>>,
    keys_a: Vec<Vec<u8>>,
    items_b: Vec<Py<PyBytes>>,
    keys_b: Vec<Vec<u8>>,
    a_marker: Py<PyBytes>,
    b_marker: Py<PyBytes>,
    split_marker: Py<PyBytes>,
}

#[pymethods]
impl Merge2 {
    #[classattr]
    #[allow(non_snake_case)]
    fn A_MARKER(py: Python) -> Bound<PyBytes> {
        PyBytes::new(py, textmerge::A_MARKER)
    }

    #[classattr]
    #[allow(non_snake_case)]
    fn B_MARKER(py: Python) -> Bound<PyBytes> {
        PyBytes::new(py, textmerge::B_MARKER)
    }

    #[classattr]
    #[allow(non_snake_case)]
    fn SPLIT_MARKER(py: Python) -> Bound<PyBytes> {
        PyBytes::new(py, textmerge::SPLIT_MARKER)
    }

    #[new]
    #[pyo3(signature = (lines_a, lines_b, a_marker = None, b_marker = None, split_marker = None))]
    fn new(
        py: Python,
        lines_a: &Bound<PyAny>,
        lines_b: &Bound<PyAny>,
        a_marker: Option<Py<PyBytes>>,
        b_marker: Option<Py<PyBytes>>,
        split_marker: Option<Py<PyBytes>>,
    ) -> PyResult<Self> {
        let (items_a, keys_a) = extract_byte_lines(lines_a)?;
        let (items_b, keys_b) = extract_byte_lines(lines_b)?;
        Ok(Self {
            items_a,
            keys_a,
            items_b,
            keys_b,
            a_marker: a_marker.unwrap_or_else(|| PyBytes::new(py, textmerge::A_MARKER).unbind()),
            b_marker: b_marker.unwrap_or_else(|| PyBytes::new(py, textmerge::B_MARKER).unbind()),
            split_marker: split_marker
                .unwrap_or_else(|| PyBytes::new(py, textmerge::SPLIT_MARKER).unbind()),
        })
    }

    #[getter]
    fn lines_a<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyList>> {
        slice_pylist(py, &self.items_a, 0, self.items_a.len())
    }

    #[getter]
    fn lines_b<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyList>> {
        slice_pylist(py, &self.items_b, 0, self.items_b.len())
    }

    #[getter]
    fn a_marker<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        self.a_marker.bind(py).clone()
    }

    #[getter]
    fn b_marker<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        self.b_marker.bind(py).clone()
    }

    #[getter]
    fn split_marker<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        self.split_marker.bind(py).clone()
    }

    /// Return raw structured merge info, without filtering empty groups.
    ///
    /// Each element is a tuple: length 1 for an unchanged region, length 2
    /// `(this, other)` for a conflict.
    fn _merge_struct<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyList>> {
        let groups = run_merge(py, &self.items_a, &self.keys_a, &self.items_b, &self.keys_b)?;
        let tuples: Vec<Bound<PyTuple>> = groups
            .iter()
            .map(|g| group_to_tuple(py, g))
            .collect::<PyResult<_>>()?;
        PyList::new(py, tuples)
    }

    /// Return structured merge info, with empty groups filtered out and
    /// optionally with conflict regions reduced via `reprocess_struct`.
    #[pyo3(signature = (reprocess = false))]
    fn merge_struct<'py>(&self, py: Python<'py>, reprocess: bool) -> PyResult<Bound<'py, PyList>> {
        let groups = run_merge(py, &self.items_a, &self.keys_a, &self.items_b, &self.keys_b)?;
        let useful: Vec<Group> = groups.into_iter().filter(|g| g.is_useful(py)).collect();
        let final_groups = if reprocess {
            reprocess_groups(py, useful)?
        } else {
            useful
        };
        let tuples: Vec<Bound<PyTuple>> = final_groups
            .iter()
            .map(|g| group_to_tuple(py, g))
            .collect::<PyResult<_>>()?;
        PyList::new(py, tuples)
    }

    /// Return `(merged_lines, had_conflicts)` where `merged_lines` is a list of
    /// byte lines, with conflict markers inserted around conflict regions.
    #[pyo3(signature = (reprocess = false))]
    fn merge_lines<'py>(
        &self,
        py: Python<'py>,
        reprocess: bool,
    ) -> PyResult<(Bound<'py, PyList>, bool)> {
        let groups = run_merge(py, &self.items_a, &self.keys_a, &self.items_b, &self.keys_b)?;
        let useful: Vec<Group> = groups.into_iter().filter(|g| g.is_useful(py)).collect();
        let final_groups = if reprocess {
            reprocess_groups(py, useful)?
        } else {
            useful
        };
        let (lines, conflicts) = render_lines(
            py,
            &final_groups,
            self.a_marker.bind(py),
            self.b_marker.bind(py),
            self.split_marker.bind(py),
        )?;
        Ok((lines, conflicts))
    }

    /// Filter empty groups out of a structured merge iterator.
    fn iter_useful(
        &self,
        py: Python<'_>,
        struct_iter: &Bound<'_, PyAny>,
    ) -> PyResult<Py<UsefulGroupsIterator>> {
        UsefulGroupsIterator::new(py, struct_iter)
    }

    /// Render structured merge info to a flat line list using this instance's
    /// conflict markers.
    fn struct_to_lines<'py>(
        &self,
        py: Python<'py>,
        struct_iter: &Bound<'py, PyAny>,
    ) -> PyResult<Bound<'py, PyList>> {
        let groups = iter_to_groups(struct_iter)?;
        let (lines, _) = render_lines(
            py,
            &groups,
            self.a_marker.bind(py),
            self.b_marker.bind(py),
            self.split_marker.bind(py),
        )?;
        Ok(lines)
    }

    /// Re-run a two-way merge over each conflict region, shrinking conflicts to
    /// their minimal diverging core.
    #[staticmethod]
    fn reprocess_struct<'py>(
        py: Python<'py>,
        struct_iter: &Bound<'py, PyAny>,
    ) -> PyResult<Bound<'py, PyList>> {
        let groups = iter_to_groups(struct_iter)?;
        let reprocessed = reprocess_groups(py, groups)?;
        let tuples: Vec<Bound<PyTuple>> = reprocessed
            .iter()
            .map(|g| group_to_tuple(py, g))
            .collect::<PyResult<_>>()?;
        PyList::new(py, tuples)
    }
}

/// Convert an iterable of `(lines,)` / `(lines_a, lines_b)` tuples into our
/// internal Group representation, retaining the original Python list objects.
fn iter_to_groups(struct_iter: &Bound<PyAny>) -> PyResult<Vec<Group>> {
    let py = struct_iter.py();
    let mut out = Vec::new();
    for item in struct_iter.try_iter()? {
        let tuple = item?.cast_into::<PyTuple>()?;
        let len = tuple.len();
        if len == 1 {
            let lines = tuple.get_item(0)?;
            let pylist = ensure_pylist(py, &lines)?;
            out.push(Group::Unchanged(pylist.unbind()));
        } else if len == 2 {
            let a = ensure_pylist(py, &tuple.get_item(0)?)?;
            let b = ensure_pylist(py, &tuple.get_item(1)?)?;
            out.push(Group::Conflict {
                a: a.unbind(),
                b: b.unbind(),
            });
        } else {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "merge struct tuples must have length 1 or 2",
            ));
        }
    }
    Ok(out)
}

fn ensure_pylist<'py>(py: Python<'py>, obj: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyList>> {
    if let Ok(list) = obj.cast::<PyList>() {
        Ok(list.clone())
    } else {
        let mut items = Vec::new();
        for item in obj.try_iter()? {
            items.push(item?);
        }
        PyList::new(py, items)
    }
}

fn reprocess_groups(py: Python, groups: Vec<Group>) -> PyResult<Vec<Group>> {
    let mut out = Vec::new();
    for group in groups {
        match group {
            Group::Unchanged(_) => out.push(group),
            Group::Conflict { a, b } => {
                let a_bound = a.bind(py);
                let b_bound = b.bind(py);
                let (items_a, keys_a) = extract_byte_lines(a_bound.as_any())?;
                let (items_b, keys_b) = extract_byte_lines(b_bound.as_any())?;
                let sub = run_merge(py, &items_a, &keys_a, &items_b, &keys_b)?;
                for g in sub.into_iter().filter(|g| g.is_useful(py)) {
                    out.push(g);
                }
            }
        }
    }
    Ok(out)
}

fn render_lines<'py>(
    py: Python<'py>,
    groups: &[Group],
    a_marker: &Bound<'py, PyBytes>,
    b_marker: &Bound<'py, PyBytes>,
    split_marker: &Bound<'py, PyBytes>,
) -> PyResult<(Bound<'py, PyList>, bool)> {
    let mut lines: Vec<Bound<PyAny>> = Vec::new();
    let mut conflicts = false;
    for group in groups {
        match group {
            Group::Unchanged(g) => {
                for item in g.bind(py).iter() {
                    lines.push(item);
                }
            }
            Group::Conflict { a, b } => {
                conflicts = true;
                lines.push(a_marker.clone().into_any());
                for item in a.bind(py).iter() {
                    lines.push(item);
                }
                lines.push(split_marker.clone().into_any());
                for item in b.bind(py).iter() {
                    lines.push(item);
                }
                lines.push(b_marker.clone().into_any());
            }
        }
    }
    Ok((PyList::new(py, lines)?, conflicts))
}

/// Extract `(state_str, line_obj)` tuples from an iterable, returning the
/// states (parsed), the line objects (unbound for re-emission) and the raw
/// line bytes (used for content comparisons inside the merge state machine).
fn extract_plan(
    plan: &Bound<PyAny>,
) -> PyResult<(Vec<textmerge::PlanState>, Vec<Py<PyAny>>, Vec<Vec<u8>>)> {
    let mut states = Vec::new();
    let mut lines = Vec::new();
    let mut line_bytes = Vec::new();
    for item in plan.try_iter()? {
        let pair = item?.cast_into::<PyTuple>()?;
        if pair.len() != 2 {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "plan items must be (state, line) pairs",
            ));
        }
        let state_str: String = pair.get_item(0)?.extract()?;
        let state = textmerge::PlanState::from_str(&state_str)
            .ok_or_else(|| pyo3::exceptions::PyAssertionError::new_err(state_str.clone()))?;
        states.push(state);
        let line_obj = pair.get_item(1)?;
        line_bytes.push(line_obj.extract::<Vec<u8>>().unwrap_or_default());
        lines.push(line_obj.unbind());
    }
    Ok((states, lines, line_bytes))
}

/// Translate a weave merge plan into structured merge groups (length-1 tuples
/// for resolved chunks, length-2 tuples for conflicts). Line objects from the
/// input plan are returned by reference, preserving identity.
#[pyfunction]
fn merge_struct_from_plan<'py>(
    py: Python<'py>,
    plan: &Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyList>> {
    let (states, lines, line_bytes) = extract_plan(plan)?;
    let groups = textmerge::merge_struct_from_plan(&states, &line_bytes);

    let mut tuples: Vec<Bound<PyTuple>> = Vec::with_capacity(groups.len());
    for group in groups {
        match group {
            textmerge::PlanGroup::Single(indices) => {
                let lst = PyList::new(py, indices.iter().map(|&i| lines[i].bind(py).clone()))?;
                tuples.push(PyTuple::new(py, [lst.into_any()])?);
            }
            textmerge::PlanGroup::Conflict { a, b } => {
                let la = PyList::new(py, a.iter().map(|&i| lines[i].bind(py).clone()))?;
                let lb = PyList::new(py, b.iter().map(|&i| lines[i].bind(py).clone()))?;
                tuples.push(PyTuple::new(py, [la.into_any(), lb.into_any()])?);
            }
        }
    }
    PyList::new(py, tuples)
}

/// Reconstruct a BASE text from a weave merge plan: emits the line objects for
/// `unchanged`, `killed-a`, `killed-b` and `killed-both` states.
#[pyfunction]
fn base_from_plan<'py>(py: Python<'py>, plan: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyList>> {
    let (states, lines, _line_bytes) = extract_plan(plan)?;
    let indices = textmerge::base_indices_from_plan(&states);
    PyList::new(py, indices.into_iter().map(|i| lines[i].bind(py).clone()))
}

/// Base class for text-mergers.
///
/// Subclasses must implement ``_merge_struct``. This is the Python-facing
/// base that `PlanWeaveMerge` and downstream `Merge3` subclass; the marker
/// bookkeeping (`struct_to_lines`, `iter_useful`, `merge_lines`,
/// `merge_struct`) is shared with `Merge2` via the module-level helpers.
///
/// `subclass` so Python subclasses can override `_merge_struct`; the
/// `merge_struct`/`merge_lines` methods dispatch back through
/// `self._merge_struct()` so the override is honoured.
#[pyclass(subclass, dict, module = "bzrformats._bzr_rs.textmerge")]
struct TextMerge;

#[pymethods]
impl TextMerge {
    #[classattr]
    #[allow(non_snake_case)]
    fn A_MARKER(py: Python) -> Bound<PyBytes> {
        PyBytes::new(py, textmerge::A_MARKER)
    }

    #[classattr]
    #[allow(non_snake_case)]
    fn B_MARKER(py: Python) -> Bound<PyBytes> {
        PyBytes::new(py, textmerge::B_MARKER)
    }

    #[classattr]
    #[allow(non_snake_case)]
    fn SPLIT_MARKER(py: Python) -> Bound<PyBytes> {
        PyBytes::new(py, textmerge::SPLIT_MARKER)
    }

    // `__new__` ignores its arguments so Python subclasses (e.g.
    // `PlanWeaveMerge`) can pass their own constructor args straight
    // through; the markers are set in `__init__`.
    #[new]
    #[pyo3(signature = (*_args, **_kwargs))]
    fn new(_args: Bound<'_, PyTuple>, _kwargs: Option<Bound<'_, PyDict>>) -> Self {
        TextMerge
    }

    /// Store the three conflict markers as instance attributes, defaulting
    /// to the standard markers when omitted.
    #[pyo3(signature = (a_marker = None, b_marker = None, split_marker = None))]
    fn __init__(
        slf: &Bound<'_, Self>,
        py: Python<'_>,
        a_marker: Option<Py<PyBytes>>,
        b_marker: Option<Py<PyBytes>>,
        split_marker: Option<Py<PyBytes>>,
    ) -> PyResult<()> {
        let a = a_marker.unwrap_or_else(|| PyBytes::new(py, textmerge::A_MARKER).unbind());
        let b = b_marker.unwrap_or_else(|| PyBytes::new(py, textmerge::B_MARKER).unbind());
        let split =
            split_marker.unwrap_or_else(|| PyBytes::new(py, textmerge::SPLIT_MARKER).unbind());
        slf.setattr("a_marker", a)?;
        slf.setattr("b_marker", b)?;
        slf.setattr("split_marker", split)?;
        Ok(())
    }

    /// Return structured merge info.  Must be implemented by subclasses.
    fn _merge_struct(&self) -> PyResult<()> {
        Err(pyo3::exceptions::PyNotImplementedError::new_err(
            "_merge_struct is abstract",
        ))
    }

    /// Render structured merge info to a flat line list using this instance's
    /// conflict markers.
    fn struct_to_lines<'py>(
        slf: &Bound<'py, Self>,
        py: Python<'py>,
        struct_iter: &Bound<'py, PyAny>,
    ) -> PyResult<Bound<'py, PyList>> {
        let groups = iter_to_groups(struct_iter)?;
        let (a, b, split) = markers(slf)?;
        let (lines, _) = render_lines(py, &groups, &a, &b, &split)?;
        Ok(lines)
    }

    /// Filter empty groups out of a structured merge iterator.
    fn iter_useful(
        &self,
        py: Python<'_>,
        struct_iter: &Bound<'_, PyAny>,
    ) -> PyResult<Py<UsefulGroupsIterator>> {
        UsefulGroupsIterator::new(py, struct_iter)
    }

    /// Produce `(merged_lines, had_conflicts)`. Dispatches through
    /// `self._merge_struct()` so a Python subclass override is honoured.
    #[pyo3(signature = (reprocess = false))]
    fn merge_lines<'py>(
        slf: &Bound<'py, Self>,
        py: Python<'py>,
        reprocess: bool,
    ) -> PyResult<(Bound<'py, PyList>, bool)> {
        let struct_iter = Self::merge_struct(slf, py, reprocess)?;
        let groups = iter_to_groups(struct_iter.as_any())?;
        let (a, b, split) = markers(slf)?;
        let (lines, conflicts) = render_lines(py, &groups, &a, &b, &split)?;
        Ok((lines, conflicts))
    }

    /// Produce structured merge info, filtering empty groups (and optionally
    /// reprocessing conflicts). Dispatches through `self._merge_struct()`.
    #[pyo3(signature = (reprocess = false))]
    fn merge_struct<'py>(
        slf: &Bound<'py, Self>,
        py: Python<'py>,
        reprocess: bool,
    ) -> PyResult<Bound<'py, PyList>> {
        let raw = slf.call_method0("_merge_struct")?;
        let useful = iter_useful_impl(py, &raw)?;
        if reprocess {
            let groups = iter_to_groups(useful.as_any())?;
            let reprocessed = reprocess_groups(py, groups)?;
            let tuples: Vec<Bound<PyTuple>> = reprocessed
                .iter()
                .map(|g| group_to_tuple(py, g))
                .collect::<PyResult<_>>()?;
            PyList::new(py, tuples)
        } else {
            Ok(useful)
        }
    }

    /// Re-run a two-way merge over each conflict region, shrinking conflicts to
    /// their minimal diverging core.
    #[staticmethod]
    fn reprocess_struct<'py>(
        py: Python<'py>,
        struct_iter: &Bound<'py, PyAny>,
    ) -> PyResult<Bound<'py, PyList>> {
        let groups = iter_to_groups(struct_iter)?;
        let reprocessed = reprocess_groups(py, groups)?;
        let tuples: Vec<Bound<PyTuple>> = reprocessed
            .iter()
            .map(|g| group_to_tuple(py, g))
            .collect::<PyResult<_>>()?;
        PyList::new(py, tuples)
    }
}

/// Read the three conflict-marker instance attributes off a `TextMerge`
/// (or subclass), which `__init__` set as `bytes`.
fn markers<'py>(
    slf: &Bound<'py, TextMerge>,
) -> PyResult<(
    Bound<'py, PyBytes>,
    Bound<'py, PyBytes>,
    Bound<'py, PyBytes>,
)> {
    let a = slf.getattr("a_marker")?.cast_into::<PyBytes>()?;
    let b = slf.getattr("b_marker")?.cast_into::<PyBytes>()?;
    let split = slf.getattr("split_marker")?.cast_into::<PyBytes>()?;
    Ok((a, b, split))
}

/// Shared `iter_useful` body: keep groups whose first list is non-empty, or
/// (for conflicts) whose second list is non-empty.
/// Lazy iterator returned by `iter_useful`. Pulls one group from the
/// source iterator per step and skips groups whose first (and, for
/// two-way groups, second) line list is empty. Mirrors the filtering
/// in `iter_useful_impl` without materialising the result.
#[pyclass]
struct UsefulGroupsIterator {
    source: Py<PyAny>,
}

impl UsefulGroupsIterator {
    fn new(py: Python<'_>, struct_iter: &Bound<'_, PyAny>) -> PyResult<Py<Self>> {
        Py::new(
            py,
            UsefulGroupsIterator {
                source: struct_iter.try_iter()?.into_any().unbind(),
            },
        )
    }
}

#[pymethods]
impl UsefulGroupsIterator {
    fn __iter__(slf: PyRef<Self>) -> PyRef<Self> {
        slf
    }

    fn __next__<'py>(&mut self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyTuple>>> {
        let source = self.source.bind(py);
        loop {
            let Some(item) = source.try_iter()?.next() else {
                return Ok(None);
            };
            let group = item?.cast_into::<PyTuple>()?;
            let len = group.len();
            let first = group.get_item(0)?;
            if first.try_iter()?.next().is_some() {
                return Ok(Some(group));
            }
            if len > 1 {
                let second = group.get_item(1)?;
                if second.try_iter()?.next().is_some() {
                    return Ok(Some(group));
                }
            }
        }
    }
}

fn iter_useful_impl<'py>(
    py: Python<'py>,
    struct_iter: &Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyList>> {
    let mut out: Vec<Bound<PyTuple>> = Vec::new();
    for item in struct_iter.try_iter()? {
        let group = item?.cast_into::<PyTuple>()?;
        let len = group.len();
        let first = group.get_item(0)?;
        if first.try_iter()?.next().is_some() {
            out.push(group);
            continue;
        }
        if len > 1 {
            let second = group.get_item(1)?;
            if second.try_iter()?.next().is_some() {
                out.push(group);
            }
        }
    }
    PyList::new(py, out)
}

/// Weave merge that takes a plan as input. Mirrors
/// `bzrformats.versionedfile.PlanWeaveMerge`. Extends the Rust `TextMerge`;
/// `_merge_struct`/`base_from_plan` drive the Rust plan helpers.
#[pyclass(
    extends = TextMerge,
    subclass,
    name = "PlanWeaveMerge",
    module = "bzrformats._bzr_rs.textmerge"
)]
struct PlanWeaveMerge {
    plan: Py<PyAny>,
}

#[pymethods]
impl PlanWeaveMerge {
    #[new]
    #[pyo3(signature = (plan, a_marker=None, b_marker=None))]
    fn new<'py>(
        py: Python<'py>,
        plan: Bound<'py, PyAny>,
        a_marker: Option<Py<PyBytes>>,
        b_marker: Option<Py<PyBytes>>,
    ) -> PyResult<PyClassInitializer<Self>> {
        let _ = (a_marker, b_marker);
        let plan_list = PyList::new(py, plan.try_iter()?.collect::<PyResult<Vec<_>>>()?)?;
        Ok(
            PyClassInitializer::from(TextMerge).add_subclass(PlanWeaveMerge {
                plan: plan_list.into_any().unbind(),
            }),
        )
    }

    /// TextMerge.__init__ stores the markers as instance attributes; mirror
    /// the Python `PlanWeaveMerge.__init__(plan, a_marker, b_marker)` which
    /// calls `TextMerge.__init__(self, a_marker, b_marker)` then `self.plan`.
    #[pyo3(signature = (plan, a_marker=None, b_marker=None))]
    fn __init__<'py>(
        slf: &Bound<'py, Self>,
        py: Python<'py>,
        plan: Bound<'py, PyAny>,
        a_marker: Option<Py<PyBytes>>,
        b_marker: Option<Py<PyBytes>>,
    ) -> PyResult<()> {
        let _ = plan;
        let a = a_marker.unwrap_or_else(|| PyBytes::new(py, textmerge::A_MARKER).unbind());
        let b = b_marker.unwrap_or_else(|| PyBytes::new(py, textmerge::B_MARKER).unbind());
        let split = PyBytes::new(py, textmerge::SPLIT_MARKER);
        slf.setattr("a_marker", a)?;
        slf.setattr("b_marker", b)?;
        slf.setattr("split_marker", split)?;
        Ok(())
    }

    #[getter]
    fn plan<'py>(&self, py: Python<'py>) -> Bound<'py, PyAny> {
        self.plan.bind(py).clone()
    }

    fn _merge_struct<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
        let lst = merge_struct_from_plan(py, self.plan.bind(py))?;
        lst.into_any().call_method0("__iter__")
    }

    fn base_from_plan<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyList>> {
        base_from_plan(py, self.plan.bind(py))
    }
}

/// Weave merge taking a `VersionedFile` and two versions. Mirrors
/// `bzrformats.versionedfile.WeaveMerge`. Extends `PlanWeaveMerge`.
#[pyclass(
    extends = PlanWeaveMerge,
    name = "WeaveMerge",
    module = "bzrformats._bzr_rs.textmerge"
)]
struct WeaveMerge;

#[pymethods]
impl WeaveMerge {
    #[new]
    #[pyo3(signature = (versionedfile, ver_a, ver_b, a_marker=None, b_marker=None))]
    fn new<'py>(
        py: Python<'py>,
        versionedfile: Bound<'py, PyAny>,
        ver_a: Bound<'py, PyAny>,
        ver_b: Bound<'py, PyAny>,
        a_marker: Option<Py<PyBytes>>,
        b_marker: Option<Py<PyBytes>>,
    ) -> PyResult<PyClassInitializer<Self>> {
        let _ = (a_marker, b_marker);
        let plan = versionedfile.call_method1("plan_merge", (ver_a, ver_b))?;
        let plan_list = PyList::new(py, plan.try_iter()?.collect::<PyResult<Vec<_>>>()?)?;
        Ok(PyClassInitializer::from(TextMerge)
            .add_subclass(PlanWeaveMerge {
                plan: plan_list.into_any().unbind(),
            })
            .add_subclass(WeaveMerge))
    }

    #[pyo3(signature = (versionedfile, ver_a, ver_b, a_marker=None, b_marker=None))]
    fn __init__<'py>(
        slf: &Bound<'py, Self>,
        py: Python<'py>,
        versionedfile: Bound<'py, PyAny>,
        ver_a: Bound<'py, PyAny>,
        ver_b: Bound<'py, PyAny>,
        a_marker: Option<Py<PyBytes>>,
        b_marker: Option<Py<PyBytes>>,
    ) -> PyResult<()> {
        let _ = (versionedfile, ver_a, ver_b);
        let a = a_marker.unwrap_or_else(|| PyBytes::new(py, textmerge::A_MARKER).unbind());
        let b = b_marker.unwrap_or_else(|| PyBytes::new(py, textmerge::B_MARKER).unbind());
        let split = PyBytes::new(py, textmerge::SPLIT_MARKER);
        slf.setattr("a_marker", a)?;
        slf.setattr("b_marker", b)?;
        slf.setattr("split_marker", split)?;
        Ok(())
    }
}

pub fn _textmerge_rs(py: Python) -> PyResult<Bound<PyModule>> {
    let m = PyModule::new(py, "textmerge")?;
    m.add_class::<Merge2>()?;
    m.add_class::<TextMerge>()?;
    m.add_class::<PlanWeaveMerge>()?;
    m.add_class::<WeaveMerge>()?;
    m.add("A_MARKER", PyBytes::new(py, textmerge::A_MARKER))?;
    m.add("B_MARKER", PyBytes::new(py, textmerge::B_MARKER))?;
    m.add("SPLIT_MARKER", PyBytes::new(py, textmerge::SPLIT_MARKER))?;
    m.add_function(pyo3::wrap_pyfunction!(merge_struct_from_plan, &m)?)?;
    m.add_function(pyo3::wrap_pyfunction!(base_from_plan, &m)?)?;
    Ok(m)
}
