# 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 functools import lru_cache
from .snowflake import AsyncQueryId, AsyncQueryStatus, check_async_query, get_default_connection
from typing import Any, Callable, Optional
from aws_lambda_powertools.utilities.parser import event_parser

from czc.logging import get_logger
import time


from pydantic import BaseModel, Field, ConfigDict, field_validator
from typing import Literal, Union
from . import snowflake
from czc import json, aws_environment


import boto3

logger = get_logger()


MAX_NAME_LENGTH = 80


class AsyncInput(BaseModel):
    """Minimum required input to handlers decorated with @async_query"""

    organization_id: Optional[str] = None

    task_token: str = Field(min_length=1)

    model_config = ConfigDict(extra='allow')


class QueryContext(BaseModel):
    """Additional information for logging purposes"""

    organization_id: Optional[str] = None
    """The organization ID of the customer that the query is related to, if relevant"""

    connection_id: Optional[str] = None
    """The connection ID of the customer that the query is related to, if relevant"""

    description: str = ''
    """A description of the query"""

    additional_context: dict = Field(default_factory=dict)
    """Additional context that will be logged as this task is processed"""

    model_config = ConfigDict(extra='allow')


class SyncResult(BaseModel):
    """The return value when we do not want to run the async query state machine"""

    output: Optional[dict] = None
    """The output of this step"""

    query_id: Optional[AsyncQueryId] = None
    """The ID of the snowflake query that has being executed."""



class AsyncResult(BaseModel):
    """The return value of an async query handler"""

    query_id: AsyncQueryId
    """The ID of the snowflake query that is being executed."""

    context: QueryContext = QueryContext()
    """Additional information for logging purposes"""

    output: Optional[dict] = None
    """The desired output of the task"""

    @property
    def serialized_output(self) -> str:
        return json.dumps(self.output or {})

    @field_validator('query_id', mode='before')
    @classmethod
    def convert_query_id(cls, value: Any) -> Any:
        if isinstance(value, (AsyncQueryId, str)):
            try:
                return AsyncQueryId.parse(value)
            except Exception:
                return value
        return value


AsyncQueryHandlerOutput = Union[SyncResult, AsyncResult, Literal[None]]


class AsyncQueryOperation(AsyncResult):
    """The event that triggers the AsyncQuery state machine"""

    token: Optional[str] = None
    """The token that represents this task. Task tokens are generated by Step Functions when tasks are assigned to a
    worker, or in the context object when a workflow enters a task state. See GetActivityTaskOutput$taskToken."""


class AsyncQueryOperationOutput(BaseModel):
    """The output of the AsyncQuery state machine"""

    status: snowflake.AsyncQueryStatus
    """The status of the query"""


class AsyncQueryOutput(BaseModel):
    """The output of a step decorated by @async_query"""

    query_id: Optional[snowflake.AsyncQueryId]
    """The query ID"""

    output: Optional[dict] = None
    """Any output set from the original async handler"""


def async_query(
    pause: Optional[int] = None,
) -> Callable[[Callable[..., AsyncQueryHandlerOutput]], Callable[..., 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]:
        @event_parser(model=AsyncInput)
        def handler_wrapper(
            event: AsyncInput,
            *args,
            **kwargs,
        ):
            if event.organization_id:
                logger.append_keys(organization_id=event.organization_id)
            full_event = event.model_dump()

            try:
                query_info = handler(full_event, *args, **kwargs)
            except Exception as err:
                logger.exception('Unhandled exception encountered')
                send_failure(event.task_token, err)
                raise

            if query_info is None:
                logger.info('No query to monitor')
                send_success(
                    event.task_token,
                    {
                        'output': {},
                        'query_id': None,
                    },
                )
                return

            if isinstance(query_info, SyncResult):
                logger.info('No query to monitor')
                send_success(
                    event.task_token,
                    {
                        'output': query_info.output if query_info.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(
                get_default_connection(readonly=False, purpose='AsyncQuery Status Check from decorator'),
                query_info.query_id,
            )
            if status == AsyncQueryStatus.success:
                logger.info('Query complete. Not triggering async state machine')
                send_success(
                    event.task_token,
                    {
                        'output': query_info.output if query_info.output else {},
                        'query_id': query_info.query_id,
                    },
                )
                return

            account_info = aws_environment.resolve_account_info()
            async_query_arn = f'arn:aws:states:us-east-1:{account_info.core_account_id}:stateMachine:cz-{account_info.default_namespace}-snowflake-async-query-stm'
            logger.info('Starting async query state machine', async_query_arn=async_query_arn)
            arn = start_state_machine(
                async_query_arn,
                AsyncQueryOperation(
                    **{
                        **query_info.model_dump(),
                        'token': event.task_token,
                        # update output to contain the query_id so that downstream steps can retrieve results
                        'output': {
                            'output': query_info.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}',
                        ),
                    }
                ),
            )
            logger.info('State machine started', async_state_machine_arn=arn)

        return handler_wrapper

    return decorated


@lru_cache(maxsize=1)
def _get_step_functions_client():
    return boto3.client('stepfunctions')


def start_state_machine(stm_arn: str, operation: AsyncQueryOperation) -> Optional[str]:
    name_parts = []
    if operation.context.organization_id is not None:
        name_parts.append(f'org_{operation.context.organization_id[0:8]}')
    if operation.context.connection_id is not None:
        name_parts.append(f'conn_{operation.context.connection_id[0:8]}')
    name_parts.append(operation.query_id.query_id)

    name = '.'.join(name_parts)[0:MAX_NAME_LENGTH]

    logger.info(f'Starting state machine.  Name: {name}')
    response = _get_step_functions_client().start_execution(stateMachineArn=stm_arn, name=name, input=json.dumps(operation))
    return response['executionArn']


def send_success(token: str, output: Optional[dict]):
    output_data = json.dumps(output or {})
    _get_step_functions_client().send_task_success(taskToken=token, output=output_data)


def send_failure(token: str, err: Exception):
    _get_step_functions_client().send_task_failure(taskToken=token, error=type(err).__name__, cause=str(err))
