import vane


# These placeholders represent PyMuPDF and tokenizer-based helpers.
# See the related runnable example linked at the end for a complete PDF pipeline.
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
MAX_CHUNK_TOKENS = 240
CHUNK_TOKEN_OVERLAP = 32




# flat_map callable: take one PDF row and yield one row per page.
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 callable: take one page row and split it into overlapping tokenizer-aligned chunks.
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 changes the row count, so declare the output schema for each stage.
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,
}


# The same connection carries both SQL queries and subsequent Python Relation operations.
con = vane.connect()


# read_blob loads file metadata and binary content into one Relation.
pdfs = con.sql("""
    SELECT filename AS source, size, content
    FROM read_blob('/data/pdfs/*.pdf')
""")


# The first flat_map call expands one PDF row into multiple page rows.
pages = pdfs.flat_map(extract_pages, schema=PAGE_SCHEMA)


# The second flat_map call expands one page row into multiple text-chunk rows.
chunks = pages.flat_map(split_chunks, schema=CHUNK_SCHEMA)


# Once the chunks form a Relation, switch back to SQL and append a vector field with AI_EMBED.
embedded = chunks.query(
    "chunks",
    """
    SELECT
        *,
        -- Generate a vector for each text chunk while retaining its location fields.
        AI_EMBED(
            text,
            provider := 'transformers',
            model := 'sentence-transformers/all-MiniLM-L6-v2',
            options := struct_pack(
                device := 'cpu',
                batch_size := 32
            )
        ) AS embedding
    FROM chunks
    """,
)


# Continue on the same Relation with AI_PROMPT to generate queryable semantic fields.
result = embedded.query(
    "embedded",
    """
    WITH enriched AS (
        SELECT
            *,
            -- AI_PROMPT returns STRUCT(topics, summary), or NULL on failure.
            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,
        -- Expand the STRUCT directly for downstream SQL filtering and presentation.
        metadata.topics AS topics,
        metadata.summary AS chunk_summary
    FROM enriched
    """,
)


# write_parquet materializes the entire Relation pipeline.
result.write_parquet("/tmp/pdf_chunks_enriched.parquet")

import pyarrow as pa
import vane


# The UDF returns one STRUCT per image, with all three fields directly accessible from SQL.
INSPECTION_TYPE = pa.struct([
    pa.field("width", pa.int32()),
    pa.field("height", pa.int32()),
    pa.field("is_usable", pa.bool_()),
])


# Define a stateless batch UDF. Vane passes two Arrow arrays per batch.
# The function returns a StructArray of the same length, and batch_size controls each batch.
@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()


# Register the Python UDF on the current connection under the SQL name inspect_image.
# parameters declares its input signature as (BLOB, INTEGER).
vane.attach_function(
    inspect_image,
    alias="inspect_image",
    connection=con,
    parameters=["BLOB", "INTEGER"],
)


# Read image files into a Relation containing metadata and BLOB values.
images = con.sql("""
    SELECT filename, size, content AS image
    FROM read_blob('/data/images/*')
""")


# Call the registered Python UDF in a SQL projection. The value 64 is the minimum short side.
inspected = images.query(
    "images",
    """
    SELECT
        *,
        -- Each row receives a STRUCT(width, height, is_usable).
        inspect_image(image, 64) AS inspection
    FROM images
    """,
)


# Route on the UDF's is_usable field in SQL and send only usable images to the model.
result = inspected.query(
    "inspected",
    """
    WITH described AS (
        SELECT
            filename,
            size,
            inspection.width AS width,
            inspection.height AS height,
            -- Call a vision Prompt for each image that passed inspection and require a fixed 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 an English-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,
        -- Expand the model-returned STRUCT into regular queryable columns.
        answer.summary AS summary,
        answer.model_confidence AS model_confidence
    FROM described
    """,
)


# Materialize the preceding UDF inspection and model calls.
result.write_parquet("/tmp/image_descriptions.parquet")

import vane




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


# Declare the types of the three intermediate results for both UDF returns and SQL parameters.
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)




# These placeholders represent audio decoding, feature construction, and token decoding.
# See the related runnable example linked at the end for a complete audio pipeline.


# First stateless batch UDF: decode audio BLOBs in different formats into 16 kHz waveforms.
@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)




# Second stateless batch UDF: convert waveforms into the fixed-shape features Whisper expects.
@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,
    )




# Define a stateful Actor UDF. Each Actor uses one GPU and processes multiple batches.
@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):
        # __init__ runs once when the Actor is created, avoiding a model load for every batch.
        self.model = load_whisper(MODEL_ID, device="cuda")


    def __call__(self, input_features):
        # __call__ processes one feature batch and explicitly selects Chinese transcription.
        return transcribe_to_token_ids(
            input_features,
            model=self.model,
            language="zh",
            task="transcribe",
        )




