# 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 pydantic import BaseModel, Field, ConfigDict
from typing import Literal, Optional, Union
from .. import snowflake
from czc import json


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

    organization_id: str = Field(min_length=1)
    task_token: str = Field(min_length=1)

    model_config = ConfigDict(extra="allow")


class SnowflakeConnectionInfo(BaseModel):
    """Information required to connect to Snowflake"""

    cluster_identifier: str
    db_user: str = "cz_readonly"
    db_name: Optional[str] = None
    purpose: Optional[str] = None
    warehouse: Optional[str] = None

    def create_connection(self) -> snowflake.SnowflakeConnection:
        return snowflake.connect(
            cluster_identifier=self.cluster_identifier,
            db_user=self.db_user,
            db_name=self.db_name,
            purpose=self.purpose,
            warehouse=self.warehouse,
        )


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"""


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

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


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

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

    snowflake_connection_info: SnowflakeConnectionInfo
    """Information necessary to connect to the snowflake cluster where the query is being run"""

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

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

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

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: snowflake.AsyncQueryId
    """The query ID"""

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