#!/usr/bin/env python3
#
# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
#

"""
CLI tool for creating a Snowflake Python stored procedure that runs a
user-supplied Python module via ``runpy.run_module(..., run_name="__main__")``
inside an SCOS server.

Thin wrapper around the programmatic API in
``snowflake.snowpark_connect.orchestration.sproc_manager``.

Examples:

  # Basic: create a stored procedure for a Python workload
  snowpark-connect-create-python-sproc my_workload.py

  # With extra Python files added to IMPORTS
  snowpark-connect-create-python-sproc my_workload.py \\
      --py-files "utils/helpers.py,utils/io.py"

  # With a requirements.txt
  snowpark-connect-create-python-sproc my_workload.py \\
      --requirements-file requirements.txt \\
      --replace

  # Glob patterns for --py-files
  snowpark-connect-create-python-sproc my_workload.py \\
      --py-files "src/*.py"

  # Dry run: print the full setup script (CREATE STAGE + PUT statements +
  # CREATE PROCEDURE) without executing anything against Snowflake. The
  # output can be copy-pasted into a Snowflake worksheet.
  snowpark-connect-create-python-sproc my_workload.py --dry-run
"""

import argparse
import logging
import os
import sys

logger = logging.getLogger("snowpark-connect-create-python-sproc")


def main():
    parser = argparse.ArgumentParser(
        description=(
            "Create a Snowflake stored procedure that runs a Python module "
            "via runpy.run_module inside an SCOS server."
        ),
        epilog=(
            "Examples:\n\n"
            "  # Basic\n"
            "  snowpark-connect-create-python-sproc my_workload.py\n\n"
            "  # With extra py files and requirements\n"
            "  snowpark-connect-create-python-sproc my_workload.py \\\n"
            '      --py-files "utils/*.py" \\\n'
            "      --requirements-file requirements.txt \\\n"
            "      --replace\n\n"
            "  # Dry run\n"
            "  snowpark-connect-create-python-sproc my_workload.py --dry-run\n"
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "module_file",
        help="Path to the main Python module (.py file). Its top-level code "
        "is executed as __main__.",
    )
    parser.add_argument(
        "--py-files",
        default=None,
        help="Comma-separated .py files or glob patterns to include in the "
        "stored procedure's IMPORTS clause.",
    )
    parser.add_argument(
        "--requirements-file",
        default=None,
        help="Path to a requirements.txt; parsed into the PACKAGES clause. "
        "Strict format: 'name' or 'name<op>version' with op in {==, >=, <=, >, <}.",
    )
    parser.add_argument(
        "--procedure-name",
        default=None,
        help="Snowflake procedure name. Defaults to <module_basename>_main.",
    )
    parser.add_argument(
        "--replace",
        action="store_true",
        help="Use CREATE OR REPLACE PROCEDURE (overwrites existing).",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print the full setup script (CREATE STAGE + PUT + CREATE "
        "PROCEDURE) without executing anything against Snowflake.",
    )
    parser.add_argument(
        "--snowpark-connect-version",
        default=None,
        help="Pin snowpark-connect to a specific version (e.g. '1.21.1'). "
        "If not set, uses the latest available in Anaconda.",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Enable verbose (DEBUG) logging.",
    )

    args = parser.parse_args()

    log_level = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(
        level=log_level,
        format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
    )

    from snowflake.snowpark_connect.orchestration.sproc_manager import (
        create_python_sproc,
        resolve_py_files,
    )

    try:
        py_files: list[str] = []
        if args.py_files:
            py_files = resolve_py_files(
                [p.strip() for p in args.py_files.split(",") if p.strip()]
            )
            logger.info(
                f"Resolved {len(py_files)} additional .py file(s): "
                f"{[os.path.basename(p) for p in py_files]}"
            )

        sql = create_python_sproc(
            module_file=args.module_file,
            py_files=py_files,
            requirements_file=args.requirements_file,
            procedure_name=args.procedure_name,
            replace=args.replace,
            dry_run=args.dry_run,
            snowpark_connect_version=args.snowpark_connect_version,
        )

        if args.dry_run:
            print(sql)  # noqa: T201

    except (FileNotFoundError, ValueError, TypeError, RuntimeError) as e:
        logger.error(str(e))
        sys.exit(1)


if __name__ == "__main__":
    main()
