"""Amazon Redshift ODBC connection configuration."""

import logging

from dataclasses import dataclass
from enum import Enum
from typing import Any

from pydantic import Field, PrivateAttr, model_validator

from data_exchange_agent.config.sections.connections.odbc.base import (
    BaseODBCConnectionConfig,
)
from data_exchange_agent.constants.connection_types import ConnectionType


logger = logging.getLogger(__name__)


# Connection string templates for Redshift ODBC
REDSHIFT_CONNECTION_STRING = "DRIVER={{{driver}}};" "SERVER={server};" "PORT={port};" "DATABASE={database};" "UID={user};" "PWD={password}"

# Default Redshift port
DEFAULT_REDSHIFT_PORT = 5439

# Default ODBC driver for Redshift
DEFAULT_REDSHIFT_DRIVER = "Amazon Redshift (x64)"


class AuthMethod(str, Enum):
    """Authentication methods for Redshift connections."""

    STANDARD = "standard"
    IAM_PROVISIONED = "iam-provisioned-cluster"
    IAM_SERVERLESS = "iam-serverless"


@dataclass(kw_only=True)
class RedshiftCredentials:
    """Base class for Redshift credentials."""

    database: str
    username: str
    port: int = DEFAULT_REDSHIFT_PORT
    odbc_driver: str | None = None
    auto_detect_driver: bool = True


@dataclass(kw_only=True)
class StandardCredentials(RedshiftCredentials):
    """Standard username/password authentication credentials."""

    host: str
    password: str


@dataclass(kw_only=True)
class IAMProvisionedCredentials(RedshiftCredentials):
    """IAM authentication credentials for provisioned Redshift clusters."""

    cluster_id: str
    region: str
    access_key_id: str
    secret_access_key: str
    session_token: str | None = None
    host: str | None = None  # Auto-resolved if not provided


@dataclass(kw_only=True)
class IAMServerlessCredentials(RedshiftCredentials):
    """IAM authentication credentials for Redshift Serverless."""

    workgroup_name: str
    region: str
    access_key_id: str
    secret_access_key: str
    session_token: str | None = None
    host: str | None = None  # Auto-resolved if not provided


@dataclass(kw_only=True)
class UnloadConfig:
    """
    Configuration for Redshift UNLOAD operations.

    When configured, the DEA can execute UNLOAD commands using these settings.
    The orchestrator sends a SELECT query and the DEA wraps it in an UNLOAD command.
    """

    s3_bucket: str
    """S3 bucket name for UNLOAD output (e.g., 'my-bucket')."""

    iam_role_arn: str
    """IAM role ARN that Redshift assumes for S3 access."""

    s3_prefix: str | None = None
    """Optional S3 prefix/path within the bucket (e.g., 'data-migration/prod')."""

    file_format: str = "parquet"
    """Output file format: 'parquet' or 'csv'."""

    max_file_size_mb: int = 256
    """Maximum file size in MB for UNLOAD output."""

    def build_s3_path(self, relative_path: str) -> str:
        """
        Build full S3 path from relative path.

        Args:
            relative_path: Relative path (e.g., "task_results/workflows/1/tables/2/data/partition_0/")

        Returns:
            Full S3 path (e.g., "s3://bucket/prefix/task_results/.../partition_0/")

        """
        bucket = self.s3_bucket.rstrip("/")
        rel_path = relative_path.strip("/")

        # Build path parts: bucket, optional prefix, relative path
        path_parts = [bucket]
        if self.s3_prefix:
            path_parts.append(self.s3_prefix.strip("/"))
        path_parts.append(rel_path)

        return f"s3://{'/'.join(path_parts)}/"

    def build_unload_query(self, select_query: str, s3_path: str) -> str:
        """
        Build a Redshift UNLOAD query.

        Args:
            select_query: The SELECT query to unload
            s3_path: Full S3 path for output

        Returns:
            Complete UNLOAD SQL statement

        """
        if self.file_format.lower() == "parquet":
            format_options = "FORMAT AS PARQUET"
        else:
            format_options = "FORMAT AS CSV DELIMITER ',' ADDQUOTES HEADER"

        return f"""UNLOAD ('{select_query.replace("'", "''")}')
TO '{s3_path}'
IAM_ROLE '{self.iam_role_arn}'
{format_options}
ALLOWOVERWRITE
PARALLEL ON
MAXFILESIZE {self.max_file_size_mb} MB"""


