import dataclasses
import logging

import click
import yaml
from loculus_client import (
    Config,
    get_or_create_group_and_return_group_id,
    get_submitted,
    regroup_and_revoke,
    submit_or_revise,
)

logger = logging.getLogger(__name__)
logging.basicConfig(
    encoding="utf-8",
    level=logging.INFO,
    format="%(asctime)s %(levelname)8s %(filename)15s%(mode)s - %(message)s ",
    datefmt="%H:%M:%S",
)


@click.command()
@click.option(
    "--metadata",
    required=False,
    type=click.Path(exists=True),
)
@click.option(
    "--sequences",
    required=False,
    type=click.Path(exists=True),
)
@click.option(
    "--mode",
    required=True,
    type=click.Choice(["submit", "revise", "regroup-and-revoke", "get-submitted"]),
)
@click.option(
    "--log-level",
    default="INFO",
    type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]),
)
@click.option(
    "--config-file",
    required=True,
    type=click.Path(exists=True),
)
@click.option(
    "--output",
    required=False,
    type=click.Path(),
)
@click.option(
    "--revoke-map",
    required=False,
    type=click.Path(exists=True),
)
def submit_to_loculus(metadata, sequences, mode, log_level, config_file, output, revoke_map):
    """
    Submit data to Loculus.
    """
    logger.setLevel(log_level)
    logging.getLogger("requests").setLevel(logging.WARNING)
    logging.getLogger("urllib3").setLevel(logging.WARNING)

    old_factory = logging.getLogRecordFactory()

    def record_factory(*args, **kwargs):
        record = old_factory(*args, **kwargs)
        record.mode = f":{mode}"
        return record

    logging.setLogRecordFactory(record_factory)

    with open(config_file, encoding="utf-8") as file:
        full_config = yaml.safe_load(file)
        relevant_config = {f.name: full_config.get(f.name, []) for f in dataclasses.fields(Config)}
        config = Config(**relevant_config)

    logger.info(f"Config: {config}")

    if mode in {"submit", "revise"}:
        logger.info(f"Starting {mode}")
        try:
            group_id = get_or_create_group_and_return_group_id(
                config, allow_creation=mode == "submit"
            )
        except ValueError as e:
            logger.error(f"Aborting {mode} due to error: {e}")
            return
        response = submit_or_revise(metadata, sequences, config, group_id, mode=mode)
        logger.info(f"Completed {mode}")

    if mode == "regroup-and-revoke":
        try:
            group_id = get_or_create_group_and_return_group_id(
                config, allow_creation=mode == "submit"
            )
        except ValueError as e:
            logger.error(f"Aborting {mode} due to error: {e}")
            return
        logger.info("Submitting new segment groups and revoking old segment groups")
        response = regroup_and_revoke(metadata, sequences, revoke_map, config, group_id)
        logger.info(f"Revoked: {len(response)} sequence entries of old segment groups")

    if mode == "get-submitted":
        logger.info("Getting submitted sequences")
        if config.segmented:
            insdc_key = [
                "insdcAccessionBase" + "_" + segment for segment in config.nucleotide_sequences
            ]
        else:
            insdc_key = ["insdcAccessionBase"]

        fields = ["hash", *insdc_key]
        get_submitted(config, output, fields)


if __name__ == "__main__":
    submit_to_loculus()
