src/utils/helpers.py

"""
霜云(Shimokumo) - 辅助函数模块

提供文本处理、时间格式化、文件操作等通用工具函数。
"""

import json
import os
import re
import uuid
from datetime import datetime, timedelta
from typing import Any, Dict, List, Optional, Union


def format_timestamp(
    dt: Optional[datetime] = None,
    fmt: str = "%Y-%m-%d %H:%M:%S",
) -> str:
    """
    格式化时间戳。

    Args:
        dt: datetime对象，为None则使用当前时间
        fmt: 格式化字符串

    Returns:
        格式化后的时间字符串

    Example:
        >>> format_timestamp()
        '2026-07-06 12:00:00'
    """
    if dt is None:
        dt = datetime.now()
    return dt.strftime(fmt)


def format_duration(seconds: float) -> str:
    """
    将秒数格式化为人类可读的时长字符串。

    Args:
        seconds: 秒数

    Returns:
        格式化的时长字符串，如 "2小时30分15秒"
    """
    if seconds < 0:
        return "0秒"

    seconds = int(seconds)
    days, seconds = divmod(seconds, 86400)
    hours, seconds = divmod(seconds, 3600)
    minutes, seconds = divmod(seconds, 60)

    parts: List[str] = []
    if days > 0:
        parts.append(f"{days}天")
    if hours > 0:
        parts.append(f"{hours}小时")
    if minutes > 0:
        parts.append(f"{minutes}分")
    if seconds > 0 or not parts:
        parts.append(f"{seconds}秒")

    return "".join(parts)


def truncate_text(text: str, max_length: int = 500, suffix: str = "...") -> str:
    """
    截断文本到指定长度。

    Args:
        text: 原始文本
        max_length: 最大长度
        suffix: 截断后缀

    Returns:
        截断后的文本

    Example:
        >>> truncate_text("这是一段很长的文本", max_length=5, suffix="...")
        '这是一段很...'
    """
    if len(text) <= max_length:
        return text
    return text[: max_length - len(suffix)] + suffix


def count_tokens_estimate(text: str) -> int:
    """
    粗略估算文本的token数量。
    对于中文，大约每个字符对应1-2个token；
    对于英文，大约每4个字符对应1个token。

    Args:
        text: 输入文本

    Returns:
        估算的token数量
    """
    if not text:
        return 0

    chinese_chars = len(re.findall(r"[\u4e00-\u9fff]", text))
    other_chars = len(text) - chinese_chars

    # 中文约1.5 token/字，英文约0.25 token/字符
    return int(chinese_chars * 1.5 + other_chars * 0.25)


def safe_json_loads(text: str, default: Any = None) -> Any:
    """
    安全地解析JSON字符串。

    Args:
        text: JSON字符串
        default: 解析失败时返回的默认值

    Returns:
        解析后的Python对象，或默认值
    """
    try:
        return json.loads(text)
    except (json.JSONDecodeError, TypeError, ValueError):
        return default


def safe_json_dumps(obj: Any, indent: int = 2, ensure_ascii: bool = False) -> str:
    """
    安全地将Python对象序列化为JSON字符串。

    Args:
        obj: 要序列化的对象
        indent: 缩进空格数
        ensure_ascii: 是否转义非ASCII字符

    Returns:
        JSON字符串
    """
    try:
        return json.dumps(obj, indent=indent, ensure_ascii=ensure_ascii)
    except (TypeError, ValueError) as e:
        return f"{{\"error\": \"JSON序列化失败: {str(e)}\"}}"


def ensure_directory(path: str) -> str:
    """
    确保目录存在，不存在则创建。

    Args:
        path: 目录路径

    Returns:
        目录的绝对路径
    """
    abs_path = os.path.abspath(path)
    os.makedirs(abs_path, exist_ok=True)
    return abs_path


def read_file_content(file_path: str, encoding: str = "utf-8") -> Optional[str]:
    """
    读取文件内容。

    Args:
        file_path: 文件路径
        encoding: 文件编码

    Returns:
        文件内容字符串，失败返回None
    """
    try:
        with open(file_path, "r", encoding=encoding) as f:
            return f.read()
    except (IOError, OSError) as e:
        return None


