from functools import wraps
from typing import Callable

from .selection import Selection
from .state import EditorState, Result, StateChange


class StateMonad:
    def __init__(self, run_fn: Callable[[EditorState], Result]):
        self._run_fn = run_fn

    def run(self, state: EditorState) -> Result:
        return self._run_fn(state)

    def commit(self, state: EditorState) -> tuple[Selection, EditorState]:
        sel, changes = self.run(state)
        state.commit_changes(changes)
        return (sel, state)

    def bind(self, f: Callable[[Selection], "StateMonad"]) -> "StateMonad":
        def run_bound(state: EditorState) -> Result:
            sel1, changes1 = self.run(state)
            expr2 = f(sel1)
            sel2, changes2 = expr2.run(state)
            return (sel2, changes1 + changes2)

        return StateMonad(run_bound)

    @staticmethod
    def pure(selection: Selection) -> "StateMonad":
        def run_pure(_: EditorState) -> Result:
            return (selection, [])

        return StateMonad(run_pure)

    def modify(change: StateChange, selection: Selection) -> "StateMonad":
        def run_modify(_: EditorState) -> Result:
            return (selection, [change])

        return StateMonad(run_modify)


def monadic(
    func: Callable[[Selection, EditorState], Result],
) -> Callable[[Selection | StateMonad], StateMonad]:
    def make_monad(sel: Selection) -> StateMonad:
        def run_fn(state: EditorState) -> Result:
            return func(sel, state)

        return StateMonad(run_fn)

    @wraps(func)
    def wrapper(arg: Selection | StateMonad) -> StateMonad:
        if isinstance(arg, StateMonad):
            return arg.bind(make_monad)
        else:
            return make_monad(arg)

    return wrapper


def group(*exprs: StateMonad) -> StateMonad:
    def run_grouped(state: EditorState) -> Result:
        all_selections = []
        all_changes = []

        for expr in exprs:
            sel, changes = expr.run(state)
            all_selections.append(sel)
            all_changes.extend(changes)

        # TODO: Check for any overlapping changes
        # TODO: Merge selections into union

        return (
            all_selections[0] if all_selections else state.current_selection,
            all_changes,
        )

    return StateMonad(run_grouped)