# Copyright (c) 2025 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

"""
Module for managing Snowflake virtual warehouses within the CloudZero platform.

This module provides utilities to define, configure, and generate SQL commands for
Snowflake virtual warehouses with environment-specific settings and constraints.
"""

from typing_extensions import Unpack
from pyfaaster.aws import handlers_decorators_v2 as faaster

from dataclasses import dataclass, field
from enum import Enum
from typing import (
    Any,
    Dict,
    Optional,
    TypedDict,
)

from czc.aws_environment import Environment, default_namespace, environment
from czc.json import serializable

from .constants import DEFAULT_NAMESPACE
from .quoting import quoted_identifier, quoted_literal

DEFAULT_WAREHOUSE_NAME = "GENERAL"


class WarehouseError(Exception):
    """Base exception class for Snowflake warehouse-related errors in the CloudZero platform."""


class InvalidDevWarehouseSizeError(WarehouseError):
    """
    Exception raised when attempting to create a development warehouse with a size
    that exceeds the permitted limits for development environments.

    This helps enforce cost controls by preventing oversized warehouses in non-production
    environments.
    """

    def __init__(self, warehouse: "CloudZeroSnowflakeWarehouse") -> None:
        super().__init__(
            f"Invalid warehouse size {warehouse.warehouse_size.value} for {warehouse.name}"
        )


class InvalidProdWarehouseNamespaceError(WarehouseError):
    """
    Exception raised when attempting to create a production warehouse with an
    unauthorized namespace.

    This enforces namespace security controls for production environments.
    """

    def __init__(self, warehouse: "CloudZeroSnowflakeWarehouse") -> None:
        super().__init__(
            f"Invalid namespace {warehouse.namespace} for {warehouse.name}"
        )


class WarehouseSize(Enum):
    """
    Enumeration of Snowflake warehouse sizes from XSMALL to X6LARGE.

    These sizes determine compute resources and credit consumption rates.
    Larger warehouses provide more compute resources but consume credits at
    proportionally higher rates.

    See: https://docs.snowflake.com/en/sql-reference/sql/create-warehouse#optional-properties-objectproperties
    """

    XSMALL = "XSMALL"
    SMALL = "SMALL"
    MEDIUM = "MEDIUM"
    LARGE = "LARGE"
    XLARGE = "XLARGE"
    XXLARGE = "XXLARGE"
    XXXLARGE = "XXXLARGE"
    X4LARGE = "X4LARGE"
    X5LARGE = "X5LARGE"
    X6LARGE = "X6LARGE"

    def is_allowed_for_environment(self, env: Environment) -> bool:
        """
        Determines if a warehouse size is permitted in a given environment.

        Args:
            env: The environment to check permission against

        Returns:
            Boolean indicating whether the size is allowed in the specified environment
        """
        if env == Environment.dev:
            return self not in [
                WarehouseSize.XXXLARGE,
                WarehouseSize.X4LARGE,
                WarehouseSize.X5LARGE,
                WarehouseSize.X6LARGE,
            ]
        return True


class ScalingPolicy(Enum):
    """
    Enumeration of Snowflake warehouse scaling policies.

    STANDARD: Maximizes query throughput by aggressively scaling up resources
    ECONOMY: Minimizes credit usage by scaling more conservatively

    See: https://docs.snowflake.com/en/sql-reference/sql/create-warehouse#scaling-policy--standard--economy
    """

    STANDARD = "STANDARD"
    ECONOMY = "ECONOMY"


class WarehouseType(Enum):
    """
    Enumeration of Snowflake warehouse types.

    STANDARD: General-purpose warehouse for all query types
    SNOWFLAKE_OPTIMIZED: Specialized for Snowpark workloads with potentially
    different performance characteristics and pricing

    See: https://docs.snowflake.com/en/sql-reference/sql/create-warehouse#warehouse-type--standard--snowpark-optimized
    """

    STANDARD = "STANDARD"
    SNOWFLAKE_OPTIMIZED = "SNOWFLAKE-OPTIMIZED"


