from typing import Any, Callable, Iterable

from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn
from tqdm import tqdm


def tqdm_progress(
    iterable: Iterable,
    description: str,
    fn_get_status: Callable[[Any], str] | None = None,
    **tqdm_kwargs,
):
    """使用 tqdm 库显示进度条

    Parameters
    ----------
    - `iterable` : `Iterable`
        - 可迭代对象
    - `description` : `str`
        - 进度条描述
    - `fn_get_status` : `Callable[[Any], str]`, optional
        - 一个回调函数，输入当前迭代的元素，返回一个用于展示在进度条后缀的状态字符串, by default None
    """
    if fn_get_status is None:
        # 简单包装，保持原有行为
        return tqdm(iterable, desc=description, **tqdm_kwargs)

    # 需要动态更新状态的情况
    pbar = tqdm(iterable, desc=description, **tqdm_kwargs)
    for item in pbar:
        status_text = fn_get_status(item)
        pbar.set_postfix_str(status_text)
        yield item


def rich_progress(
    iterable: Iterable,
    description: str,
    color: str = 'green',
    fn_get_status: Callable[[Any], str] | None = None,
):
    """使用 rich 库显示一个美观的进度条

    Parameters
    ----------
    - `iterable` : `Iterable`
        - 可迭代对象
    - `description` : `str`
        - 进度条描述
    - `color` : `str`, optional
        - 进度条颜色, by default 'green'
    - `fn_get_status` : `Callable[[Any], str]`, optional
        - 一个回调函数，输入当前迭代的元素，返回一个用于展示在进度条旁边的状态字符串, by default None
    """
    try:
        total = len(iterable)  # pyright: ignore[reportArgumentType]
    except TypeError:
        total = None  # For iterables without a defined length (e.g., generators)

    progress_columns = [
        TextColumn('[progress.description]{task.description}'),
        BarColumn(),
        TaskProgressColumn(),
    ]
    if total is not None:
        progress_columns.append(TextColumn('{task.completed}/{task.total}'))
    else:
        # If total is None, just show completed count
        progress_columns.append(TextColumn('{task.completed}'))

    if fn_get_status:
        progress_columns.append(TextColumn('[cyan]({task.fields[status]})'))

    with Progress(*progress_columns, transient=True) as progress:
        task = progress.add_task(f'[{color}]{description}', total=total, status='')
        for item in iterable:
            if fn_get_status:
                status_text = fn_get_status(item)
                progress.update(task, status=status_text)

            yield item

            progress.advance(task)
