import asyncio
import logging
from argparse import ArgumentParser, Namespace
from asyncio import Event, Task
from signal import SIGINT, SIGTERM

from lisa.globalfit.framework.args import add_bus_args, add_logging_args
from lisa.globalfit.framework.bus import PullEventBus
from lisa.globalfit.framework.buses import DefaultPullEventBus
from lisa.globalfit.framework.exceptions import DeserializationError
from lisa.globalfit.framework.msg.streams import STREAM_INFRA
from lisa.globalfit.framework.pipelinerunner.backends import argo
from lisa.globalfit.framework.pipelinerunner.workflow import (
    WorkflowBackendDetails,
    WorkflowResult,
    WorkflowSpec,
)

logger = logging.getLogger(__name__)

SUBJECT_SUBMIT = "infra.pipelinerunner.submit"
SUBJECT_RESULT = "infra.pipelinerunner.result"
SUBJECT_NOTIFICATION = "infra.pipelinerunner.notification"
SUBJECT_SHUTDOWN = "infra.workflow-submitter.control.shutdown"


class WorkflowSubmitter:
    CONSUMER_NAME = "workflow-submitter"
    MSG_FETCH_TIMEOUT_S = 5

    def __init__(self, bus: PullEventBus, options: WorkflowBackendDetails) -> None:
        self.bus = bus
        self.should_stop = Event()
        self.options = options

    async def setup(self) -> None:
        # TODO Handle bus disconnects after setup.
        await self.bus.connect()
        await self.bus.add_stream(STREAM_INFRA)

    async def run(self):
        logger.info("starting workflow submitter")
        tasks = (
            Task(self.handle_msg_submit()),
            Task(self.handle_msg_result()),
            Task(self.handle_msg_shutdown()),
        )
        await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED)

    async def handle_msg_submit(self) -> None:
        sub = await self.bus.subscribe(SUBJECT_SUBMIT, f"{self.CONSUMER_NAME}-sumbit")
        while not self.should_stop.is_set():
            msg = await sub.fetch_message_with_timeout(self.MSG_FETCH_TIMEOUT_S)
            if msg is None:
                continue
            await msg.in_progress()
            try:
                spec = WorkflowSpec.decode(msg.data)
                await self.submit_workflow(spec)
                await msg.ack()
            except DeserializationError as err:
                logger.error(f"got invalid workflow submission (discarding): {err}")
                await msg.discard()
            except KeyError as err:
                logger.error(f"could not submit workflow (retrying): {err}")
                await msg.nack(delay=5)
            except ValueError as err:
                logger.error(f"could not submit workflow (discarding): {err}")
                await msg.discard()
        logger.debug(f"unsubscribing from {SUBJECT_SUBMIT!r}")
        await sub.unsubscribe()

    async def handle_msg_result(self) -> None:
        sub = await self.bus.subscribe(
            SUBJECT_NOTIFICATION, f"{self.CONSUMER_NAME}-result"
        )
        while not self.should_stop.is_set():
            msg = await sub.fetch_message_with_timeout(self.MSG_FETCH_TIMEOUT_S)
            if msg is None:
                continue
            await msg.in_progress()
            name, status, duration = msg.data.decode().split()
            logger.info(
                f"workflow {name} ended with status {status} ({float(duration):.1f}s)"
            )
            result = WorkflowResult(name, status, self.options, [])
            await self.bus.publish(SUBJECT_RESULT, result.encode())
            await msg.ack()
        logger.debug(f"unsubscribing from {SUBJECT_NOTIFICATION!r}")
        await sub.unsubscribe()

    async def handle_msg_shutdown(self) -> None:
        shutdown_sub = await self.bus.subscribe(
            SUBJECT_SHUTDOWN,
            f"{self.CONSUMER_NAME}-shutdown",
        )
        while not self.should_stop.is_set():
            _ = await shutdown_sub.fetch_message()
            # TODO Check message contents.
            logger.debug("setting stop flag")
            self.should_stop.set()
        logger.debug(f"unsubscribing from {SUBJECT_SHUTDOWN!r}")
        await shutdown_sub.unsubscribe()

    async def submit_workflow(self, spec: WorkflowSpec) -> None:
        await argo.submit_workflow(spec, self.options)

    async def shutdown(self) -> None:
        logger.info("shutting down DAG submitter")
        await self.bus.close()


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

    add_logging_args(parser)
    add_bus_args(parser)

    parser.add_argument("--url", type=str, required=True)

    return parser.parse_args()


async def run_submitter(args: Namespace) -> None:
    bus = DefaultPullEventBus(servers=args.servers)
    options = WorkflowBackendDetails(url=args.url)
    submitter = WorkflowSubmitter(bus, options)
    try:
        await submitter.setup()
        await submitter.run()
    except asyncio.CancelledError:
        pass
    await submitter.shutdown()


def main():
    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",
    )

    loop = asyncio.get_event_loop()
    if logging.root.level == logging.DEBUG:
        loop.set_debug(True)
    task = asyncio.ensure_future(run_submitter(args))

    for signal in [SIGINT, SIGTERM]:
        loop.add_signal_handler(signal, task.cancel)

    loop.run_until_complete(task)


if __name__ == "__main__":
    main()