class WarehouseConfigurationKwargs(TypedDict, total=False):
    """
    Type definition for warehouse configuration parameters.

    These parameters correspond directly to Snowflake warehouse properties and
    affect warehouse behavior, scaling, resource utilization, and cost.

    All parameters are optional since they have default values in Snowflake.

    See: https://docs.snowflake.com/en/sql-reference/sql/create-warehouse#optional-properties-objectproperties
    """

    initially_suspended: bool
    min_cluster_count: int
    max_cluster_count: int
    auto_resume: bool
    auto_suspend: int
    scaling_policy: ScalingPolicy
    statement_timeout_in_seconds: int
    enable_query_acceleration: bool
    query_acceleration_max_scale_factor: int
    resource_monitor: str
    comment: str
    warehouse_type: WarehouseType
    max_concurrency_level: int
    statement_queued_timeout_in_seconds: int


@dataclass(frozen=True)
class CloudZeroSnowflakeWarehouse:
    """
    Represents a Snowflake virtual warehouse with CloudZero-specific configurations.

    This class encapsulates all properties needed to define and create a warehouse
    in Snowflake while enforcing CloudZero's environment-specific constraints, naming
    conventions, and security controls.
    """

    identifier: str
    environment: Environment
    namespace: str
    tags: dict = field(default_factory=dict)
    warehouse_size: WarehouseSize = WarehouseSize.MEDIUM
    initially_suspended: bool = True
    min_cluster_count: int = 1
    max_cluster_count: int = 3
    auto_resume: Optional[bool] = None
    auto_suspend: Optional[int] = None
    scaling_policy: Optional[ScalingPolicy] = None
    statement_timeout_in_seconds: Optional[int] = None
    enable_query_acceleration: Optional[bool] = None
    query_acceleration_max_scale_factor: Optional[int] = None
    resource_monitor: Optional[str] = None
    comment: Optional[str] = None
    warehouse_type: Optional[WarehouseType] = None
    max_concurrency_level: Optional[int] = None
    statement_queued_timeout_in_seconds: Optional[int] = None

    def __post_init__(self):
        """
        Validates warehouse configuration after initialization.

        Enforces organizational policies like:
        - Size restrictions for development environments
        - Namespace restrictions for production environments

        Raises:
            InvalidDevWarehouseSizeError: When a dev warehouse exceeds size limits (XXXLARGE and above)
            InvalidProdWarehouseNamespaceError: When a production warehouse uses an unauthorized namespace
        """
        if not self.warehouse_size.is_allowed_for_environment(self.environment):
            raise InvalidDevWarehouseSizeError(self)

        if self.environment == Environment.prod and self.namespace != "live":
            raise InvalidProdWarehouseNamespaceError(self)

    @classmethod
    @faaster.namespace_aware
    def default(
        cls,
        *,
        identifier: str = DEFAULT_WAREHOUSE_NAME,
        warehouse_size: WarehouseSize = WarehouseSize.MEDIUM,
        NAMESPACE: str = DEFAULT_NAMESPACE,
        **kwargs: Unpack[WarehouseConfigurationKwargs],
    ) -> "CloudZeroSnowflakeWarehouse":
        """
        Factory method to create a warehouse with default settings based on environment.

        This method simplifies warehouse creation by automatically determining the
        appropriate environment and namespace from the current context, while allowing
        customization of specific properties as needed.

        Args:
            identifier: Unique name for the warehouse (defaults to GENERAL)
            warehouse_size: Size of the warehouse (defaults to MEDIUM)
            NAMESPACE: Namespace to use (typically determined by faaster.namespace_aware)
            **kwargs: Additional warehouse configuration properties

        Returns:
            A properly configured CloudZeroSnowflakeWarehouse instance
        """
        env = environment(NAMESPACE)
        warehouse_namespace = default_namespace(env)
        return cls(
            identifier=identifier,
            environment=env,
            namespace=warehouse_namespace,
            warehouse_size=warehouse_size,
            **kwargs,
        )

    @property
    def name(self) -> str:
        """
        Generates the fully qualified warehouse name following CloudZero conventions.

        The naming convention <ENVIRONMENT>_<NAMESPACE>_<IDENTIFIER> provides clarity
        about the warehouse's purpose and environment, making it easier to identify
        and manage warehouses across different contexts.

        Returns:
            Snowflake warehouse name following CloudZero's naming convention
        """
        return f"{self.environment.value}_{self.namespace}_{self.identifier}".upper()

    def to_config_dict(self) -> Dict[str, Any]:
        """
        Converts warehouse configuration to a dictionary for SQL generation.

        Transforms the warehouse properties into the format required by Snowflake,
        filtering out null values and internal properties not relevant to Snowflake.

        Returns:
            Dictionary of Snowflake-compatible configuration properties
        """

        def is_warehouse_config(key: str) -> bool:
            return key not in ["identifier", "environment", "namespace", "tags"]

        return {
            k.upper(): v
            for k, v in serializable(self).items()
            if v is not None and is_warehouse_config(k)
        }

    def to_create_sql(self, *, force: bool = False) -> str:
        """
        Generates SQL to create this warehouse in Snowflake.

        Creates properly formatted SQL with appropriate quoting of identifiers
        and literals according to Snowflake's SQL syntax requirements.

        Args:
            force: If True, generates SQL to replace existing warehouse if present
                  If False, generates SQL that only creates if not exists

        Returns:
            SQL statement to create the warehouse

        See: https://docs.snowflake.com/en/sql-reference/sql/create-warehouse#syntax
        """
        warehouse_name = quoted_identifier(self.name)

        create_sql = (
            f"CREATE OR REPLACE WAREHOUSE {warehouse_name}"
            if force
            else f"CREATE WAREHOUSE IF NOT EXISTS {warehouse_name}"
        )

        config_parts = sorted(
            [f"{k} = {quoted_literal(v)}" for k, v in self.to_config_dict().items()]
        )
        config_fragment = f"WITH {' '.join(config_parts)}"

        return f"{create_sql}\n{config_fragment}"