# Third stateless batch UDF: decode the model's token IDs into Chinese strings.
@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()


# Register the decoding and resampling UDF with the SQL signature decode_resample_16k(BLOB).
vane.attach_function(
    decode_resample_16k,
    alias="decode_resample_16k",
    connection=con,
    parameters=["BLOB"],
)


# Register the feature-generation UDF, matching its input type to the previous stage's return.
vane.attach_function(
    prepare_whisper_features,
    alias="prepare_whisper_features",
    connection=con,
    parameters=[RESAMPLED_AUDIO_TYPE],
)


# Register the Actor instance. Every SQL call reuses the model already loaded in the Actor.
vane.attach_function(
    WhisperTranscriber(),
    alias="whisper_transcribe_zh",
    connection=con,
    parameters=[INPUT_FEATURES_TYPE],
)


# Register the token-decoding UDF, converting INTEGER[] into VARCHAR.
vane.attach_function(
    decode_whisper_tokens,
    alias="decode_whisper_tokens",
    connection=con,
    parameters=[TOKEN_IDS_TYPE],
)


# Read a Parquet Relation containing the audio STRUCT and sample metadata.
source = con.sql("""
    SELECT *
    FROM read_parquet('/data/chinese-speech/*.parquet')
""")


# Connect the registered UDFs through staged SQL CTEs and the final projection.
result = source.query(
    "source",
    """
    WITH audio AS (
        SELECT
            * EXCLUDE (audio),
            -- Extract the raw audio BLOB from the audio STRUCT.
            audio.bytes AS audio_bytes
        FROM source
    ),
    resampled AS (
        SELECT
            * EXCLUDE (audio_bytes),
            -- Stage 1: call a Python UDF to decode and resample the audio to 16 kHz.
            decode_resample_16k(audio_bytes) AS waveform
        FROM audio
    ),
    featured AS (
        SELECT
            * EXCLUDE (waveform),
            -- Stage 2: call a Python UDF to generate Whisper input tensors.
            prepare_whisper_features(waveform) AS input_features
        FROM resampled
    ),
    tokens AS (
        SELECT
            * EXCLUDE (input_features),
            -- Stage 3: call an Actor UDF to run model inference on the GPU.
            whisper_transcribe_zh(input_features) AS token_ids
        FROM featured
    )
    SELECT
        * EXCLUDE (token_ids),
        -- Stage 4: call a Python UDF to decode token IDs into the final text.
        decode_whisper_tokens(token_ids) AS transcription
    FROM tokens
    """,
)


# Write only the business fields and transcript. The write materializes the entire UDF chain.
result.write_parquet("/tmp/chinese_audio_transcriptions.parquet")

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




# Stateful batch callable: Vane creates and manages the Actor after it is passed to map_batches.
class Detector:
    def __init__(self):
        # Each Actor loads the YOLO model only once during its lifetime.
        self.model = load_yolo("yolo11n.pt", device="cuda")


    def __call__(self, table):
        # __call__ receives one Arrow Table batch and detects objects in all of its video frames.
        return run_object_detection(
            table,
            model=self.model,
            frame_column="frame",
            fields={
                "label": "boxes.cls",
                "confidence": "boxes.conf",
                "bbox": "boxes.xyxy",
            },
            # Retain source fields and frame indices with the detection results.
            pass_through=(
                "source_id",
                "video_path",
                "frame_index",
                "frame",
            ),
        )




# Stateless batch callable: expand each frame's object list into multiple cropped-result rows.
def crop_object_batch(table):
    # A frame may contain multiple objects, so expand it into rows and crop each object.
    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 decodes the videos, and read_datasource organizes the output as a frame Relation.
# Each row represents one frame and carries its source, path, frame index, and frame tensor.
frames = read_datasource(
    VideoFrameSource(
        video_paths,
        height=640,
        width=640,
    ),
    con=con,
)


# The first map_batches call runs Detector. Vane creates a GPU Actor and reuses its model.
detected = frames.map_batches(
    Detector,
    # schema describes every column in the Relation produced by the detection stage.
    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,
)


# The second map_batches call runs a stateless function to expand and crop objects on the CPU.
objects = detected.map_batches(
    crop_object_batch,
    # After expansion, each row represents one object and adds a cropped 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,
    },
)


# Select the final output columns explicitly; the crop stage has already dropped the frame tensor.
result = objects.project(
    "source_id, video_path, frame_index, features, object"
)


# write_parquet materializes the preceding read, detection, and crop stages.
result.write_parquet("/tmp/video_objects.parquet")

from io import BytesIO


from PIL import Image, ImageOps


import vane


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
}"""


@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()
routes = connection.sql(
    """
    WITH prepared AS (
        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.*,
        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 *,
               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},
)


routes.write_parquet("claim_routes.parquet")