import dataclasses as dc
import inspect
import logging
import time
from typing import Optional, Sequence, Any, Type, TypeVar, Union, List

from databricks.sdk import WarehousesAPI
from databricks.sdk.errors import ResourceDoesNotExist
from databricks.sdk.service.sql import (
    State, EndpointInfo,
    EndpointTags, EndpointTagPair, EndpointInfoWarehouseType,
    GetWarehouseResponse, GetWarehouseResponseWarehouseType,
    Disposition, Format,
    ExecuteStatementRequestOnWaitTimeout, StatementParameterListItem,
    WarehouseAccessControlRequest, WarehousePermissionLevel
)

from .statement_result import StatementResult
from ..workspaces import Workspace, WorkspaceService
from yggdrasil.concurrent.threading import Job
from ...dataclasses.expiring import ExpiringDict
from ...pyutils.equality import dicts_equal
from yggdrasil.dataclasses.waiting import WaitingConfig, WaitingConfigArg

_CREATE_ARG_NAMES = {_ for _ in inspect.signature(WarehousesAPI.create).parameters.keys()}
_EDIT_ARG_NAMES = {_ for _ in inspect.signature(WarehousesAPI.edit).parameters.keys()}

__all__ = [
    "SQLWarehouse"
]


LOGGER = logging.getLogger(__name__)
CACHE_MAP: dict[str, ExpiringDict[str, "SQLWarehouse"]] = {}
T = TypeVar("T")


def safeGetWarehouseResponse(src: Union[GetWarehouseResponse, EndpointInfo]) -> GetWarehouseResponse:
    if isinstance(src, GetWarehouseResponse):
        return src

    payload = _copy_common_fields(src, GetWarehouseResponse, skip={"warehouse_type"})

    payload["warehouse_type"] = _safe_map_enum(
        GetWarehouseResponseWarehouseType,
        getattr(src, "warehouse_type", None),
    )

    return GetWarehouseResponse(**payload)


def safeEndpointInfo(src: Union[GetWarehouseResponse, EndpointInfo]) -> EndpointInfo:
    if isinstance(src, EndpointInfo):
        return src

    payload = _copy_common_fields(src, EndpointInfo, skip={"warehouse_type"})

    payload["warehouse_type"] = _safe_map_enum(
        EndpointInfoWarehouseType,
        getattr(src, "warehouse_type", None),
    )

    return EndpointInfo(**payload)


def set_cached_warehouse(
    workspace: Workspace,
    warehouse: "SQLWarehouse"
) -> None:
    host = workspace.safe_host
    existing = CACHE_MAP.get(host)

    if existing is None:
        existing = CACHE_MAP[host] = ExpiringDict(default_ttl=3600)

    existing[warehouse.warehouse_name] = warehouse


def get_cached_warehouse(
    workspace: Workspace,
    warehouse_name: str,
) -> Optional["SQLWarehouse"]:
    host = workspace.safe_host
    existing = CACHE_MAP.get(host)

    return existing.get(warehouse_name) if existing else None



def _safe_map_enum(dst_enum: Type[T], src_val: Any) -> Optional[T]:
    """
    Best-effort mapping:
      - if src_val is already a dst_enum member -> return it
      - try dst_enum(src_val) for value-based enums
      - try dst_enum(src_val.value) if src is enum-like
      - try dst_enum[src_val.name] if src is enum-like
      - try dst_enum[str(src_val)] as a last resort
    If nothing works -> None
    """
    if src_val is None:
        return None

    # already the right type
    if isinstance(src_val, dst_enum):
        return src_val

    # direct constructor (works if src_val is value-compatible)
    try:
        return dst_enum(src_val)  # type: ignore[misc]
    except Exception:
        pass

    # enum-like: .value
    try:
        return dst_enum(src_val.value)  # type: ignore[misc]
    except Exception:
        pass

    # enum-like: .name
    try:
        return dst_enum[src_val.name]  # type: ignore[index]
    except Exception:
        pass

    # string fallback
    try:
        return dst_enum(str(src_val))  # type: ignore[misc]
    except Exception:
        return None


def _copy_common_fields(src: Any, dst_cls: Type[T], *, skip: set[str] = frozenset()) -> dict:
    dst_field_names = {f.name for f in dc.fields(dst_cls)}
    payload = {}
    for name in dst_field_names:
        if name in skip:
            continue
        payload[name] = getattr(src, name, None)
    return payload


