# Copyright (c) CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com

from ..snowflake import AsyncQueryStatus, check_async_query
import pyfaaster.aws.handlers_decorators_v2 as faaster
from typing import Callable, Optional
from aws_lambda_powertools.utilities.parser import event_parser

from .types import (
    AsyncInput,
    AsyncQueryHandlerOutput,
    AsyncQueryOperation,
    QueryContext,
    SyncResult,
)
from . import state_machine
from czc.logging import get_logger
import time

logger = get_logger()


def async_query(pause: Optional[int] = None):
    """
    Triggers a new state machine to monitor a long-running query.
    Wraps a function which returns `AsyncQueryOperation`

    Args:
        pause: If set, the decorator will wait for the given number of seconds before triggering the async state machine.
               If the query is completed in that time, then the results are returned without the machine triggering.
    """

    def decorated(
        handler: Callable[..., AsyncQueryHandlerOutput],
    ) -> Callable[..., None]:
        @faaster.environ_aware(required=["ASYNC_QUERY_STM_ARN"])
        @event_parser(model=AsyncInput)
        def handler_wrapper(
            event: AsyncInput,
            *args,
            ASYNC_QUERY_STM_ARN: str,
            **kwargs,
        ):
            logger.append_keys(organization_id=event.organization_id)
            full_event = event.model_dump()

            query_info = handler(full_event, *args, **kwargs)
            if query_info is None:
                logger.info("No query to monitor")
                state_machine.send_success(
                    event.task_token,
                    {
                        "output": {},
                        "query_id": None,
                    },
                )
                return

            if isinstance(query_info, SyncResult):
                logger.info("No query to monitor")
                state_machine.send_success(
                    event.task_token,
                    {
                        "output": query_info.task_output
                        if query_info.task_output
                        else {},
                        "query_id": None,
                    },
                )
                return

            logger.append_keys(context={**query_info.context.model_dump()})
            if pause is not None:
                logger.info(f"Pause for {pause} seconds")
                time.sleep(pause)

            # check if we can return without running the other state machine
            status = check_async_query(
                query_info.snowflake_connection_info.create_connection(),
                query_info.query_id,
            )
            if status == AsyncQueryStatus.success:
                logger.info("Query complete.  Not triggering async state machine")
                state_machine.send_success(
                    event.task_token,
                    {
                        "output": query_info.task_output
                        if query_info.task_output
                        else {},
                        "query_id": query_info.query_id,
                    },
                )
                return

            logger.info("Starting async query state machine")
            async_query_operation_data = {
                **query_info.model_dump(),
                "token": event.task_token,
                # update output to contain the query_id so that downstream steps can retrieve results
                "task_output": {
                    "output": query_info.task_output,
                    "query_id": query_info.query_id,
                },
                "context": query_info.context
                if query_info.context
                else QueryContext(
                    organization_id=event.organization_id,
                    description=f"Run query {query_info.query_id.query_id}",
                ),
            }
            arn = state_machine.start(
                ASYNC_QUERY_STM_ARN, AsyncQueryOperation(**async_query_operation_data)
            )
            logger.info("State machine started", async_state_machine_arn=arn)

        return handler_wrapper

    return decorated
