#!/usr/bin/env python
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The Geneva Authors

"""
Upload manifest and add columns for BLIP caption UDF to Geneva.

This script runs in the UDF's dependency environment (via uv run)
and uploads the manifest to the Geneva connection, then adds columns
to the specified table.

Usage:
    export GENEVA_TABLE_NAME=oxford_pets_shared_abc123
    cd e2e/oxford-pets/udfs/blip
    uv run python upload_manifest.py --bucket gs://bucket/path

Environment Variables:
    GENEVA_TABLE_NAME: Table name to add columns to (required)
    CAPTION_COL: Custom column name for caption (default: "caption")
"""

import argparse
import logging
import os

import geneva
from caption_udf import GenCaption
from manifest import create_manifest

logging.basicConfig(level=logging.INFO)
_LOG = logging.getLogger(__name__)


def get_columns() -> dict[str, callable]:
    """
    Get column definitions for this UDF package.

    Column names can be customized via environment variables:
    - CAPTION_COL: Custom name for caption column (default: "caption")

    Returns:
        Dictionary mapping column names to UDF functions
    """
    return {
        os.getenv("CAPTION_COL", "caption"): GenCaption(),
    }


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Upload BLIP caption UDF manifest and add columns to table"
    )
    parser.add_argument(
        "--bucket", required=True, help="Geneva bucket path (e.g., gs://bucket/path)"
    )
    parser.add_argument(
        "--manifest-name",
        default="blip-udfs-v1",
        help="Manifest name (default: blip-udfs-v1)",
    )
    args = parser.parse_args()

    # Get table name from environment
    table_name = os.getenv("GENEVA_TABLE_NAME")
    if not table_name:
        raise ValueError(
            "GENEVA_TABLE_NAME environment variable must be set "
            "(e.g., export GENEVA_TABLE_NAME=oxford_pets_shared_abc123)"
        )

    # Create manifest in UDF's environment
    _LOG.info(f"Creating manifest '{args.manifest_name}' in UDF environment")
    manifest = create_manifest(name=args.manifest_name)

    # Connect and upload
    _LOG.info(f"Connecting to {args.bucket}")
    conn = geneva.connect(args.bucket)

    _LOG.info(f"Uploading manifest '{args.manifest_name}'")
    conn.define_manifest(args.manifest_name, manifest)

    _LOG.info(f"✓ Successfully uploaded manifest '{args.manifest_name}' to {args.bucket}")
    _LOG.info(f"  Dependencies: {manifest.pip}")
    _LOG.info(f"  Modules: {manifest.py_modules}")

    # Add columns to table
    _LOG.info(f"Opening table '{table_name}'")
    tbl = conn.open_table(table_name)

    columns = get_columns()
    _LOG.info(f"Adding columns: {list(columns.keys())}")
    tbl.add_columns(columns)

    _LOG.info(f"✓ Successfully added columns to table '{table_name}'")
    _LOG.info(f"  Columns: {list(columns.keys())}")


if __name__ == "__main__":
    main()
