import logging
import os
from typing import Dict, List, Literal, Optional

import certifi
from pydantic import Field, ValidationInfo, field_validator, model_validator

from datahub.configuration.common import (
    AllowDenyPattern,
    ConfigModel,
    HiddenFromDocs,
    TransparentSecretStr,
)
from datahub.configuration.source_common import (
    EnvConfigMixin,
    PlatformInstanceConfigMixin,
)
from datahub.configuration.time_window_config import BaseTimeWindowConfig
from datahub.ingestion.api.incremental_properties_helper import (
    IncrementalPropertiesConfigMixin,
)
from datahub.ingestion.source.ge_profiling_config import GEProfilingConfig
from datahub.ingestion.source.state.stale_entity_removal_handler import (
    StatefulStaleMetadataRemovalConfig,
)
from datahub.ingestion.source.state.stateful_ingestion_base import (
    StatefulIngestionConfigBase,
)
from datahub.ingestion.source.usage.usage_common import BaseUsageConfig
from datahub.ingestion.source_config.operation_config import is_profiling_enabled

logger = logging.getLogger(__name__)


class DremioConnectionConfig(ConfigModel):
    hostname: Optional[str] = Field(
        default=None,
        description="Hostname or IP Address of the Dremio server",
    )

    port: int = Field(
        default=9047,
        description="Port of the Dremio REST API",
    )

    username: Optional[str] = Field(
        default=None,
        description="Dremio username",
    )

    authentication_method: Optional[str] = Field(
        default="PAT",
        description="Authentication method: 'password' or 'PAT' (Personal Access Token)",
    )

    password: Optional[TransparentSecretStr] = Field(
        default=None,
        description="Dremio password or Personal Access Token",
    )

    tls: bool = Field(
        default=True,
        description="Whether the Dremio REST API port is encrypted",
    )

    disable_certificate_verification: Optional[bool] = Field(
        default=False,
        description="Disable TLS certificate verification",
    )

    path_to_certificates: str = Field(
        default=certifi.where(),
        description="Path to SSL certificates",
    )

    is_dremio_cloud: bool = Field(
        default=False,
        description="Whether this is a Dremio Cloud instance",
    )

    # https://docs.dremio.com/cloud/reference/api/#api-endpoint
    dremio_cloud_region: Literal["US", "EU"] = Field(
        default="US",
        description="Dremio Cloud region ('US' or 'EU')",
    )

    dremio_cloud_project_id: Optional[str] = Field(
        default=None,
        description="ID of Dremio Cloud Project. Found in Project Settings in the Dremio Cloud UI",
    )

    @field_validator("authentication_method", mode="after")
    @classmethod
    def validate_auth_method(cls, value: str) -> str:
        allowed_methods = ["password", "PAT"]
        if value not in allowed_methods:
            raise ValueError(
                f"authentication_method must be one of {allowed_methods}",
            )
        return value

    @field_validator("password", mode="after")
    @classmethod
    def validate_password(
        cls, value: Optional[TransparentSecretStr], info: ValidationInfo
    ) -> Optional[TransparentSecretStr]:
        if info.data.get("authentication_method") == "PAT" and not value:
            raise ValueError(
                "Password (Personal Access Token) is required when using PAT authentication",
            )
        return value


class ProfileConfig(GEProfilingConfig):
    query_timeout: int = Field(
        default=300, description="Time before cancelling Dremio profiling query"
    )
    include_field_median_value: HiddenFromDocs[bool] = Field(
        # Hidden because median causes a number of issues in Dremio.
        default=False,
    )


class DremioSourceMapping(EnvConfigMixin, PlatformInstanceConfigMixin, ConfigModel):
    platform: str = Field(
        description="Source connection made by Dremio (e.g. S3, Snowflake)",
    )
    source_name: str = Field(
        description="Alias of platform in Dremio connection",
    )


class DremioSourceTypeOverride(ConfigModel):
    platform: str = Field(
        description="DataHub platform name to emit for this Dremio source type (e.g. `kafka`, `iceberg`, `snowflake`).",
    )
    category: Optional[Literal["database", "file_object_storage"]] = Field(
        default=None,
        description="Whether the Dremio source uses dot-notation database/schema paths (`database`) or slash-notation file paths (`file_object_storage`). Defaults to `unknown`.",
    )


