"""帮助OPENAI API 推理有关的函数"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../src/notebooks/04_llm_api.ipynb.

# %% auto 0
__all__ = [
    "make_info_df",
    "get_pydantic_version",
    "is_url",
    "is_local_file_path",
    "download_file",
    "local_video_to_base64_uri",
    "separate_think_and_other",
    "get_env_bool",
    "Endpoint",
    "flatten_dict",
]

# %% ../src/notebooks/04_llm_api.ipynb 4
import pandas as pd


def make_info_df(df, n_samples=2):
    """
    构造包含 column_name, dtype 和前 n_samples 行样本值的表格
    """
    info_df = pd.DataFrame({"column_name": df.columns, "dtype": df.dtypes.values})
    # 添加前 n_samples 行样本值
    for i in range(n_samples):
        col = f"sample{i + 1}"
        info_df[col] = df.iloc[i].values
    return info_df


# %% ../src/notebooks/04_llm_api.ipynb 5
import pydantic


def get_pydantic_version():
    """
    判断当前使用的是 Pydantic 1 还是 Pydantic 2
    :return: 版本号 1 或 2
    """
    try:
        if pydantic.VERSION.startswith("1"):
            return 1
        elif pydantic.VERSION.startswith("2"):
            return 2
        else:
            raise ValueError(f"不支持的 Pydantic 版本: {pydantic.VERSION}")
    except AttributeError:
        # Pydantic 1 可能没有 VERSION 属性，通过其他方式判断
        try:
            from pydantic.main import ModelMetaclass

            return 1
        except ImportError:
            return 2


# %% ../src/notebooks/04_llm_api.ipynb 7
import re
from urllib.parse import urlparse
import os


def is_url(s: str) -> bool:
    """判断是否为URL（包含http/https/ftp/file等协议）"""
    # 1. 语法校验：协议前缀 + ://（file:// 也符合）
    url_pattern = r"^[a-zA-Z][a-zA-Z0-9+.-]*://[^\s]*$"
    if not re.match(url_pattern, s.strip()):
        return False
    # 2. 解析协议，确认是合法URL协议
    parsed = urlparse(s)
    valid_schemes = {"http", "https", "ftp", "ftps", "sftp", "ssh", "telnet", "file"}
    return parsed.scheme.lower() in valid_schemes


def is_local_file_path(s: str) -> bool:
    """判断是否为本地文件路径（非URL格式）"""
    # 排除URL（避免与file://混淆）
    if is_url(s):
        return False
    # 检查是否符合操作系统的文件路径格式
    # 简化判断：包含路径分隔符，或符合盘符（Windows）/根目录（Linux/macOS）特征
    s_stripped = s.strip()
    if not s_stripped:
        return False
    # Windows路径特征：盘符（如C:）+ 反斜杠或斜杠
    windows_pattern = r"^[a-zA-Z]:[\\/].*"
    # Linux/macOS路径特征：根目录/或相对路径./../
    unix_pattern = r"^(/|\./|\.\./).*"
    # 还可以通过尝试解析路径是否合法进一步验证（可选）
    try:
        # 尝试规范化路径（若报错则不是有效路径）
        os.path.normpath(s_stripped)
        return (
            re.match(windows_pattern, s_stripped) is not None
            or re.match(unix_pattern, s_stripped) is not None
        )
    except:
        return False


# %% ../src/notebooks/04_llm_api.ipynb 9
import aiohttp
import tempfile
import os
from typing import Optional


async def download_file(video_url: str, verbose: bool = False) -> Optional[str]:
    """
    异步下载文件到本地临时文件，并返回临时文件路径

    :param video_url: 要下载的文件的URL
    :param verbose: 是否打印详细信息
    :return: 本地临时文件路径（下载失败返回None）
    """
    # 创建临时文件（默认在系统临时目录，关闭后自动删除，这里手动控制删除时机）
    # mode='wb' 以二进制写模式打开，suffix保留原文件后缀（可选）
    try:
        # 提取URL中的文件名后缀（可选，用于临时文件保留后缀）
        url_path = video_url.split("/")[-1]
        suffix = os.path.splitext(url_path)[1] if "." in url_path else ""

        # 创建临时文件，delete=False表示不自动删除，需要手动管理
        with tempfile.NamedTemporaryFile(
            mode="wb", suffix=suffix, delete=False
        ) as temp_file:
            temp_file_path = temp_file.name  # 保存临时文件路径

        # 异步下载文件
        async with aiohttp.ClientSession() as session:
            async with session.get(
                video_url, timeout=aiohttp.ClientTimeout(total=60)
            ) as response:
                if response.status != 200:
                    if verbose:
                        print(f"下载失败，状态码：{response.status}")
                    os.unlink(temp_file_path)  # 删除无效临时文件
                    return None

                # 分块写入临时文件
                with open(temp_file_path, "wb") as f:
                    async for chunk in response.content.iter_chunked(
                        1024 * 1024
                    ):  # 1MB 块
                        f.write(chunk)

        if verbose:
            print(f"文件下载成功，临时路径：{temp_file_path}")
        return temp_file_path

    except aiohttp.ClientError as e:
        if verbose:
            print(f"网络错误：{str(e)}")
    except Exception as e:
        if verbose:
            print(f"下载失败：{str(e)}")
    finally:
        # 若临时文件存在但未正常写入，清理文件
        if (
            "temp_file_path" in locals()
            and os.path.exists(temp_file_path)
            and os.path.getsize(temp_file_path) == 0
        ):
            os.unlink(temp_file_path)

    return None


# %% ../src/notebooks/04_llm_api.ipynb 11
import base64
import os
import asyncio

try:
    import aiofiles
except ImportError:
    print("请先安装 'aiofiles' 库: !pip install aiofiles")


async def local_video_to_base64_uri(file_path: str) -> str:
    """
    异步地将本地视频文件转换为 Base64 编码的数据 URI。

    格式遵循: data:video/<视频格式>;base64,<Base64编码>，其中，
    视频格式: 支持 mp4, avi, mov。
    Base64编码: 视频文件的 Base64 编码。

    Args:
        file_path: 本地视频文件的路径。

    Returns:
        一个字符串，包含了视频的 Base64 编码数据 URI。

    Raises:
        ValueError: 如果视频格式不是 'mp4', 'avi', 或 'mov' 之一。
        FileNotFoundError: 如果在 file_path 指定的路径下找不到文件。
    """
    supported_formats = {"mp4", "avi", "mov"}

    # 从路径中获取文件扩展名作为视频格式
    file_extension = file_path.split(".")[-1].lower()

    if file_extension not in supported_formats:
        raise ValueError(
            f"不支持的视频格式: '{file_extension}'。支持的格式为: {', '.join(supported_formats)}"
        )

    if not os.path.exists(file_path):
        raise FileNotFoundError(f"文件未找到: {file_path}")

    # 异步读取文件
    async with aiofiles.open(file_path, "rb") as video_file:
        video_bytes = await video_file.read()

    # 进行 Base64 编码
    base64_encoded_video = base64.b64encode(video_bytes).decode("utf-8")

    # 拼接成最终的 Data URI 字符串
    return f"data:video/{file_extension};base64,{base64_encoded_video}"


# %% ../src/notebooks/04_llm_api.ipynb 13
import re
from typing import Optional, Tuple


def separate_think_and_other(text: str) -> Tuple[Optional[str], str]:
    """
    从文本中分离 <think> 标签内容和其余内容。

    Args:
        text: 包含或不包含 <think> 标签的原始字符串。

    Returns:
        一个元组，包含两部分：
        - 第一个元素是所有 <think> 标签内内容的合并字符串（以换行符分隔），如果没有则为 None。
        - 第二个元素是去除 <think> 标签后剩余内容的合并字符串。
    """
    # 使用 re.findall 找到所有 <think> 标签的内容
    # re.DOTALL 标志让 '.' 可以匹配包括换行符在内的任何字符
    think_parts = re.findall(r"<think>(.*?)</think>", text, re.DOTALL)

    if think_parts:
        # 合并所有 <think> 标签内的内容，使用换行符分隔
        think_content = "\n".join(part.strip() for part in think_parts)
    else:
        think_content = None

    # 使用 re.split 来获取 <think> 标签外的内容
    other_parts = re.split(r"<think>.*?</think>", text, flags=re.DOTALL)
    # 合并所有标签外的内容，并清理首尾及中间多余的空白
    other_content = "\n".join(part.strip() for part in other_parts if part.strip())

    return think_content, other_content


# %% ../src/notebooks/04_llm_api.ipynb 15
import os
import asyncio
import time
from openai import AsyncOpenAI, APIError, AsyncAzureOpenAI, OpenAI, AzureOpenAI
from typing import Optional, Tuple
import os


def get_env_bool(env_var, default=False):
    env_val = os.getenv(env_var)
    if env_val is None:
        return default
    true_values = ("true", "1", "yes", "on", "y", "t")
    return env_val.strip().lower() in true_values


class Endpoint:
    def __init__(
        self,
        base_url: Optional[str] = None,
        ip: Optional[str] = None,
        port: Optional[int] = None,
        api_key: Optional[str] = None,
        model_name_or_path: Optional[str] = None,
        use_azure: Optional[bool] = None,  # https://github.com/openai/openai-python
        api_version: Optional[str] = None,
    ):
        self.model_name_or_path = (
            model_name_or_path or os.getenv("OPENAI_MODEL") or "gemini-2.5-flash"
        )

        # Handle IPv6 addresses by wrapping them in brackets for URL compatibility
        if ip and ":" in ip:
            ip = f"[{ip}]"
        if ip and port:
            self.base_url = f"http://{ip}:{port}/v1"
        else:
            self.base_url = (
                base_url
                or os.getenv("OPENAI_BASE_URL")
                or os.getenv("AZURE_OPENAI_ENDPOINT")
            )
        api_key = (
            api_key
            or os.getenv("OPENAI_API_KEY")
            or os.getenv("AZURE_OPENAI_API_KEY")
            or "not-needed"
        )
        api_version = (
            api_version
            or os.getenv("OPENAI_API_VERSION")
            or os.getenv("AZURE_OPENAI_API_VERSION")
            or "2024-03-01-preview"
        )
        use_azure = use_azure or get_env_bool("OPENAI_USE_AZURE")
        if use_azure:
            self.async_client = AsyncAzureOpenAI(
                azure_endpoint=self.base_url,
                api_key=api_key,
                api_version=api_version,  # type: ignore
            )
            self.client = AzureOpenAI(
                azure_endpoint=self.base_url,
                api_key=api_key,
                api_version=api_version,  # type: ignore
            )
        else:
            self.async_client = AsyncOpenAI(base_url=self.base_url, api_key=api_key)
            self.client = OpenAI(base_url=self.base_url, api_key=api_key)

    async def chat_completions_create(self, **kwargs):
        return await self.async_client.chat.completions.create(
            model=self.model_name_or_path, **kwargs
        )

    def chat_completions_create_sync(self, **kwargs):
        return self.client.chat.completions.create(
            model=self.model_name_or_path, **kwargs
        )


# %% ../src/notebooks/04_llm_api.ipynb 17
def flatten_dict(d: dict, level: int, parent_key: str = "", sep: str = ".") -> dict:
    items = []
    for k, v in d.items():
        new_key = parent_key + sep + k if parent_key else k
        if isinstance(v, dict) and v and level > 0:
            items.extend(flatten_dict(v, level - 1, new_key, sep=sep).items())
        else:
            items.append((new_key, v))
    return dict(items)