@dc.dataclass
class SQLWarehouse(WorkspaceService):
    warehouse_id: Optional[str] = None
    warehouse_name: Optional[str] = None

    _details: Optional[EndpointInfo] = dc.field(default=None, repr=False, hash=False, compare=False)

    def __post_init__(self):
        if self.warehouse_name and not self.warehouse_id:
            found = self.find_warehouse(warehouse_name=self.warehouse_name)

            self.warehouse_id = found.warehouse_id
            self.warehouse_name = found.warehouse_name
            self.details = found.details

    def client(self):
        return self.workspace.sdk().warehouses

    @property
    def details(self) -> EndpointInfo:
        if self._details is None:
            self.refresh()
        return self._details

    def latest_details(self):
        return self.client().get(id=self.warehouse_id)

    def refresh(self):
        self.details = self.latest_details()
        return self

    @details.setter
    def details(self, value: Union[GetWarehouseResponse, EndpointInfo]):
        self._details = safeEndpointInfo(value)

        if self._details is not None:
            self.warehouse_id = self._details.id
            self.warehouse_name = self._details.name

    @property
    def state(self):
        return self.latest_details().state

    @property
    def is_serverless(self):
        return self.details.enable_serverless_compute

    @property
    def is_running(self):
        return self.state == State.RUNNING

    @property
    def is_pending(self):
        return self.state in {
            State.DELETING, State.STARTING, State.STOPPING
        }

    def wait_for_status(
        self,
        wait: Optional[WaitingConfigArg] = None
    ):
        """
        Polls until not pending, using wait.sleep(iteration, start).

        WaitingConfig:
          - timeout: total wall-clock seconds (0 => no timeout)
          - interval: base sleep seconds (0 => busy/no-sleep)
          - backoff: exponential factor (>= 1)
          - max_interval: cap for sleep seconds (0 => no cap)

        Returns self.
        """
        wait = WaitingConfig.default() if wait is None else WaitingConfig.check_arg(wait)

        start = time.time()
        iteration = 0

        if wait.timeout:
            while self.is_pending:
                wait.sleep(iteration=iteration, start=start)
                iteration += 1

        return self

    def start(
        self,
        wait: WaitingConfigArg = True,
        raise_error: bool = True
    ):
        if self.warehouse_id:
            if not self.is_running:
                try:
                    response = self.client().start(id=self.warehouse_id)
                except Exception:
                    if raise_error:
                        raise
                    return self

                if wait:
                    wait = WaitingConfig.check_arg(wait)
                    response.result(timeout=wait.timeout_timedelta)

        return self

    def stop(self):
        if self.is_running:
            return self.client().stop(id=self.warehouse_id)
        return self

    def find_warehouse(
        self,
        warehouse_id: Optional[str] = None,
        warehouse_name: Optional[str] = None,
        find_default: bool = False,
        raise_error: bool = True,
    ):
        if warehouse_id:
            if warehouse_id == self.warehouse_id:
                return self

            return SQLWarehouse(
                workspace=self.workspace,
                warehouse_id=warehouse_id,
                warehouse_name=warehouse_name
            )
        elif warehouse_name:
            warehouse_name = warehouse_name or self.warehouse_name

            cached = get_cached_warehouse(
                workspace=self.workspace,
                warehouse_name=warehouse_name
            )

            if cached is not None:
                return cached

            for warehouse in self.list_warehouses():
                if warehouse.warehouse_name == warehouse_name:
                    set_cached_warehouse(
                        workspace=self.workspace,
                        warehouse=warehouse
                    )

                    if warehouse_name == self.warehouse_name:
                        self.warehouse_id = warehouse_id
                        return self

                    return warehouse

            if raise_error:
                raise ResourceDoesNotExist(
                    f"Cannot find SQL warehouse {warehouse_name!r}"
                )

            return None
        elif self.warehouse_id:
            return self
        elif self.warehouse_name:
            return self.find_warehouse(
                warehouse_name=self.warehouse_name,
                raise_error=raise_error,
                find_default=False
            )

        if find_default:
            return self.find_default(raise_error=raise_error)
        if raise_error:
            raise ResourceDoesNotExist(
                f"Cannot find SQL warehouse, no parameters given"
            )
        return None

    def find_default(self, raise_error: bool = True):
        classic = get_cached_warehouse(
            workspace=self.workspace,
            warehouse_name=DEFAULT_ALL_PURPOSE_CLASSIC_NAME
        )

        if classic is not None:
            if classic.state not in {State.RUNNING, State.STARTING, State.STOPPING}:
                classic.start(wait=False, raise_error=False)
            else:
                return classic

        serverless = get_cached_warehouse(
            workspace=self.workspace,
            warehouse_name=DEFAULT_ALL_PURPOSE_SERVERLESS_NAME
        )

        if serverless is not None:
            return serverless

        first_found = None

        for warehouse in self.list_warehouses():
            if first_found is None:
                first_found = warehouse

            if warehouse.warehouse_name == DEFAULT_ALL_PURPOSE_CLASSIC_NAME:
                classic = warehouse

                set_cached_warehouse(
                    workspace=self.workspace,
                    warehouse=classic
                )

                if classic._details.state == State.RUNNING:
                    return classic
                elif classic._details.state != State.STARTING:
                    classic.start(wait=False, raise_error=False)

                if serverless is not None:
                    return serverless
            elif warehouse.warehouse_name == DEFAULT_ALL_PURPOSE_SERVERLESS_NAME:
                serverless = warehouse

                set_cached_warehouse(
                    workspace=self.workspace,
                    warehouse=serverless
                )

                if classic is not None:
                    return serverless

        if serverless is None:
            try:
                created = self.create(
                    name=DEFAULT_ALL_PURPOSE_SERVERLESS_NAME,
                    permissions=["users"],
                )

                set_cached_warehouse(workspace=self.workspace, warehouse=created)

                return created
            except:
                pass

        if classic is None:
            Job.make(
                self.create,
                name=DEFAULT_ALL_PURPOSE_CLASSIC_NAME,
                permissions=["users"],
            ).fire_and_forget()

        if serverless is not None:
            return serverless

        if first_found is not None:
            return first_found

        if raise_error:
            raise ResourceDoesNotExist(
                f"Cannot find default SQL warehouse, no parameters given"
            )

        return self.create_or_update(
            name=self.warehouse_name or DEFAULT_ALL_PURPOSE_SERVERLESS_NAME,
            permissions=["users"]
        )

    def list_warehouses(self):
        for info in self.client().list():
            warehouse = SQLWarehouse(
                workspace=self.workspace,
                warehouse_id=info.id,
                warehouse_name=info.name,
                _details=info
            )

            yield warehouse

    def _check_details(
        self,
        keys: Sequence[str],
        update: bool,
        details: Optional[EndpointInfo] = None,
        **warehouse_specs
    ):
        if details is None:
            details = EndpointInfo(**{
                k: v
                for k, v in warehouse_specs.items()
                if k in keys
            })
        else:
            kwargs = {
                **details.as_shallow_dict(),
                **warehouse_specs
            }

            details = EndpointInfo(
                **{
                    k: v
                    for k, v in kwargs.items()
                    if k in keys
                },
            )

        if details.cluster_size is None:
            details.cluster_size = "X-Small"

        if details.warehouse_type is None:
            details.warehouse_type = EndpointInfoWarehouseType.PRO

        if not details.name:
            if details.enable_serverless_compute:
                details.name = DEFAULT_ALL_PURPOSE_SERVERLESS_NAME
            else:
                details.name = DEFAULT_ALL_PURPOSE_CLASSIC_NAME

        if details.enable_serverless_compute is None:
            if details.name == DEFAULT_ALL_PURPOSE_CLASSIC_NAME:
                details.enable_serverless_compute = False
            elif details.name == DEFAULT_ALL_PURPOSE_SERVERLESS_NAME:
                details.enable_serverless_compute = True
            else:
                details.enable_serverless_compute = "verless" in details.name.lower()

        if details.enable_serverless_compute:
            details.warehouse_type = EndpointInfoWarehouseType.PRO

        if not details.auto_stop_mins:
            if details.name == DEFAULT_ALL_PURPOSE_CLASSIC_NAME:
                details.auto_stop_mins = 30

        default_tags = self.workspace.default_tags(update=update)

        if details.tags is None:
            details.tags = EndpointTags(custom_tags=[
                EndpointTagPair(key=k, value=v)
                for k, v in default_tags.items()
            ])
        else:
            tags = {
                pair.key: pair.value
                for pair in details.tags.custom_tags
            }

            tags.update(default_tags)

        if details.tags is not None and not isinstance(details.tags, EndpointTags):
            details.tags = EndpointTags(custom_tags=[
                EndpointTagPair(key=k, value=v)
                for k, v in default_tags.items()
            ])

        if not details.max_num_clusters:
            details.max_num_clusters = 8

        return details

    def create_or_update(
        self,
        warehouse_id: Optional[str] = None,
        name: Optional[str] = None,
        permissions: Optional[List[WarehouseAccessControlRequest | str]] = None,
        wait: Optional[WaitingConfig] = None,
        **warehouse_specs
    ):
        name = name or self.warehouse_name
        found = self.find_warehouse(
            warehouse_id=warehouse_id,
            warehouse_name=name,
            raise_error=False,
        )

        if found is not None:
            return found.update(
                name=name,
                permissions=permissions,
                wait=wait,
                **warehouse_specs
            )

        return self.create(
            name=name,
            permissions=permissions,
            wait=wait,
            **warehouse_specs
        )

    def create(
        self,
        name: Optional[str] = None,
        *,
        permissions: Optional[List[WarehouseAccessControlRequest | str]] = None,
        wait: Optional[WaitingConfigArg] = None,
        **warehouse_specs
    ):
        name = name or self.warehouse_name

        wait = WaitingConfig.check_arg(wait)

        details = self._check_details(
            keys=_CREATE_ARG_NAMES,
            update=False,
            name=name,
            **warehouse_specs
        )

        update_details = {
            k: v
            for k, v in details.as_shallow_dict().items()
            if k in _CREATE_ARG_NAMES
        }

        if wait or permissions:
            info = self.client().create_and_wait(
                timeout=wait.timeout_timedelta,
                **update_details
            )
        else:
            info = self.client().create(**update_details)

            update_details["id"] = info.response.id

            info = EndpointInfo(**update_details)

        self.update_permissions(permissions=permissions, wait=wait)

        created = SQLWarehouse(
            workspace=self.workspace,
            warehouse_id=info.id,
            warehouse_name=info.name,
            _details=info
        )

        return created

    def update(
        self,
        wait: Optional[WaitingConfigArg] = None,
        permissions: Optional[List[WarehouseAccessControlRequest | str]] = None,
        **warehouse_specs
    ):
        if not warehouse_specs:
            return self

        wait = WaitingConfig.check_arg(wait)

        existing_details = {
            k: v
            for k, v in self.details.as_shallow_dict().items()
            if k in _EDIT_ARG_NAMES
        }

        update_details = {
            k: v
            for k, v in (
                self._check_details(
                    details=self.details,
                    update=True,
                    keys=_EDIT_ARG_NAMES,
                    **warehouse_specs
                )
                .as_shallow_dict()
                .items()
            )
            if k in _EDIT_ARG_NAMES
        }

        same = dicts_equal(
            existing_details,
            update_details,
            keys=_EDIT_ARG_NAMES,
        )

        if not same:
            LOGGER.debug(
                "Updating %s with %s",
                self, update_details
            )

            if wait.timeout:
                self.details = self.client().edit_and_wait(
                    timeout=wait.timeout_timedelta,
                    **update_details
                )
            else:
                _ = self.client().edit(**update_details)

                self.details = EndpointInfo(**update_details)

        self.update_permissions(permissions=permissions, wait=wait)

        LOGGER.info(
            "Updated %s",
            self
        )

        return self

    def update_permissions(
        self,
        permissions: Optional[List[WarehouseAccessControlRequest | str]] = None,
        *,
        wait: WaitingConfigArg = True,
        warehouse_id: Optional[str] = None
    ):
        warehouse_id = warehouse_id or self.warehouse_id

        permissions = [
            self.check_permission(_)
            for _ in permissions
        ] if permissions else []

        if permissions and warehouse_id:
            if wait:
                self.client().update_permissions(
                    warehouse_id=warehouse_id,
                    access_control_list=permissions
                )
            else:
                Job.make(
                    self.client().update_permissions,
                    warehouse_id=warehouse_id,
                    access_control_list=permissions
                ).fire_and_forget()

        return self

    @staticmethod
    def check_permission(
        permission: WarehouseAccessControlRequest | str
    ):
        if isinstance(permission, str):
            if permission == "users":
                return WarehouseAccessControlRequest(
                    group_name=permission,
                    permission_level=WarehousePermissionLevel.CAN_USE
                )
            elif "@" in permission:
                return WarehouseAccessControlRequest(
                    user_name=permission,
                    permission_level=WarehousePermissionLevel.CAN_USE
                )
            else:
                return WarehouseAccessControlRequest(
                    group_name=permission,
                    permission_level=WarehousePermissionLevel.CAN_MANAGE
                )

        return permission

    def delete(self):
        if self.warehouse_id:
            LOGGER.debug(
                "Deleting %s",
                self
            )

            self.client().delete(id=self.warehouse_id)

            LOGGER.info(
                "Deleted %s",
                self
            )

    def execute(
        self,
        statement: Optional[str] = None,
        *,
        warehouse_id: Optional[str] = None,
        warehouse_name: Optional[str] = None,
        byte_limit: Optional[int] = None,
        disposition: Optional[Disposition] = None,
        format: Optional[Format] = None,
        on_wait_timeout: Optional[ExecuteStatementRequestOnWaitTimeout] = None,
        parameters: Optional[List[StatementParameterListItem]] = None,
        row_limit: Optional[int] = None,
        wait_timeout: Optional[str] = None,
        catalog_name: Optional[str] = None,
        schema_name: Optional[str] = None,
        wait: WaitingConfigArg = True
    ) -> StatementResult:
        """Execute a SQL statement via Spark or Databricks SQL Statement Execution API.

        Engine resolution:
        - If `engine` is not provided and a Spark session is active -> uses Spark.
        - Otherwise uses Databricks SQL API (warehouse).

        Waiting behavior (`wait_result`):
        - If True (default): returns a StatementResult in terminal state (SUCCEEDED/FAILED/CANCELED).
        - If False: returns immediately with the initial handle (caller can `.wait()` later).

        Args:
            statement: SQL statement to execute. If None, a `SELECT *` is generated from the table params.
            warehouse_id: Warehouse override (for API engine).
            warehouse_name: Warehouse name override (for API engine).
            byte_limit: Optional byte limit for results.
            disposition: Result disposition mode (API engine).
            format: Result format (API engine).
            on_wait_timeout: Timeout behavior for waiting (API engine).
            parameters: Optional statement parameters (API engine).
            row_limit: Optional row limit for results (API engine).
            wait_timeout: API wait timeout value.
            catalog_name: Optional catalog override for API engine.
            schema_name: Optional schema override for API engine.
            wait: Whether to block until completion (API engine).

        Returns:
            StatementResult
        """
        if format is None:
            format = Format.ARROW_STREAM

        if disposition is None:
            disposition = Disposition.EXTERNAL_LINKS
        elif format in (Format.CSV, Format.ARROW_STREAM):
            disposition = Disposition.EXTERNAL_LINKS

        instance = self.find_warehouse(warehouse_id=warehouse_id, warehouse_name=warehouse_name)
        warehouse_id = warehouse_id or instance.warehouse_id
        workspace_client = instance.workspace.sdk()

        LOGGER.debug(
            "API SQL executing query:\n%s",
            statement
        )

        response = workspace_client.statement_execution.execute_statement(
            statement=statement,
            warehouse_id=warehouse_id,
            byte_limit=byte_limit,
            disposition=disposition,
            format=format,
            on_wait_timeout=on_wait_timeout,
            parameters=parameters,
            row_limit=row_limit,
            wait_timeout=wait_timeout,
            catalog=catalog_name,
            schema=schema_name,
        )

        execution = StatementResult(
            workspace_client=workspace_client,
            warehouse_id=warehouse_id,
            statement_id=response.statement_id,
            disposition=disposition,
            _response=response,
        )

        LOGGER.info(
            "API SQL executed %s",
            execution
        )

        return execution.wait(wait=wait)


DEFAULT_ALL_PURPOSE_CLASSIC_NAME = "Yggdrasil All Purpose"
DEFAULT_ALL_PURPOSE_SERVERLESS_NAME = DEFAULT_ALL_PURPOSE_CLASSIC_NAME + " Serverless"
