import semchunk
import tiktoken
from typing import Union, List
from rich.progress import Progress, SpinnerColumn, TextColumn


def chunker(
    inputs: Union[str, List[str]],
    chunk_size: int = 512,
    model: str = "gpt-4",
    processes: int = 1,
    memoize: bool = True,
    progress: bool = False,
    max_token_chars: int = None,
    progress_bar: bool = True
) -> Union[List[str], List[List[str]]]:
    """
    Takes a string or a list of strings and returns the chunked content.

    Example:
        ```python
        chunk("Hello, world!")
        # ["Hello, world!"]
        ```

    Args:
        inputs: Union[str, List[str]]: The input to chunk.
        chunk_size: int: The size of the chunks to return.
        model: str: The model to use for chunking.
        processes: int: The number of processes to use for chunking.
        memoize: bool: Whether to memoize the chunking process.
        progress: bool: Whether to show a progress bar.
        max_token_chars: int: The maximum number of characters to use for chunking.

    Returns:
        Union[List[str], List[List[str]]]: The chunked content.
    """

    try:
        tokenizer = tiktoken.encoding_for_model(model)


        if progress_bar:
            with Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                transient=True
            ) as progress:
                task_id = progress.add_task("Chunking...", total=None)

                chunker = semchunk.chunkerify(
                    tokenizer,
                    chunk_size=chunk_size,
                    max_token_chars=max_token_chars,
                    memoize=memoize,
                )
                for i, chunk in enumerate(chunker(inputs)):
                    progress.update(task_id, completed=i + 1)

        else:
            chunker = semchunk.chunkerify(
                tokenizer,
                chunk_size=chunk_size,
                max_token_chars=max_token_chars,
                memoize=memoize,
            )

        # Handle single input case (str)
        if isinstance(inputs, str):
            inputs = [inputs]  # Convert to list for uniform handling

        if not isinstance(inputs, list):
            raise TypeError(
                "inputs must be a string or a list of strings"
            )

        texts = []
        for item in inputs:
            # Handle string input directly
            if isinstance(item, str):
                texts.append(item)
            else:
                raise TypeError(f"Unsupported input type: {type(item)}")

        # Chunk the content, using processes and progress bar as needed
        if len(texts) == 1:
            return chunker(texts[0])  # Single input, return the chunked result
        else:

            if progress_bar:
                progress.update(task_id, completed=1)

            return chunker(
                texts, processes=processes, progress=progress
            )  # Multiple inputs

    except Exception as e:
        # Detailed error logging
        print(f"Error in chunk function: {str(e)}")
        raise e
    

if __name__ == "__main__":
    print(chunker("Hello, world!", progress_bar=True))
