import vane


# 以下占位函数分别代表 PyMuPDF 和基于 tokenizer 的辅助逻辑。
# 完整的 PDF 处理流程请参考文末链接的可运行示例。
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
MAX_CHUNK_TOKENS = 240
CHUNK_TOKEN_OVERLAP = 32




# flat_map 回调：接收一条 PDF 记录，按页生成多条记录。
def extract_pages(row):
    for page_number, text in parse_pdf(bytes(row["content"])):
        yield {
            "source": row["source"],
            "size": row["size"],
            "page_number": page_number,
            "text": text,
        }




# flat_map 回调：接收一条页面记录，将其切分为按 tokenizer 对齐且相互重叠的文本块。
def split_chunks(row):
    for chunk_index, text in enumerate(
        split_text_by_tokens(
            row["text"],
            model=EMBEDDING_MODEL,
            max_tokens=MAX_CHUNK_TOKENS,
            overlap_tokens=CHUNK_TOKEN_OVERLAP,
        )
    ):
        yield {
            "source": row["source"],
            "size": row["size"],
            "page_number": row["page_number"],
            "chunk_index": chunk_index,
            "text": text,
        }




# flat_map 会改变行数，因此每个阶段都需要声明输出 schema。
PAGE_SCHEMA = {
    "source": vane.sqltypes.VARCHAR,
    "size": vane.sqltypes.BIGINT,
    "page_number": vane.sqltypes.INTEGER,
    "text": vane.sqltypes.VARCHAR,
}
CHUNK_SCHEMA = {
    **PAGE_SCHEMA,
    "chunk_index": vane.sqltypes.INTEGER,
}


# 同一个连接既执行 SQL 查询，也承接后续的 Python Relation 操作。
con = vane.connect()


# read_blob 将文件元数据和二进制内容读入同一个 Relation。
pdfs = con.sql("""
    SELECT filename AS source, size, content
    FROM read_blob('/data/pdfs/*.pdf')
""")


# 第一次 flat_map 将一条 PDF 记录展开为多条页面记录。
pages = pdfs.flat_map(extract_pages, schema=PAGE_SCHEMA)


# 第二次 flat_map 将一条页面记录展开为多条文本块记录。
chunks = pages.flat_map(split_chunks, schema=CHUNK_SCHEMA)


# 文本块形成 Relation 后切回 SQL，通过 AI_EMBED 增加向量字段。
embedded = chunks.query(
    "chunks",
    """
    SELECT
        *,
        -- 为每个文本块生成向量，同时保留定位字段。
        AI_EMBED(
            text,
            provider := 'transformers',
            model := 'sentence-transformers/all-MiniLM-L6-v2',
            options := struct_pack(
                device := 'cpu',
                batch_size := 32
            )
        ) AS embedding
    FROM chunks
    """,
)


# 在同一个 Relation 上继续调用 AI_PROMPT，生成可查询的语义字段。
result = embedded.query(
    "embedded",
    """
    WITH enriched AS (
        SELECT
            *,
            -- AI_PROMPT 返回 STRUCT(topics, summary)，失败时返回 NULL。
            AI_PROMPT(
                text,
                return_format := json '{
                    "type": "object",
                    "properties": {
                        "topics": {
                            "type": "array",
                            "items": {"type": "string"}
                        },
                        "summary": {"type": "string"}
                    },
                    "required": ["topics", "summary"],
                    "additionalProperties": false
                }',
                system_message :=
                    'Extract the main topics from the text chunk and summarize it in one sentence.',
                provider := 'vllm',
                model := 'Qwen/Qwen2.5-7B-Instruct',
                on_error := 'ignore',
                options := struct_pack(
                    max_tokens := 128,
                    temperature := 0.0
                )
            ) AS metadata
        FROM embedded
    )
    SELECT
        source,
        size,
        page_number,
        chunk_index,
        text,
        embedding,
        -- 直接展开 STRUCT，供后续 SQL 筛选和展示。
        metadata.topics AS topics,
        metadata.summary AS chunk_summary
    FROM enriched
    """,
)


# write_parquet 触发整条 Relation 流水线执行。
result.write_parquet("/tmp/pdf_chunks_enriched.parquet")

import pyarrow as pa
import vane


# UDF 为每张图片返回一个 STRUCT，三个字段都可直接在 SQL 中访问。
INSPECTION_TYPE = pa.struct([
    pa.field("width", pa.int32()),
    pa.field("height", pa.int32()),
    pa.field("is_usable", pa.bool_()),
])