class DremioSourceConfig(
    DremioConnectionConfig,
    StatefulIngestionConfigBase,
    BaseTimeWindowConfig,
    IncrementalPropertiesConfigMixin,
    EnvConfigMixin,
    PlatformInstanceConfigMixin,
):
    # StatefulLineageConfigMixin / StatefulProfilingConfigMixin are
    # intentionally not inherited — their fields are unused on Dremio
    # (we use enable_stateful_time_window and profile_if_updated_since_days)
    # and their validators would emit warnings for dead config.

    domain: Optional[str] = Field(
        default=None,
        description="Domain for all source objects.",
    )

    include_system_tables: bool = Field(
        default=True,
        description="Whether to include system tables and schemas (INFORMATION_SCHEMA, SYS) in ingestion. ",
    )

    # Backward compatibility: usage config parameter (hidden from docs)
    usage: HiddenFromDocs[BaseUsageConfig] = Field(
        default=BaseUsageConfig(),
        description="Usage extraction configuration. For backward compatibility, this provides the same fields as the flattened time window config.",
    )

    source_mappings: Optional[List[DremioSourceMapping]] = Field(
        default=None,
        description="Mappings from Dremio sources to DataHub platforms and datasets.",
    )

    source_type_mappings: Dict[str, "DremioSourceTypeOverride"] = Field(
        default_factory=dict,
        description=(
            "Register additional Dremio source-type strings on top of the built-in mapping, "
            "for Dremio ARP connectors or source types not yet known to DataHub. "
            "Keys are the Dremio `type` field (case-insensitive); values declare the DataHub "
            "platform and optional category. "
            'Example: `{"MYORG_KAFKA": {"platform": "kafka", "category": "database"}}`.'
        ),
    )

    schema_pattern: AllowDenyPattern = Field(
        default=AllowDenyPattern.allow_all(),
        description="Regex patterns for schemas to filter",
    )

    dataset_pattern: AllowDenyPattern = Field(
        default=AllowDenyPattern.allow_all(),
        description="Regex patterns for tables and views to filter in ingestion. Specify regex to match the entire table name in dremio.schema.table format. e.g. to match all tables starting with customer in Customer database and public schema, use the regex 'dremio.public.customer.*'",
    )

    profile_pattern: AllowDenyPattern = Field(
        default=AllowDenyPattern.allow_all(),
        description="Regex patterns for tables to profile",
    )
    profiling: ProfileConfig = Field(
        default=ProfileConfig(),
        description="Configuration for profiling",
    )

    def is_profiling_enabled(self) -> bool:
        return self.profiling.enabled and is_profiling_enabled(
            self.profiling.operation_config
        )

    max_workers: int = Field(
        default=5 * (os.cpu_count() or 4),
        description="Number of worker threads to use for parallel processing",
    )

    include_query_lineage: bool = Field(
        default=False,
        description="Whether to include query-based lineage information.",
    )

    incremental_lineage: bool = Field(
        default=False,
        description="When enabled, lineage aspects are emitted as PATCH operations rather than full "
        "overwrites. This preserves any lineage edges that were manually added in DataHub "
        "between runs. Disable if you want each run to fully replace lineage. "
        "PATCH emission requires a GMS that supports patch aspects.",
    )

    enable_stateful_time_window: bool = Field(
        default=False,
        description="Enable stateful time window tracking for query lineage/usage extraction. "
        "When enabled, subsequent runs will skip time windows already fully processed, "
        "avoiding redundant API calls. Requires stateful_ingestion to be configured.",
    )

    stateful_ingestion: Optional[StatefulStaleMetadataRemovalConfig] = Field(
        default=None,
        description="Stateful ingestion config with stale-entity removal support.",
    )

    ingest_owner: bool = Field(
        default=True,
        description="Ingest Owner from source. This will override Owner info entered from UI",
    )

    @model_validator(mode="after")
    def _warn_if_stateful_time_window_without_stateful_ingestion(
        self,
    ) -> "DremioSourceConfig":
        # Warn when the user opts into time-window skipping but hasn't
        # enabled stateful_ingestion (the run-skip handler is only
        # installed in that case, so the flag would otherwise no-op).
        if self.enable_stateful_time_window and (
            self.stateful_ingestion is None or not self.stateful_ingestion.enabled
        ):
            logger.warning(
                "`enable_stateful_time_window: true` has no effect because "
                "`stateful_ingestion.enabled` is not set. Add a "
                "`stateful_ingestion:` block with `enabled: true` (and a "
                "configured `pipeline_name`) to activate run-to-run "
                "time-window skipping."
            )
        return self