def write_file_content(
    file_path: str,
    content: str,
    encoding: str = "utf-8",
) -> bool:
    """
    写入文件内容。

    Args:
        file_path: 文件路径
        content: 要写入的内容
        encoding: 文件编码

    Returns:
        是否写入成功
    """
    try:
        dir_path = os.path.dirname(file_path)
        if dir_path:
            os.makedirs(dir_path, exist_ok=True)
        with open(file_path, "w", encoding=encoding) as f:
            f.write(content)
        return True
    except (IOError, OSError) as e:
        return False


def generate_unique_id(prefix: str = "") -> str:
    """
    生成唯一ID。

    Args:
        prefix: ID前缀

    Returns:
        带前缀的唯一ID字符串

    Example:
        >>> generate_unique_id("chat")
        'chat_3fa85f64-5717-4562-b3fc-2c963f66afa6'
    """
    uid = str(uuid.uuid4())[:8]
    if prefix:
        return f"{prefix}_{uid}"
    return uid


def clean_html_text(html: str) -> str:
    """
    清理HTML标签，提取纯文本。

    Args:
        html: HTML字符串

    Returns:
        清理后的纯文本
    """
    # 移除script和style标签及其内容
    text = re.sub(r"<script[^>]*>[\s\S]*?</script>", "", html, flags=re.IGNORECASE)
    text = re.sub(r"<style[^>]*>[\s\S]*?</style>", "", text, flags=re.IGNORECASE)
    # 移除所有HTML标签
    text = re.sub(r"<[^>]+>", "", text)
    # 解码HTML实体
    text = text.replace("&nbsp;", " ")
    text = text.replace("&lt;", "<")
    text = text.replace("&gt;", ">")
    text = text.replace("&amp;", "&")
    text = text.replace("&quot;", '"')
    text = text.replace("&#39;", "'")
    # 合并多余空白
    text = re.sub(r"\s+", " ", text).strip()
    return text


def extract_urls(text: str) -> List[str]:
    """
    从文本中提取URL列表。

    Args:
        text: 输入文本

    Returns:
        URL列表
    """
    url_pattern = r"https?://[^\s<>\"]+|[a-zA-Z0-9-]+\.[a-zA-Z]{2,}[^\s<>\"]*"
    return re.findall(url_pattern, text)


def chunk_text(
    text: str,
    chunk_size: int = 1000,
    overlap: int = 100,
) -> List[str]:
    """
    将长文本按段落分块。

    Args:
        text: 输入文本
        chunk_size: 每块最大字符数
        overlap: 块间重叠字符数

    Returns:
        文本块列表
    """
    if len(text) <= chunk_size:
        return [text]

    chunks: List[str] = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        # 尝试在句号、换行等处断句
        if end < len(text):
            for sep in ["\n", "。", ".", "！", "!", "？", "?"]:
                sep_pos = text.rfind(sep, start + chunk_size // 2, end)
                if sep_pos > start:
                    end = sep_pos + 1
                    break
        chunks.append(text[start:end])
        start = end - overlap
    return chunks


def calculate_similarity(text1: str, text2: str) -> float:
    """
    计算两段文本的简单相似度（基于字符重叠）。

    Args:
        text1: 文本1
        text2: 文本2

    Returns:
        相似度分数 (0.0 ~ 1.0)
    """
    if not text1 or not text2:
        return 0.0

    set1 = set(text1.lower())
    set2 = set(text2.lower())
    intersection = set1 & set2
    union = set1 | set2

    if not union:
        return 1.0

    return len(intersection) / len(union)


def is_media_url(url: str) -> Dict[str, Any]:
    """
    检测URL是否为媒体资源URL，并识别媒体类型。

    Args:
        url: 要检测的URL

    Returns:
        包含媒体信息的字典，如 {"is_media": True, "type": "video", "format": "mp4"}
    """
    media_extensions = {
        "video": {
            "mp4", "avi", "mkv", "mov", "wmv", "flv", "webm", "m4v", "ts", "3gp"
        },
        "audio": {
            "mp3", "wav", "flac", "aac", "ogg", "wma", "m4a", "opus", "ape", "alac"
        },
        "image": {
            "jpg", "jpeg", "png", "gif", "webp", "bmp", "svg", "ico", "tiff"
        },
    }

    # 提取文件扩展名
    ext_match = re.search(r"\.(\w+)(?:\?|$)", url.lower())
    if not ext_match:
        return {"is_media": False}

    ext = ext_match.group(1)

    for media_type, extensions in media_extensions.items():
        if ext in extensions:
            return {
                "is_media": True,
                "type": media_type,
                "format": ext,
            }

    return {"is_media": False}