# 定义无状态批量 UDF；Vane 每批传入两个 Arrow Array。
# 函数返回等长的 StructArray，batch_size 控制单批大小。
@vane.func.batch(return_dtype=INSPECTION_TYPE, batch_size=32)
def inspect_image(image, minimum_side):
    return inspect_image_blobs(image, minimum_side)


con = vane.connect()


# 将 Python UDF 注册到当前连接，SQL 函数名为 inspect_image。
# parameters 声明函数签名为 (BLOB, INTEGER)。
vane.attach_function(
    inspect_image,
    alias="inspect_image",
    connection=con,
    parameters=["BLOB", "INTEGER"],
)


# 将图片文件读入包含元数据和 BLOB 值的 Relation。
images = con.sql("""
    SELECT filename, size, content AS image
    FROM read_blob('/data/images/*')
""")


# 在 SQL 投影中调用已注册的 Python UDF；64 表示最短边阈值。
inspected = images.query(
    "images",
    """
    SELECT
        *,
        -- 每行得到一个 STRUCT(width, height, is_usable)。
        inspect_image(image, 64) AS inspection
    FROM images
    """,
)


# 在 SQL 中按 is_usable 分流，只把可用图片交给模型。
result = inspected.query(
    "inspected",
    """
    WITH described AS (
        SELECT
            filename,
            size,
            inspection.width AS width,
            inspection.height AS height,
            -- 对通过检查的图片调用视觉 Prompt，并要求返回固定结构的 STRUCT。
            AI_PROMPT(
                'Describe the main subject in the image and provide a confidence score from 0 to 1.',
                image,
                return_format := json '{
                    "type": "object",
                    "properties": {
                        "summary": {"type": "string"},
                        "model_confidence": {"type": "number"}
                    },
                    "required": ["summary", "model_confidence"],
                    "additionalProperties": false
                }',
                system_message := 'Return only a Chinese-language result that matches the requested structure.',
                provider := 'openai',
                model := 'gpt-4o-mini',
                on_error := 'ignore',
                options := struct_pack(
                    use_chat_completions := true,
                    max_output_tokens := 128,
                    temperature := 0.0
                )
            ) AS answer
        FROM inspected
        WHERE inspection.is_usable
    )
    SELECT
        filename,
        size,
        width,
        height,
        -- 将模型返回的 STRUCT 展开为普通的可查询列。
        answer.summary AS summary,
        answer.model_confidence AS model_confidence
    FROM described
    """,
)


# 写出时执行前面的 UDF 检查和模型调用。
result.write_parquet("/tmp/image_descriptions.parquet")

import vane




MODEL_ID = "openai/whisper-tiny"
SAMPLING_RATE = 16000
BATCH_SIZE = 128
NUM_GPU_ACTORS = 1


# 声明三个中间结果的类型，同时用于 UDF 返回值和 SQL 参数。
FEATURE_MELS = 80
FEATURE_FRAMES = 3000
RESAMPLED_AUDIO_TYPE = vane.list_type(vane.sqltypes.FLOAT)
INPUT_FEATURES_TYPE = vane.tensor_type(
    vane.sqltypes.FLOAT,
    (FEATURE_MELS, FEATURE_FRAMES),
)
TOKEN_IDS_TYPE = vane.list_type(vane.sqltypes.INTEGER)




# 以下占位函数分别代表音频解码、特征构建和 token 解码。
# 完整的音频处理流程请参考文末链接的可运行示例。


# 第一个无状态批量 UDF：将不同格式的音频 BLOB 解码并重采样为 16 kHz 波形。
@vane.func.batch(return_dtype=RESAMPLED_AUDIO_TYPE, batch_size=BATCH_SIZE)
def decode_resample_16k(audio_bytes):
    return decode_and_resample(audio_bytes, sample_rate=SAMPLING_RATE)




# 第二个无状态批量 UDF：将波形转换为 Whisper 要求的定长特征。
@vane.func.batch(return_dtype=INPUT_FEATURES_TYPE, batch_size=BATCH_SIZE)
def prepare_whisper_features(waveform):
    return build_whisper_features(
        waveform,
        model=MODEL_ID,
        sample_rate=SAMPLING_RATE,
    )




