"""CLI to create the S3 Sink connector.

See https://docs.confluent.io/current/connect/kafka-connect-s3
"""

__all__ = ["create_s3_sink"]

import json
import time
from typing import List

import click

from kafkaconnect.connect import Connect
from kafkaconnect.s3_sink.config import S3Config
from kafkaconnect.topics import Topic


@click.command("s3-sink")
@click.argument("topiclist", nargs=-1, required=False)
@click.option(
    "-n",
    "--name",
    "name",
    envvar="KAFKA_CONNECT_NAME",
    default="s3-sink",
    show_default=True,
    help=(
        "Name of the connector. Alternatively set via the "
        "$KAFKA_CONNECT_NAME env var."
    ),
)
@click.option(
    "-b",
    "--bucket-name",
    "s3_bucket_name",
    envvar="KAFKA_CONNECT_S3_BUCKET_NAME",
    default="mybucket",
    show_default=True,
    help=(
        "S3 bucket name. Must exist already. Alternatively set via the "
        "$KAFKA_CONNECT_S3_BUCKECT_NAME env var."
    ),
)
@click.option(
    "-r",
    "--region",
    "s3_region",
    envvar="KAFKA_CONNECT_S3_REGION",
    default="us-east-1",
    show_default=True,
    help=(
        "S3 region.Alternatively set via the $KAFKA_CONNECT_S3_REGION env var."
    ),
)
@click.option(
    "-id",
    "--aws-access-key-id",
    "aws_access_key_id",
    envvar="AWS_ACCESS_KEY_ID",
    default="",
    show_default=True,
    help=(
        "The AWS access key ID used to authenticate personal AWS credentials."
    ),
)
@click.option(
    "-secret",
    "--aws-secret-access-key",
    "aws_secret_access_key",
    envvar="AWS_SECRET_ACCESS_KEY",
    default="",
    show_default=True,
    help=(
        "The secret access key used to authenticate personal AWS credentials."
    ),
)
@click.option(
    "-d",
    "--topics-dir",
    "topics_dir",
    envvar="KAFKA_CONNECT_S3_TOPICS_DIR",
    default="topics",
    show_default=True,
    help=(
        "Top level directory to store the data ingested from Kafka. "
        "Alternatively set via the $KAFKA_CONNECT_S3_TOPICS_DIR env var."
    ),
)
@click.option(
    "-s",
    "--s3-schema-compatibility",
    "s3_schema_compatibility",
    envvar="KAFKA_CONNECT_S3_SCHEMA_COMPATIBILITY",
    default="NONE",
    show_default=True,
    help=(
        "The S3 schema compatibility mode. "
        "The supported configurations are NONE, BACKWARD, FORWARD and FULL."
    ),
)
@click.option(
    "--flush-size",
    "flush_size",
    envvar="KAFKA_CONNECT_S3_FLUSH_SIZE",
    default="3600",
    show_default=True,
    help=(
        "Number of records written to store before invoking file commits. "
        "By default this is set to six times the number of records expected "
        "for an output stream of 1Hz within the default rotate_interval_ms "
        "value. This way the rotate_interval_ms configuration takes "
        "precedence over the flush_size configuration. But flush_size "
        "still works as a maximum limit to invoke file commits when the "
        "connect-s3-sink consumer accumulates 3600 records. "
        "Alternatively set via the $KAFKA_CONNECT_S3_FLUSH_SIZE env var."
    ),
)
@click.option(
    "--rotate-interval-ms",
    "rotate_interval_ms",
    envvar="KAFKA_CONNECT_S3_ROTATE_INTERVAL_MS",
    default="600000",
    show_default=True,
    help=(
        "The time interval in milliseconds to invoke file commits. "
        "Use this option to control the size of the objects in S3. "
        "For example, if the output data stream is 1Hz, a rotate "
        "interval of 600 seconds will create a file with aproximatelly"
        "600 records if less than flush_size. Note that the lag of the"
        "connect-s3-sink consumer will increase until it accumulates "
        "records within 600 seconds and it will decrease again after "
        "the file is commited. For an hourly partitioner this configuration "
        "should create 6 parquet files at the destination path. "
        "Alternatively set via the $KAFKA_CONNECT_S3_ROTATE_INTERVAL_MS "
        "env var."
    ),
)
@click.option(
    "-p",
    "--partition-duration-ms",
    "partition_duration_ms",
    envvar="KAFKA_CONNECT_S3_PARTITION_DURATION_MS",
    default="3600000",
    show_default=True,
    help=(
        "The duration of a partition in milliseconds used by "
        "``TimeBasedPartitioner``. Alternatively set via the "
        "$KAFKA_CONNECT_S3_PARTITION_DURATION_MS env var."
    ),
)
@click.option(
    "--path-format",
    "path_format",
    envvar="KAFKA_CONNECT_S3_PATH_FORMAT",
    default="'year'=YYYY/'month'=MM/'day'=dd/'hour'=HH",
    show_default=True,
    help=(
        "Pattern used to format the path in the S3 object name. Portion of "
        " the path generated by the S3 connector’s partitioner. The default "
        " is for an hourly partitioner. Alternatively set via the "
        "$KAFKA_CONNECT_S3_PATH_FORMAT env var."
    ),
)
@click.option(
    "--topic-regex",
    "topic_regex",
    envvar="KAFKA_CONNECT_TOPIC_REGEX",
    default=".*",
    show_default=True,
    help=(
        "Regex for selecting topics. Alternatively set via the "
        "$KAFKA_CONNECT_TOPIC_REGEX env var."
    ),
)
@click.option(
    "--dry-run",
    is_flag=True,
    help=(
        "Show the InfluxDB Sink Connector configuration but do not create "
        "the connector."
    ),
)
@click.option(
    "--auto-update",
    is_flag=True,
    help=(
        "Check for new topics and update the connector. See also the "
        "``--check-interval`` option."
    ),
)
@click.option(
    "-v",
    "--validate",
    is_flag=True,
    help="Validate the connector configuration before creating.",
)
@click.option(
    "-c",
    "--check-interval",
    "check_interval",
    envvar="KAFKA_CONNECT_CHECK_INTERVAL",
    default="15000",
    show_default=True,
    help=(
        "The interval, in milliseconds, to check for new topics and update"
        "the connector. Alternatively set via the "
        "$KAFKA_CONNECT_CHECK_INTERVAL env var."
    ),
)
@click.option(
    "-e",
    "--excluded-topics",
    "excluded_topics",
    envvar="KAFKA_CONNECT_EXCLUDED_TOPICS",
    default="",
    show_default=True,
    help=(
        "Comma separated list of topics to exclude from "
        "selection. Alternatively set via the "
        "$KAFKA_CONNECT_EXCLUDED_TOPICS env var."
    ),
)
@click.option(
    "--locale",
    "locale",
    envvar="KAFKA_CONNECT_S3_LOCALE",
    default="en-US",
    show_default=True,
    help=(
        "The locale to use when partitioning with `TimeBasedPartitioner`. "
        "Alternatively set via the $KAFKA_CONNECT_S3_LOCALE env var."
    ),
)
@click.option(
    "--timezone",
    "timezone",
    envvar="KAFKA_CONNECT_S3_TIMEZONE",
    default="UTC",
    show_default=True,
    help=(
        "The timezone to use when partitioning with `TimeBasedPartitioner`."
        "Alternatively set via the $KAFKA_CONNECT_S3_TIMEZONE"
    ),
)
@click.option(
    "--timestamp-extractor",
    "timestamp_extractor",
    envvar="KAFKA_CONNECT_TIMESTAMP_EXTRACTOR",
    default="Record",
    show_default=True,
    help=(
        "The extractor determines how to obtain a timestamp from each record. "
        "Values can be Wallclock to use the system time when "
        "the record is processed, Record to use the timestamp of the "
        "Kafka record denoting when it was produced or stored by the broker, "
        "RecordField to extract the timestamp from one of the fields in the "
        "record’s value as specified by the timestamp_field configuration "
        "property. Alternatively set via the "
        "$KAFKA_CONNECT_TIMESTAMP_EXTRACTOR env var."
    ),
)
@click.option(
    "--timestamp-field",
    "timestamp_field",
    envvar="KAFKA_CONNECT_TIMESTAMP_FIELD",
    default="time",
    show_default=True,
    help=(
        "The record field to be used as timestamp by the timestamp extractor. "
        "Only applies if timestamp_extractor is set to `RecordField`. "
        "Alternatively set via the $KAFKA_CONNECT_TIMESTAMP_FIELD env var."
    ),
)
@click.option(
    "-t",
    "--tasks-max",
    "tasks_max",
    envvar="KAFKA_CONNECT_TASKS_MAX",
    default="1",
    show_default=True,
    help=(
        "Number of Kafka Connect tasks. Alternatively set via the "
        "$KAFKA_CONNECT_TASKS_MAX env var."
    ),
)
@click.pass_context
def create_s3_sink(
    ctx: click.Context,
    topiclist: tuple,
    name: str,
    s3_bucket_name: str,
    s3_region: str,
    aws_access_key_id: str,
    aws_secret_access_key: str,
    topics_dir: str,
    s3_schema_compatibility: str,
    flush_size: str,
    rotate_interval_ms: str,
    partition_duration_ms: str,
    path_format: str,
    topic_regex: str,
    dry_run: bool,
    auto_update: bool,
    validate: bool,
    check_interval: str,
    excluded_topics: str,
    locale: str,
    timezone: str,
    timestamp_extractor: str,
    timestamp_field: str,
    tasks_max: str,
) -> int:
    """Create an instance of the S3 Sink connector.

    A list of topics can be specified using the TOPICLIST argument.
    If not, topics are discovered from Kafka. Use the ``--topic-regex`` and
    ``--excluded_topics`` options to help in selecting the topics
    that you want to write to S3. To check for new topics and update
    the connector configuration use the
    ``--auto-update`` and ``--check-interval`` options.
    """
    # Get configuration from the main command
    if ctx.parent:
        config = ctx.parent.obj["config"]
    # Connector configuration
    s3config = S3Config(
        name=name,
        s3_bucket_name=s3_bucket_name,
        s3_region=s3_region,
        aws_access_key_id=aws_access_key_id,
        aws_secret_access_key=aws_secret_access_key,
        topics_dir=topics_dir,
        s3_schema_compatibility=s3_schema_compatibility,
        flush_size=int(flush_size),
        rotate_interval_ms=int(rotate_interval_ms),
        partition_duration_ms=int(partition_duration_ms),
        path_format=path_format,
        locale=locale,
        timezone=timezone,
        timestamp_extractor=timestamp_extractor,
        timestamp_field=timestamp_field,
        tasks_max=int(tasks_max),
    )
    # The variadic argument is a tuple
    topics: List[str] = list(topiclist)
    if not topics:
        click.echo("Discoverying Kafka topics...")
        topics = Topic(config.broker_url, topic_regex, excluded_topics).names
        n = 0 if not topics else len(topics)
        click.echo(f"Found {n} topics.")
    connect = Connect(connect_url=config.connect_url)
    if topics:
        s3config.update_topics(topics)
        # --validate option
        if validate:
            click.echo(
                connect.validate(
                    name=s3config.connector_class,
                    connect_config=s3config.asjson(),
                )
            )
            return 0
        # --dry-run option returns the connector configuration
        if dry_run:
            click.echo(s3config.asjson())
            return 0
        # Validate configuration before creating the connector
        validation = connect.validate(
            name=s3config.connector_class, connect_config=s3config.asjson(),
        )
        try:
            error_count = json.loads(validation)["error_count"]
            click.echo(f"Validation returned {error_count} error(s).")
            if error_count > 0:
                click.echo(
                    "Use the ``--validate`` option to return the validation "
                    "results."
                )
                return 0
        except Exception:
            click.echo(validation)
            return 1
        click.echo(f"Uploading {name} connector configuration...")
        connect.create_or_update(name=name, connect_config=s3config.asjson())
    if auto_update:
        while True:
            time.sleep(int(check_interval) / 1000)
            try:
                # Current list of topics from Kafka
                current_topics = Topic(
                    config.broker_url, topic_regex, excluded_topics
                ).names
                new_topics = list(set(current_topics) - set(topics))
                if new_topics:
                    click.echo("Found new topics, updating the connector...")
                    s3config.update_topics(current_topics)
                    connect.create_or_update(
                        name=name, connect_config=s3config.asjson()
                    )
                    topics = current_topics
            except KeyboardInterrupt:
                raise click.ClickException("Interruped.")
    return 0
