from functools import reduce
from typing import Literal

import pandas as pd


def merge_many(
    dfs: list[pd.DataFrame],
    on: list[str],
    how: str = "inner",
) -> pd.DataFrame:
    return reduce(lambda df1, df2: df1.merge(df2, on=on, how=how), dfs)


def merge(
    left: pd.DataFrame,
    right: pd.DataFrame,
    on: list[str] | None = None,
    how: Literal["left", "inner", "right", "outer"] = "inner",
    suffixes: tuple[str, str] = ("_x", "_y"),
):
    """Merge DataFrames as an equi-join.

    Args:
        left (pd.DataFrame): Left Dataframe
        right (pd.DataFrame): Right DataFrame
        on (list[str]): Columns to join on. Columns must be present in both DataFrames.
        how (JoinType, optional): Join type ("inner", "left", "right", "outer"
            or "cross"). Defaults to "inner".
        suffixes (tuple[str, str], optional): Suffixes to use for colliding column
            names. Defaults to ("_x", "_y").
    """
    if on is None:
        on = []
    # pandas has improved merge performance in 3, but we can
    # enable it for pandas 2 using mode.copy_on_write option
    with pd.option_context("mode.copy_on_write", True):
        if len(on) == 0:
            return pd.merge(
                left,
                right,
                how="cross",
                suffixes=suffixes,
            )

        return pd.merge(
            left,
            right,
            on=on,
            how=how,
            suffixes=suffixes,
        )


def align(
    left: pd.DataFrame,
    right: pd.DataFrame,
    left_on: list[str] | None = None,
    right_on: list[str] | None = None,
    on: list[str] | None = None,
    how: str = "inner",
):
    if on:
        if left_on is not None or right_on is not None:
            # raise same error as pd.merge
            raise pd.errors.MergeError(
                'Can only pass argument "on" OR "left_on" '
                'and "right_on", not a combination of both.'
            )
        left_on = on
        right_on = on

    left = left.set_index(left_on)
    right = right.set_index(right_on)

    aligned_left, aligned_right = left.align(right, join=how, axis="index")

    aligned_left = aligned_left.reset_index()
    aligned_right = aligned_right.reset_index()

    return aligned_left, aligned_right