# 定义有状态 Actor UDF；每个 Actor 使用一张 GPU，并连续处理多个批次。
@vane.cls.batch(
    actor_number=NUM_GPU_ACTORS,
    gpus=1.0,
    return_dtype=TOKEN_IDS_TYPE,
    batch_size=BATCH_SIZE,
)
class WhisperTranscriber:
    def __init__(self):
        # Actor 创建时只执行一次 __init__，避免每批重复加载模型。
        self.model = load_whisper(MODEL_ID, device="cuda")


    def __call__(self, input_features):
        # __call__ 处理一个特征批次，并显式指定中文转写。
        return transcribe_to_token_ids(
            input_features,
            model=self.model,
            language="zh",
            task="transcribe",
        )




# 第三个无状态批量 UDF：将模型输出的 token ID 解码为中文字符串。
@vane.func.batch(return_dtype=vane.sqltypes.VARCHAR, batch_size=BATCH_SIZE)
def decode_whisper_tokens(token_ids):
    return decode_token_ids(token_ids, model=MODEL_ID)




con = vane.connect()


# 注册解码与重采样 UDF，对应 SQL 签名 decode_resample_16k(BLOB)。
vane.attach_function(
    decode_resample_16k,
    alias="decode_resample_16k",
    connection=con,
    parameters=["BLOB"],
)


# 注册特征生成 UDF，其输入类型与上一阶段的返回类型一致。
vane.attach_function(
    prepare_whisper_features,
    alias="prepare_whisper_features",
    connection=con,
    parameters=[RESAMPLED_AUDIO_TYPE],
)


# 注册 Actor 实例；后续 SQL 调用都会复用已经加载的模型。
vane.attach_function(
    WhisperTranscriber(),
    alias="whisper_transcribe_zh",
    connection=con,
    parameters=[INPUT_FEATURES_TYPE],
)


# 注册 token 解码 UDF，将 INTEGER[] 转换为 VARCHAR。
vane.attach_function(
    decode_whisper_tokens,
    alias="decode_whisper_tokens",
    connection=con,
    parameters=[TOKEN_IDS_TYPE],
)


# 读取包含 audio STRUCT 和样本元数据的 Parquet Relation。
source = con.sql("""
    SELECT *
    FROM read_parquet('/data/chinese-speech/*.parquet')
""")


# 通过分阶段的 SQL CTE 和最终投影串联已注册的 UDF。
result = source.query(
    "source",
    """
    WITH audio AS (
        SELECT
            * EXCLUDE (audio),
            -- 从 audio STRUCT 中提取原始音频 BLOB。
            audio.bytes AS audio_bytes
        FROM source
    ),
    resampled AS (
        SELECT
            * EXCLUDE (audio_bytes),
            -- 阶段 1：调用 Python UDF，将音频解码并重采样为 16 kHz。
            decode_resample_16k(audio_bytes) AS waveform
        FROM audio
    ),
    featured AS (
        SELECT
            * EXCLUDE (waveform),
            -- 阶段 2：调用 Python UDF，生成 Whisper 输入张量。
            prepare_whisper_features(waveform) AS input_features
        FROM resampled
    ),
    tokens AS (
        SELECT
            * EXCLUDE (input_features),
            -- 阶段 3：调用 Actor UDF，在 GPU 上执行模型推理。
            whisper_transcribe_zh(input_features) AS token_ids
        FROM featured
    )
    SELECT
        * EXCLUDE (token_ids),
        -- 阶段 4：调用 Python UDF，将 token ID 解码为最终文本。
        decode_whisper_tokens(token_ids) AS transcription
    FROM tokens
    """,
)


# 只写出业务字段和转写文本；写出操作会触发整条 UDF 链执行。
result.write_parquet("/tmp/chinese_audio_transcriptions.parquet")

import vane
from vane.datasource import read_datasource
from vane.datasource.video_reader import VideoFrameSource




# 有状态批量回调：传给 map_batches 后，由 Vane 创建并管理 Actor。
class Detector:
    def __init__(self):
        # 每个 Actor 在生命周期内只加载一次 YOLO 模型。
        self.model = load_yolo("yolo11n.pt", device="cuda")


    def __call__(self, table):
        # __call__ 接收一个 Arrow Table 批次，对其中所有视频帧执行目标检测。
        return run_object_detection(
            table,
            model=self.model,
            frame_column="frame",
            fields={
                "label": "boxes.cls",
                "confidence": "boxes.conf",
                "bbox": "boxes.xyxy",
            },
            # 保留来源字段和帧序号，方便定位检测结果。
            pass_through=(
                "source_id",
                "video_path",
                "frame_index",
                "frame",
            ),
        )




