import os
import re
import sys
import tempfile
import unicodedata
from collections.abc import Iterable
from itertools import islice
from pathlib import Path
from typing import TypeVar

import numpy as np
from numpy.typing import NDArray

from qwen3_embed.common.types import NumpyArray

T = TypeVar("T")


def last_token_pool(input_array: NumpyArray, attention_mask: NDArray[np.int64]) -> NumpyArray:
    """Extract embedding from the last non-padding token position.

    Qwen3-Embedding uses last-token pooling (NOT CLS/mean pooling).
    Handles both left-padding and right-padding.

    Args:
        input_array: Model output, shape (batch_size, seq_len, hidden_dim).
        attention_mask: Attention mask, shape (batch_size, seq_len).

    Returns:
        Pooled embeddings, shape (batch_size, hidden_dim).
    """
    left_padding = bool(attention_mask[:, -1].sum() == attention_mask.shape[0])
    if left_padding:
        return input_array[:, -1]

    sequence_lengths = attention_mask.sum(axis=1).astype(np.int64) - 1
    batch_size = input_array.shape[0]
    return input_array[np.arange(batch_size), sequence_lengths]


def normalize(input_array: NumpyArray, p: int = 2, dim: int = 1, eps: float = 1e-12) -> NumpyArray:
    # Calculate the Lp norm along the specified dimension
    norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
    norm = np.maximum(norm, eps)  # Avoid division by zero
    normalized_array = input_array / norm
    return normalized_array


def mean_pooling(input_array: NumpyArray, attention_mask: NDArray[np.int64]) -> NumpyArray:
    input_mask_expanded = np.expand_dims(attention_mask, axis=-1).astype(np.int64)
    input_mask_expanded = np.tile(input_mask_expanded, (1, 1, input_array.shape[-1]))
    sum_embeddings = np.sum(input_array * input_mask_expanded, axis=1)
    sum_mask = np.sum(input_mask_expanded, axis=1)
    pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
    return pooled_embeddings


def iter_batch[T](iterable: Iterable[T], size: int) -> Iterable[list[T]]:
    """
    >>> list(iter_batch([1,2,3,4,5], 3))
    [[1, 2, 3], [4, 5]]
    """
    source_iter = iter(iterable)
    while source_iter:
        b = list(islice(source_iter, size))
        if len(b) == 0:
            break
        yield b


def define_cache_dir(cache_dir: str | None = None) -> Path:
    """
    Define the cache directory for qwen3_embed
    """
    if cache_dir is None:
        default_cache_dir = os.path.join(tempfile.gettempdir(), "qwen3_embed_cache")
        cache_path = Path(os.getenv("QWEN3_EMBED_CACHE_PATH", default_cache_dir))
    else:
        cache_path = Path(cache_dir)
    cache_path.mkdir(parents=True, exist_ok=True)

    return cache_path


def get_all_punctuation() -> set[str]:
    return set(
        chr(i) for i in range(sys.maxunicode) if unicodedata.category(chr(i)).startswith("P")
    )


def remove_non_alphanumeric(text: str) -> str:
    return re.sub(r"[^\w\s]", " ", text, flags=re.UNICODE)
