aimbat.utils

def delete_sampledata() -> None:
    """Delete sample data."""

    logger.info(f"Deleting sample data in {settings.sampledata_dir}.")

    shutil.rmtree(settings.sampledata_dir)

def download_sampledata(force: bool = False) -> None:
    """Download sample data."""

    logger.info(
        f"Downloading sample data from {settings.sampledata_src} to {settings.sampledata_dir}."
    )

    if (
        settings.sampledata_dir.exists()
        and len(os.listdir(settings.sampledata_dir)) != 0
    ):
        if force is True:
            delete_sampledata()
        else:
            raise FileExistsError(
                f"The directory {settings.sampledata_dir} already exists and is non-empty."
            )

    with urlopen(settings.sampledata_src) as zipresp:
        with ZipFile(BytesIO(zipresp.read())) as zfile:
            zfile.extractall(settings.sampledata_dir)

def json_to_table(
    data: dict[str, Any] | list[dict[str, Any]],
    title: str | None = None,
    formatters: dict[str, Callable[[Any], str]] | None = None,
    skip_keys: list[str] | None = None,
    column_order: list[str] | None = None,
    column_kwargs: dict[str, dict[str, Any]] | None = None,
    common_column_kwargs: dict[str, Any] | None = None,
) -> None:
    """Print a JSON dict or list of dicts as a rich table.

    For a single dict the table has ``Key`` and ``Value`` columns with one row
    per key-value pair.  For a list of dicts the keys become column headers and
    each list item becomes a row.

    Args:
        data: A single JSON dict or a list of JSON dicts.
        title: Optional title displayed above the table.
        formatters: Optional mapping of key names to callables that receive the
            raw value and return a string for display.
        skip_keys: Optional list of keys to exclude from the table.
        column_order: Optional list of keys defining the display order.  Keys
            not listed are appended after in their original order.  For a single
            dict this controls row order; for a list of dicts it controls column
            order.
        column_kwargs: Optional mapping of key names to keyword arguments
            forwarded to ``Table.add_column`` (e.g. ``style``, ``justify``,
            ``min_width``).  A ``"header"`` entry overrides the displayed column
            header name.  For a single dict the special keys ``"Key"`` and
            ``"Value"`` target those header columns.
        common_column_kwargs: Optional keyword arguments applied to every
            column, merged with any per-column entries in ``column_kwargs``.
            Per-column values take precedence over these defaults.

    Examples:
        >>> json_to_table({"name": "Alice", "age": 30}, title="Person")
        >>> json_to_table([{"id": 1}, {"id": 2}], formatters={"id": str})
        >>> json_to_table({"name": "Alice", "secret": "x"}, skip_keys=["secret"])
        >>> json_to_table(
        ...     [{"id": 1, "name": "Alice"}],
        ...     column_order=["name", "id"],
        ...     column_kwargs={"id": {"justify": "right", "style": "cyan"}},
        ... )
    """
    formatters = formatters or {}
    skip = set(skip_keys or [])
    column_kwargs = column_kwargs or {}
    common_column_kwargs = common_column_kwargs or {}
    console = Console()
    table = make_table(title=title)

    def _sorted_keys(keys: list[str]) -> list[str]:
        """Return keys reordered by column_order, with remaining keys appended."""
        if not column_order:
            return keys
        ordered = [k for k in column_order if k in keys]
        rest = [k for k in keys if k not in set(column_order)]
        return ordered + rest

    if isinstance(data, dict):
        key_kw = {**common_column_kwargs, **column_kwargs.get("Key", {})}
        val_kw = {**common_column_kwargs, **column_kwargs.get("Value", {})}
        table.add_column(key_kw.pop("header", "Key"), **key_kw)
        table.add_column(val_kw.pop("header", "Value"), **val_kw)
        keys = _sorted_keys([k for k in data if k not in skip])
        for key in keys:
            formatted = (
                formatters[key](data[key]) if key in formatters else str(data[key])
            )
            table.add_row(str(key), formatted)
    else:
        if not data:
            console.print(table)
            return
        columns = _sorted_keys([k for k in data[0].keys() if k not in skip])
        for col in columns:
            col_kw = {**common_column_kwargs, **column_kwargs.get(col, {})}
            table.add_column(col_kw.pop("header", str(col)), **col_kw)
        for item in data:
            row = []
            for col in columns:
                value = item.get(col)
                formatted = formatters[col](value) if col in formatters else str(value)
                row.append(formatted)
            table.add_row(*row)

    console.print(table)

def string_to_uuid(
    session: Session,
    id: str,
    aimbat_class: type[AimbatTypes],
    custom_error: str | None = None,
) -> UUID:
    """Determine a UUID from a string containing the first few characters.

    Args:
        session: Database session.
        id: Input string to find UUID for.
        aimbat_class: Aimbat class to use to find UUID.
        custom_error: Overrides the default error message.

    Returns:
        The full UUID.

    Raises:
        ValueError: If the UUID could not be determined.
    """
    statement = select(aimbat_class.id).where(
        func.replace(cast(aimbat_class.id, String), "-", "").like(
            f"{id.replace('-', '')}%"
        )
    )
    uuid_set = set(session.exec(statement).all())
    if len(uuid_set) == 1:
        return uuid_set.pop()
    if len(uuid_set) == 0:
        raise ValueError(
            custom_error or f"Unable to find {aimbat_class.__name__} using id: {id}."
        )
    raise ValueError(f"Found more than one {aimbat_class.__name__} using id: {id}")

def uuid_shortener[T: AimbatTypes](
    session: Session,
    aimbat_obj: T | type[T],
    min_length: int = 2,
    str_uuid: str | None = None,
) -> str:
    """Calculates the shortest unique prefix for a UUID, returning with dashes.

    Args:
        session: An active SQLModel/SQLAlchemy session.
        aimbat_obj: Either an instance of a SQLModel or the SQLModel class itself.
        min_length: The starting character length for the shortened ID.
        str_uuid: The full UUID string. Required only if `aimbat_obj` is a class.

    Returns:
        str: The shortest unique prefix string, including hyphens where applicable.
    """

    if isinstance(aimbat_obj, type):
        model_class = aimbat_obj
        if str_uuid is None:
            raise ValueError("str_uuid must be provided when aimbat_obj is a class.")
        target_full = str(UUID(str_uuid))
    else:
        model_class = type(aimbat_obj)
        target_full = str(aimbat_obj.id)

    prefix_clean = target_full.replace("-", "")[:min_length]

    # select with a WHERE clause that removes dashes and compares the cleaned prefix
    statement = select(model_class.id).where(
        func.replace(cast(model_class.id, String), "-", "").like(f"{prefix_clean}%")
    )

    # Store results as standard hyphenated strings
    results = session.exec(statement).all()
    relevant_pool = [str(uid) for uid in results]

    if target_full not in relevant_pool:
        raise ValueError(f"ID {target_full} not found in table {model_class.__name__}")

    current_length = min_length
    while current_length < len(target_full):
        candidate = target_full[:current_length]
        if candidate.endswith("-"):
            current_length += 1
            candidate = target_full[:current_length]

        matches = [u for u in relevant_pool if u.startswith(candidate)]
        if len(matches) == 1:
            return candidate
        current_length += 1

    return target_full