import pyarrow.fs as pafs
import argparse
import logging

from apache_beam.options.pipeline_options import PipelineOptions

from deltacat.examples.experimental.iceberg.converter.beam import app


if __name__ == "__main__":
    logging.getLogger().setLevel(logging.INFO)
    parser = argparse.ArgumentParser(
        description="DeltaCat Beam Iceberg Converter Example using REST Catalog",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Start REST catalog server first (Iceberg 1.6.0):
  docker run -d -p 8181:8181 --name iceberg-rest-catalog tabulario/iceberg-rest:1.6.0

  # Install PySpark for read and rewrite modes:
  pip install pyspark

  # Write sample data with DeltaCAT data file converter (automatic merge by key):
  python main.py --mode write --table-name "deltacat.hello_world"

  # Read data back (uses Spark SQL to read positional deletes):
  python main.py --mode read --table-name "deltacat.hello_world"

  # Rewrite table to materialize positional deletes:
  python main.py --mode rewrite --table-name "deltacat.hello_world"

  # Use custom REST catalog server:
  python main.py --mode write --rest-uri http://localhost:9000 --table-name "deltacat.hello_world"

  # Use custom warehouse path:
  python main.py --mode write --warehouse-path /tmp/my_warehouse --table-name "deltacat.hello_world"
        """,
    )

    parser.add_argument(
        "--mode",
        default="write",
        choices=["write", "read", "rewrite"],
        help="Pipeline mode: 'write' to write data, 'read' to read data, 'rewrite' to materialize positional deletes (default: write). "
        "   Note: Beam writes may fail on tables processed by external tools.",
    )

    parser.add_argument(
        "--rest-uri",
        default="http://localhost:8181",
        help="REST catalog server URI (default: http://localhost:8181).",
    )

    parser.add_argument(
        "--warehouse-path",
        default=None,
        help="Custom warehouse path (default: temporary directory).",
    )

    parser.add_argument(
        "--table-name",
        default=None,
        help="Table name to use (default: autogenerated table name).",
    )

    parser.add_argument(
        "--deltacat-converter-interval",
        type=float,
        default=5.0,
        help="DeltaCat converter monitoring interval in seconds (default: 5.0).",
    )

    parser.add_argument(
        "--ray-inactivity-timeout",
        type=int,
        default=20,
        help="Ray cluster shutdown timeout after inactivity in seconds (default: 20).",
    )

    parser.add_argument(
        "--max-converter-parallelism",
        type=int,
        default=1,
        help="Maximum converter task parallelism - number of concurrent converter tasks (default: 1).",
    )

    args, beam_args = parser.parse_known_args()

    beam_options = PipelineOptions(
        beam_args,
        save_main_session=True,
    )

    print("DeltaCAT Beam Iceberg Upsert Example")
    print("=" * 50)
    print(f"Mode: {args.mode}")
    print(f"REST Catalog URI: {args.rest_uri}")
    print(f"Warehouse Path: {args.warehouse_path or 'temporary directory'}")
    print(f"Table Name: {args.table_name}")
    print(f"Converter Interval: {args.deltacat_converter_interval}s")
    print(f"Ray Inactivity Timeout: {args.ray_inactivity_timeout}s")
    print(f"Max Converter Parallelism: {args.max_converter_parallelism}")
    print()

    # Remind user about prerequisites
    if args.mode == "write":
        print("Prerequisites:")
        print("   Make sure the Iceberg REST catalog server is running:")
        print(
            "   docker run -d -p 8181:8181 --name iceberg-rest-catalog tabulario/iceberg-rest:1.6.0"
        )
        print()
    elif args.mode in ["read", "rewrite"]:
        print("Prerequisites:")
        print("   Make sure the Iceberg REST catalog server is running:")
        print(
            "   docker run -d -p 8181:8181 --name iceberg-rest-catalog tabulario/iceberg-rest:1.6.0"
        )
        print("   PySpark is required for this mode:")
        print("   pip install pyspark")
        print()

    app.run(
        beam_options=beam_options,
        mode=args.mode,
        rest_catalog_uri=args.rest_uri,
        warehouse_path=args.warehouse_path,
        table_name=args.table_name,
        deltacat_converter_interval=args.deltacat_converter_interval,
        ray_inactivity_timeout=args.ray_inactivity_timeout,
        filesystem=pafs.LocalFileSystem(),
        max_converter_parallelism=args.max_converter_parallelism,
    )