def get_cluster():
    """
    Returns the Snowflake cluster name for the current environment.

    Uses a standardized naming convention that makes cluster association with
    environments clear and consistent.

    Returns:
        Formatted cluster name for the current environment
    """
    return f"snowflake-environment-{get_environment()}".lower()


@faaster.namespace_aware
def get_environment(*, NAMESPACE: str = DEFAULT_NAMESPACE) -> str:
    """
    Retrieves the current environment value based on the namespace.

    Args:
        NAMESPACE: Current namespace context, typically injected through NAMESPACE environment variable

    Returns:
        String representation of the current environment
    """
    return environment(NAMESPACE).value


def get_default_warehouse(*, warehouse_name: str = DEFAULT_WAREHOUSE_NAME) -> str:
    """
    Gets the fully qualified name of a default warehouse for the current context.

    Simplifies accessing the standard warehouse name when the default configuration
    is sufficient, eliminating the need to manually construct the full name.

    Args:
        warehouse_name: Base identifier for the warehouse

    Returns:
        Fully qualified warehouse name following CloudZero naming convention
    """
    return CloudZeroSnowflakeWarehouse.default(identifier=warehouse_name).name


def generate_warehouse_name(name: str, environment: str, namespace: str) -> str:
    """
    Generates a warehouse name for the specified parameters.

    Useful when you need to construct a warehouse name in a specific
    environment/namespace context different from the current one.

    Args:
        name: Base identifier for the warehouse
        environment: Target environment value
        namespace: Target namespace

    Returns:
        Fully qualified warehouse name following CloudZero naming convention
    """
    return CloudZeroSnowflakeWarehouse(name, Environment(environment), namespace).name
