import asyncio
import json
import logging
import sys
from argparse import ArgumentParser, Namespace
from pathlib import Path

from lisa.globalfit.framework.args import add_bus_args, add_common_args
from lisa.globalfit.framework.buses import DefaultPullEventBus
from lisa.globalfit.framework.pipelinerunner.workflow import (
    WorkflowSpec,
)

logger = logging.getLogger(__name__)

INBOX = "infra.pipelinerunner.submit"


def get_args() -> Namespace:
    parser = ArgumentParser()

    add_common_args(parser)
    add_bus_args(parser)

    parser.add_argument(
        "workflow",
        type=str,
    )

    parser.add_argument(
        "--labels",
        type=str,
        nargs="+",
    )

    parser.add_argument(
        "--parameters",
        "-p",
        action="extend",
        nargs="+",
        default=[],
        help="workflow parameters as key=value pairs",
    )
    parser.add_argument(
        "--parameter-file",
        type=Path,
        help="JSON file containing key-value workflow parameters",
    )

    return parser.parse_args()


def parse_workflow_arguments(
    parameters: list[str], parameter_file: Path | None
) -> dict[str, str]:
    params = dict(param.split("=") for param in parameters)
    if parameter_file is not None:
        with open(parameter_file) as f:
            file_params = json.load(f)
        for k, v in file_params.items():
            if k not in params:
                params[k] = v
    logger.debug(f"resolved parameters: {params}")
    return params


async def main_async():
    args = get_args()
    logging.basicConfig(
        level=args.log_level,
        format="%(asctime)s.%(msecs)03d [%(levelname)s] %(name)s %(message)s",
        datefmt=r"%Y-%m-%dT%H:%M:%S",
    )

    if args.workdir is None:
        logger.error("missing --workdir argument")
        sys.exit(1)

    try:
        params = parse_workflow_arguments(args.parameters, args.parameter_file)
    except FileNotFoundError as err:
        logger.error(err)
        sys.exit(1)

    bus = DefaultPullEventBus(servers=args.servers)
    spec = WorkflowSpec(
        resource=args.workflow,
        labels=args.labels if args.labels is not None else [],
        parameters={"workdir": args.workdir, **params},
    )
    await bus.connect()
    await bus.publish(INBOX, spec.encode())
    logger.info(f"published {spec}")


def main():
    return asyncio.run(main_async())


if __name__ == "__main__":
    main()