class RedshiftConnectionConfig(BaseODBCConnectionConfig):
    """
    Configuration class for Amazon Redshift ODBC connection settings.

    Supports multiple authentication modes:
    1. Standard Authentication: Uses username and password directly
    2. IAM Authentication (Provisioned Cluster): Uses AWS credentials to get temporary password
    3. IAM Authentication (Serverless): Uses AWS credentials with workgroup

    For IAM Authentication, the system will use boto3 to generate temporary credentials.

    References:
        - https://docs.aws.amazon.com/redshift/latest/mgmt/generating-user-credentials.html
        - https://docs.aws.amazon.com/redshift/latest/mgmt/options-for-providing-iam-credentials.html

    """

    # Override parent fields with Redshift-specific defaults
    # Make database and host optional since they can come from credentials
    driver_name: str = Field(default=ConnectionType.REDSHIFT, description="Driver name")
    database: str = Field(default="", description="Database name")
    host: str = Field(default="", description="Database host address")
    port: int = Field(default=DEFAULT_REDSHIFT_PORT, ge=1, le=65535, description="Database port number")

    # Redshift uses credentials object internally but we need to handle both cases
    # The credentials are resolved during initialization via from_config_dict factory
    credentials: StandardCredentials | IAMProvisionedCredentials | IAMServerlessCredentials | None = Field(
        default=None,
        exclude=True,  # Don't include in serialization
        description="Credentials model for the specific authentication method",
    )
    unload_config: UnloadConfig | None = Field(
        default=None,
        exclude=True,  # Don't include in serialization
        description="Optional configuration for UNLOAD operations",
    )

    # Private attributes for internal state
    _resolved_odbc_driver: str = PrivateAttr()
    _auth_method: AuthMethod = PrivateAttr()

    @property
    def auth_method(self) -> AuthMethod:
        """Get the authentication method."""
        return self._auth_method

    @model_validator(mode="before")
    @classmethod
    def populate_fields_from_credentials(cls, data: Any) -> Any:
        """Populate connection fields from credentials before validation."""
        if not isinstance(data, dict):
            return data

        credentials = data.get("credentials")
        if credentials is None:
            return data

        # Populate database if not set
        if not data.get("database") and hasattr(credentials, "database"):
            data["database"] = credentials.database

        # Populate username if not set
        if not data.get("username") and hasattr(credentials, "username"):
            data["username"] = credentials.username

        # Populate port if still at default and credentials has different port
        if data.get("port", DEFAULT_REDSHIFT_PORT) == DEFAULT_REDSHIFT_PORT:
            if hasattr(credentials, "port") and credentials.port != DEFAULT_REDSHIFT_PORT:
                data["port"] = credentials.port

        # Standard credentials also provide host and password
        if isinstance(credentials, StandardCredentials):
            if not data.get("host"):
                data["host"] = credentials.host
            if not data.get("password"):
                data["password"] = credentials.password

        return data

    def model_post_init(self, __context: Any) -> None:
        """Post-initialization setup for private attributes."""
        if self.credentials is None:
            self._auth_method = AuthMethod.STANDARD
            self._resolved_odbc_driver = self._select_odbc_driver(None, True)
        else:
            # Set auth method based on credentials type
            if isinstance(self.credentials, StandardCredentials):
                self._auth_method = AuthMethod.STANDARD
            elif isinstance(self.credentials, IAMProvisionedCredentials):
                self._auth_method = AuthMethod.IAM_PROVISIONED
            elif isinstance(self.credentials, IAMServerlessCredentials):
                self._auth_method = AuthMethod.IAM_SERVERLESS
            else:
                self._auth_method = AuthMethod.STANDARD

            # Select ODBC driver
            self._resolved_odbc_driver = self._select_odbc_driver(
                self.credentials.odbc_driver,
                self.credentials.auto_detect_driver,
            )

    @model_validator(mode="after")
    def validate_redshift_config(self) -> "RedshiftConnectionConfig":
        """Validate the Redshift connection configuration."""
        if self.credentials is not None:
            self._validate_credentials()
        return self

    def _validate_credentials(self) -> None:
        """Validate credentials based on type."""
        if isinstance(self.credentials, StandardCredentials):
            if not self.credentials.host:
                raise ValueError("Host is required for standard authentication.")
            if not self.credentials.password:
                raise ValueError("Password is required for standard authentication.")

        elif isinstance(self.credentials, IAMProvisionedCredentials):
            if not self.credentials.cluster_id:
                raise ValueError("cluster_id is required for IAM provisioned cluster authentication.")
            if not self.credentials.region:
                raise ValueError("region is required for IAM authentication.")
            if not self.credentials.access_key_id:
                raise ValueError("access_key_id is required for IAM authentication.")
            if not self.credentials.secret_access_key:
                raise ValueError("secret_access_key is required for IAM authentication.")

        elif isinstance(self.credentials, IAMServerlessCredentials):
            if not self.credentials.workgroup_name:
                raise ValueError("workgroup_name is required for IAM serverless authentication.")
            if not self.credentials.region:
                raise ValueError("region is required for IAM authentication.")
            if not self.credentials.access_key_id:
                raise ValueError("access_key_id is required for IAM authentication.")
            if not self.credentials.secret_access_key:
                raise ValueError("secret_access_key is required for IAM authentication.")

    @classmethod
    def from_config_dict(
        cls,
        database: str,
        username: str,
        auth_method: str | AuthMethod = AuthMethod.STANDARD,
        # Common optional parameters
        port: int = DEFAULT_REDSHIFT_PORT,
        odbc_driver: str | None = None,
        auto_detect_driver: bool = True,
        # Standard auth parameters
        host: str | None = None,
        password: str | None = None,
        # IAM Provisioned Cluster parameters
        cluster_id: str | None = None,
        # IAM Serverless parameters
        workgroup_name: str | None = None,
        # IAM common parameters
        region: str | None = None,
        access_key_id: str | None = None,
        secret_access_key: str | None = None,
        session_token: str | None = None,
        # UNLOAD configuration parameters (optional)
        unload_s3_bucket: str | None = None,
        unload_iam_role_arn: str | None = None,
        unload_s3_prefix: str | None = None,
        unload_file_format: str = "parquet",
        unload_max_file_size_mb: int = 256,
        **extra_options: str | int | float | bool,
    ) -> "RedshiftConnectionConfig":
        """
        Create a RedshiftConnectionConfig from a flat configuration dictionary.

        This factory method is used by the TOML configuration loader to construct
        a RedshiftConnectionConfig from flat parameters read from the configuration file.

        Args:
            database: Database name
            username: Username for the database
            auth_method: Authentication method (AuthMethod enum or string:
                "standard", "iam-provisioned-cluster", or "iam-serverless")
            port: Port number (default: 5439)
            odbc_driver: ODBC driver name (optional)
            auto_detect_driver: Whether to auto-detect the ODBC driver (default: True)
            host: Hostname (required for standard auth, optional for IAM)
            password: Password (required for standard auth)
            cluster_id: Cluster ID (required for IAM provisioned cluster)
            workgroup_name: Workgroup name (required for IAM serverless)
            region: AWS region (required for IAM auth)
            access_key_id: AWS access key ID (required for IAM auth)
            secret_access_key: AWS secret access key (required for IAM auth)
            session_token: AWS session token (optional for IAM auth)
            unload_s3_bucket: S3 bucket for UNLOAD output (optional)
            unload_iam_role_arn: IAM role ARN for UNLOAD (optional)
            unload_s3_prefix: S3 prefix/path within the bucket for UNLOAD (optional)
            unload_file_format: File format for UNLOAD (default: "parquet")
            unload_max_file_size_mb: Max file size for UNLOAD (default: 256)
            **extra_options: Extra options to add to the ODBC connection string

        Returns:
            RedshiftConnectionConfig: Configured connection instance

        Raises:
            ValueError: If required parameters for the selected auth method are missing

        """
        # Normalize auth_method to enum (accepts both string and enum)
        if isinstance(auth_method, str):
            try:
                auth_method_enum = AuthMethod(auth_method)
            except ValueError as err:
                raise ValueError(
                    f"Invalid auth_method '{auth_method}'. " f"Must be one of: {', '.join([e.value for e in AuthMethod])}"
                ) from err
        else:
            auth_method_enum = auth_method

        # Construct the appropriate credentials object based on auth_method
        credentials: StandardCredentials | IAMProvisionedCredentials | IAMServerlessCredentials

        if auth_method_enum == AuthMethod.STANDARD:
            if not host:
                raise ValueError("'host' is required for standard authentication")
            if not password:
                raise ValueError("'password' is required for standard authentication")
            credentials = StandardCredentials(
                database=database,
                username=username,
                host=host,
                password=password,
                port=port,
                odbc_driver=odbc_driver,
                auto_detect_driver=auto_detect_driver,
            )
        elif auth_method_enum == AuthMethod.IAM_PROVISIONED:
            if not cluster_id:
                raise ValueError("'cluster_id' is required for IAM provisioned cluster authentication")
            if not region:
                raise ValueError("'region' is required for IAM authentication")
            if not access_key_id:
                raise ValueError("'access_key_id' is required for IAM authentication")
            if not secret_access_key:
                raise ValueError("'secret_access_key' is required for IAM authentication")
            credentials = IAMProvisionedCredentials(
                database=database,
                username=username,
                cluster_id=cluster_id,
                region=region,
                access_key_id=access_key_id,
                secret_access_key=secret_access_key,
                session_token=session_token,
                host=host,
                port=port,
                odbc_driver=odbc_driver,
                auto_detect_driver=auto_detect_driver,
            )
        elif auth_method_enum == AuthMethod.IAM_SERVERLESS:
            if not workgroup_name:
                raise ValueError("'workgroup_name' is required for IAM serverless authentication")
            if not region:
                raise ValueError("'region' is required for IAM authentication")
            if not access_key_id:
                raise ValueError("'access_key_id' is required for IAM authentication")
            if not secret_access_key:
                raise ValueError("'secret_access_key' is required for IAM authentication")
            credentials = IAMServerlessCredentials(
                database=database,
                username=username,
                workgroup_name=workgroup_name,
                region=region,
                access_key_id=access_key_id,
                secret_access_key=secret_access_key,
                session_token=session_token,
                host=host,
                port=port,
                odbc_driver=odbc_driver,
                auto_detect_driver=auto_detect_driver,
            )

        # Create UnloadConfig if UNLOAD parameters are provided
        unload_config: UnloadConfig | None = None
        if unload_s3_bucket and unload_iam_role_arn:
            unload_config = UnloadConfig(
                s3_bucket=unload_s3_bucket,
                iam_role_arn=unload_iam_role_arn,
                s3_prefix=unload_s3_prefix,
                file_format=unload_file_format,
                max_file_size_mb=unload_max_file_size_mb,
            )
            logger.info(
                "UNLOAD configuration enabled: s3_bucket=%s, s3_prefix=%s, iam_role=%s",
                unload_s3_bucket,
                unload_s3_prefix,
                unload_iam_role_arn,
            )

        # Resolve host, username, and password based on authentication method
        resolved_host: str
        resolved_username: str
        resolved_password: str

        if isinstance(credentials, StandardCredentials):
            resolved_host = credentials.host
            resolved_username = credentials.username
            resolved_password = credentials.password

        elif isinstance(credentials, (IAMProvisionedCredentials | IAMServerlessCredentials)):
            if isinstance(credentials, IAMProvisionedCredentials):
                logger.info(
                    "Using IAM authentication for Redshift (method: %s)",
                    auth_method_enum.value,
                )
                logger.info(
                    "Cluster ID: %s, Region: %s",
                    credentials.cluster_id,
                    credentials.region,
                )
            else:
                logger.info(
                    "Using IAM authentication for Redshift (method: %s)",
                    auth_method_enum.value,
                )
                logger.info(
                    "Workgroup: %s, Region: %s",
                    credentials.workgroup_name,
                    credentials.region,
                )

            # Get IAM credentials - we need a temporary instance for this
            temp_instance = cls.__new__(cls)
            resolved_host, resolved_username, resolved_password = temp_instance._get_iam_credentials(credentials)
            logger.info(
                "IAM authentication successful. Host: %s, IAM User: %s",
                resolved_host,
                resolved_username,
            )
            logger.debug("Temporary password obtained from AWS")

        else:
            raise ValueError(f"Unsupported credentials type: {type(credentials)}")

        # Validate required fields
        if not resolved_host:
            raise ValueError("Host is required. For IAM auth, ensure cluster_id/workgroup_name and region are provided.")

        return cls(
            driver_name=ConnectionType.REDSHIFT,
            username=resolved_username,
            password=resolved_password,
            database=database,
            host=resolved_host,
            port=port,
            extra_options=dict(extra_options),
            credentials=credentials,
            unload_config=unload_config,
        )

    def _get_iam_credentials(self, credentials: IAMProvisionedCredentials | IAMServerlessCredentials) -> tuple[str, str, str]:
        """
        Get temporary credentials using IAM authentication.

        Args:
            credentials: IAM credentials model

        Returns:
            tuple[str, str, str]: (host, iam_username, temporary_password)

        Raises:
            ValueError: If required IAM parameters are missing
            Exception: If AWS API calls fail

        """
        # Lazy import boto3 to avoid loading AWS SDK at module import time.
        # This improves startup performance and reduces memory footprint.
        # boto3 is only loaded when IAM authentication is actually used.
        import boto3

        from botocore.exceptions import ClientError

        # Create boto3 session with credentials
        session_kwargs = {
            "aws_access_key_id": credentials.access_key_id,
            "aws_secret_access_key": credentials.secret_access_key,
            "region_name": credentials.region,
        }
        if credentials.session_token:
            session_kwargs["aws_session_token"] = credentials.session_token

        session = boto3.Session(**session_kwargs)

        try:
            if isinstance(credentials, IAMProvisionedCredentials):
                return self._get_provisioned_cluster_credentials(session, credentials)
            elif isinstance(credentials, IAMServerlessCredentials):
                return self._get_serverless_credentials(session, credentials)
            else:
                raise ValueError(f"Unsupported IAM credentials type: {type(credentials)}")

        except ClientError as e:
            error_code = e.response.get("Error", {}).get("Code", "Unknown")
            error_message = e.response.get("Error", {}).get("Message", str(e))
            raise ConnectionError(f"Failed to get Redshift IAM credentials. AWS Error ({error_code}): {error_message}") from e

    def _get_provisioned_cluster_credentials(self, session, credentials: IAMProvisionedCredentials) -> tuple[str, str, str]:
        """
        Get credentials for a provisioned Redshift cluster.

        Args:
            session: boto3 Session with AWS credentials
            credentials: IAM provisioned credentials model

        Returns:
            tuple[str, str, str]: (host, iam_username, temporary_password)

        """
        redshift_client = session.client("redshift")

        # Get cluster endpoint (use provided host or auto-resolve)
        if credentials.host:
            host = credentials.host
            logger.info("Using provided cluster endpoint: %s", host)
        else:
            logger.info(
                "Calling AWS Redshift DescribeClusters for cluster: %s",
                credentials.cluster_id,
            )
            clusters = redshift_client.describe_clusters(ClusterIdentifier=credentials.cluster_id)
            cluster = clusters["Clusters"][0]
            host = cluster["Endpoint"]["Address"]
            logger.info("Resolved cluster endpoint: %s", host)

        # Get temporary credentials
        logger.info(
            "Calling AWS Redshift GetClusterCredentials for user: %s",
            credentials.username,
        )
        response = redshift_client.get_cluster_credentials(
            ClusterIdentifier=credentials.cluster_id,
            DbUser=credentials.username,
            DbName=credentials.database,
            DurationSeconds=3600,  # 1 hour
        )

        # AWS returns a modified username (often with IAM: prefix) that must be used
        iam_username = response["DbUser"]
        temporary_password = response["DbPassword"]

        logger.info(
            "IAM credentials obtained successfully. IAM User: %s (original: %s)",
            iam_username,
            credentials.username,
        )
        logger.debug("Password length: %s", len(temporary_password))

        return host, iam_username, temporary_password

    def _get_serverless_credentials(self, session, credentials: IAMServerlessCredentials) -> tuple[str, str, str]:
        """
        Get credentials for Redshift Serverless.

        Args:
            session: boto3 Session with AWS credentials
            credentials: IAM serverless credentials model

        Returns:
            tuple[str, str, str]: (host, iam_username, temporary_password)

        """
        redshift_serverless_client = session.client("redshift-serverless")

        # Get workgroup endpoint (use provided host or auto-resolve)
        if credentials.host:
            host = credentials.host
            logger.info("Using provided workgroup endpoint: %s", host)
        else:
            logger.info(
                "Calling AWS Redshift Serverless GetWorkgroup for workgroup: %s",
                credentials.workgroup_name,
            )
            workgroup = redshift_serverless_client.get_workgroup(workgroupName=credentials.workgroup_name)
            host = workgroup["workgroup"]["endpoint"]["address"]
            logger.info("Resolved workgroup endpoint: %s", host)

        # Get temporary credentials
        response = redshift_serverless_client.get_credentials(
            workgroupName=credentials.workgroup_name,
            dbName=credentials.database,
        )

        # AWS returns a modified username that must be used
        iam_username = response.get("dbUser", credentials.username)
        temporary_password = response["dbPassword"]

        logger.info(
            "IAM credentials obtained successfully. IAM User: %s",
            iam_username,
        )
        logger.debug("Password length: %s", len(temporary_password))

        return host, iam_username, temporary_password

    def build_connection_string(self) -> str:
        """
        Build the ODBC connection string for Redshift.

        Returns:
            str: The ODBC connection string

        """
        connection_string = REDSHIFT_CONNECTION_STRING.format(
            driver=self._resolved_odbc_driver,
            server=self.host,
            port=self.port,
            database=self.database,
            user=self.username,
            password=self.password,
        )

        # Add extra options
        for key, value in self.extra_options.items():
            stripped_key = str(key).strip()
            stripped_value = str(value).strip()
            connection_string += f";{stripped_key}={stripped_value}"

        return connection_string

    def _select_odbc_driver(self, requested_driver: str | None, auto_detect: bool) -> str:
        """
        Select the appropriate ODBC driver.

        Args:
            requested_driver: The driver requested by the user (can be None)
            auto_detect: Whether to validate driver availability

        Returns:
            str: The ODBC driver name to use

        """
        # If auto-detect is disabled, just use the requested driver or default
        if not auto_detect:
            driver = requested_driver or DEFAULT_REDSHIFT_DRIVER
            logger.debug("Using Redshift ODBC driver without validation: %s", driver)
            return driver

        # If auto-detect is enabled, validate the driver
        if requested_driver:
            # Check if requested driver is available
            from shared.odbc.driver_utils import is_driver_available

            if is_driver_available(requested_driver):
                logger.info("Using requested Redshift ODBC driver: %s", requested_driver)
                return requested_driver
            else:
                logger.warning(
                    "Requested Redshift ODBC driver '%s' not found. Using default: %s",
                    requested_driver,
                    DEFAULT_REDSHIFT_DRIVER,
                )
                return DEFAULT_REDSHIFT_DRIVER
        else:
            logger.info(
                "No ODBC driver specified. Using default Redshift driver: %s",
                DEFAULT_REDSHIFT_DRIVER,
            )
            return DEFAULT_REDSHIFT_DRIVER

    def _repr_fields(self) -> str:
        """Get string representation of fields (for use in __repr__)."""
        base_repr = super()._repr_fields()
        auth_details = f"auth_method='{self._auth_method.value}'"

        if isinstance(self.credentials, IAMProvisionedCredentials):
            auth_details += f", cluster_id='{self.credentials.cluster_id}', region='{self.credentials.region}'"
        elif isinstance(self.credentials, IAMServerlessCredentials):
            auth_details += f", workgroup_name='{self.credentials.workgroup_name}', " f"region='{self.credentials.region}'"

        return f"{auth_details}, {base_repr}"