# 无状态批量回调：将每帧的目标列表展开为多条裁剪结果。
def crop_object_batch(table):
    # 一帧可能包含多个目标，因此需要展开为多行，并分别裁剪。
    return crop_detected_objects(
        table,
        frame_column="frame",
        features_column="features",
        pass_through=("source_id", "video_path", "frame_index"),
        object_column="object",
    )




con = vane.connect()


# VideoFrameSource 负责视频解码，read_datasource 将输出组织成帧 Relation。
# 每行表示一帧，并携带来源、路径、帧序号和帧张量。
frames = read_datasource(
    VideoFrameSource(
        video_paths,
        height=640,
        width=640,
    ),
    con=con,
)


# 第一次 map_batches 调用 Detector；Vane 创建 GPU Actor 并复用其中的模型。
detected = frames.map_batches(
    Detector,
    # schema 描述检测阶段生成的 Relation 中的全部字段。
    schema={
        "source_id": vane.sqltypes.VARCHAR,
        "video_path": vane.sqltypes.VARCHAR,
        "frame_index": vane.sqltypes.BIGINT,
        "frame": FRAME_TYPE,
        "features": FEATURE_LIST_TYPE,
    },
    batch_size=16,
    actor_number=1,
    gpus=1.0,
)


# 第二次 map_batches 在 CPU 上调用无状态函数，展开并裁剪目标。
objects = detected.map_batches(
    crop_object_batch,
    # 展开后每行表示一个目标，并增加裁剪后的 PNG BLOB。
    schema={
        "source_id": vane.sqltypes.VARCHAR,
        "video_path": vane.sqltypes.VARCHAR,
        "frame_index": vane.sqltypes.BIGINT,
        "features": FEATURE_TYPE,
        "object": vane.sqltypes.BLOB,
    },
)


# 显式选择最终输出字段；裁剪阶段已经移除帧张量。
result = objects.project(
    "source_id, video_path, frame_index, features, object"
)


# write_parquet 触发前面的读取、检测和裁剪阶段执行。
result.write_parquet("/tmp/video_objects.parquet")

from io import BytesIO


from PIL import Image, ImageOps


import vane


# Constrain model output to directly queryable structured fields.
damage_schema = """{
  "type": "object",
  "properties": {
    "damage_part": {"type": "string"},
    "severity": {"type": "string", "enum": ["low", "medium", "high"]},
    "confidence": {"type": "number"}
  },
  "required": ["damage_part", "severity", "confidence"],
  "additionalProperties": false
}"""


# Correct orientation, validate dimensions, and normalize the image to JPEG.
@vane.func(return_dtype="BLOB")
def prepare_image(raw: bytes) -> bytes | None:
    try:
        with Image.open(BytesIO(raw)) as image:
            image = ImageOps.exif_transpose(image).convert("RGB")
            if min(image.size) < 480:
                return None
            image.thumbnail((1024, 1024))
            output = BytesIO()
            image.save(output, "JPEG", quality=85)
            return output.getvalue()
    except (OSError, TypeError):
        return None


connection = vane.connect()
# Compose relational processing, the Python UDF, the AI Function, and business rules.
routes = connection.sql(
    """
    WITH prepared AS (
        -- Join, filter, and preprocess each image with the Python UDF.
        SELECT c.claim_id, c.claim_amount,
               prepare_image(i.content) AS image
        FROM pending_claims AS c
        LEFT JOIN claim_images AS i USING (claim_id)
        WHERE c.status = 'pending'
    )
    SELECT
        claim_id,
        assessment.*,
        -- Combine the model result and claim amount for business routing.
        CASE
            WHEN image IS NULL THEN 'Request more photos'
            WHEN assessment IS NULL OR assessment.confidence < 0.65 THEN 'Manual review'
            WHEN assessment.severity = 'high' OR claim_amount >= 50000 THEN 'Claims review'
            ELSE 'Automatic routing'
        END AS next_step
    FROM (
        SELECT *,
               -- Convert the preprocessed image into a structured assessment.
               CASE WHEN image IS NOT NULL THEN ai_prompt(
                   'Identify the damaged vehicle part and severity.',
                   image,
                   return_format := $damage_schema,
                   system_message := 'Use only visual evidence; lower confidence when uncertain.',
                   provider := 'openai',
                   model := 'gpt-4o-mini',
                   on_error := 'ignore'
               ) END AS assessment
        FROM prepared
    ) AS assessed
    """,
    params={"damage_schema": damage_schema},
)


# Materialize the complete processing plan.
routes.write_parquet("claim_routes.parquet")