"""
Databricks Warehouse Adapter.

Provides pre-warming capabilities for Databricks serverless SQL warehouses
to eliminate cold-start delays before dbt command execution.
"""

import asyncio
import logging
from pathlib import Path
from typing import Any, Callable

import requests
import yaml

logger = logging.getLogger(__name__)


class DatabricksProfileError(Exception):
    """Raised when Databricks profile configuration is invalid or missing."""

    pass


class DatabricksWarehouseAdapter:
    """
    Warehouse adapter for Databricks serverless SQL warehouses.

    This adapter pre-warms Databricks warehouses by starting them via API
    and polling until they reach RUNNING state. This eliminates the ~30s
    cold-start delay that dbt would otherwise experience.

    The adapter reads connection info from dbt profiles (profiles.yml) and
    uses the Databricks SQL Warehouses API for control operations.
    """

    def __init__(self, project_dir: Path):
        """
        Initialize Databricks warehouse adapter.

        Args:
            project_dir: Path to the dbt project directory
        """
        self.project_dir = project_dir
        self._connection_info = None  # Lazy-loaded connection details
        self._is_running = False  # Track if we've already started the warehouse

    async def prewarm(self, progress_callback: Callable[[int, int, str], Any] | None = None) -> None:
        """
        Pre-warm the Databricks serverless warehouse.

        Starts the warehouse if not already running and waits for it to reach
        RUNNING state. This operation is idempotent - calling it multiple times
        is safe and won't cause issues.

        Args:
            progress_callback: Optional callback for progress updates (current, total, message)

        Raises:
            DatabricksProfileError: If profile configuration is invalid
            RuntimeError: If warehouse fails to start or times out
        """
        # If we've already started the warehouse in this session, skip
        if self._is_running:
            logger.debug("Warehouse already pre-warmed in this session, skipping")
            return

        logger.info("Pre-warming Databricks serverless warehouse...")

        # Report initial progress
        logger.info(f"Progress callback is: {progress_callback}")
        if progress_callback:
            logger.info("Invoking initial progress callback: 'Initializing warehouse...'")
            try:
                result = progress_callback(0, 1, "Initializing warehouse...")
                logger.info(f"Progress callback result type: {type(result)}")
                if asyncio.iscoroutine(result):
                    await result
                logger.info("Initial progress callback completed")
            except Exception as e:
                logger.warning(f"Progress callback error: {e}")
        else:
            logger.warning("No progress callback provided to prewarm")

        # Get connection info from dbt profile
        try:
            instance, token, warehouse_id = await self._get_connection_info()
        except DatabricksProfileError as e:
            logger.error(f"Failed to get Databricks connection info: {e}")
            raise

        headers = {"Authorization": f"Bearer {token}"}
        warehouse_url = f"https://{instance}/api/2.0/sql/warehouses/{warehouse_id}"

        # Check current warehouse state
        try:
            resp = await asyncio.to_thread(requests.get, warehouse_url, headers=headers)
            if resp.status_code != 200:
                raise RuntimeError(f"Failed to fetch warehouse info: {resp.text}")

            warehouse = resp.json()
            current_state = warehouse.get("state")
            warehouse_name = warehouse.get("name", warehouse_id)

            logger.info(f"Warehouse '{warehouse_name}' current state: {current_state}")

            # If already running, we're done
            if current_state == "RUNNING":
                logger.info("Warehouse already running, no pre-warming needed")
                self._is_running = True
                return

            # Verify it's a serverless warehouse
            is_serverless = warehouse.get("warehouse_type") == "PRO" and warehouse.get("enable_serverless_compute", False)
            if not is_serverless:
                logger.warning("Warehouse is not serverless, pre-warming may not be beneficial")

            # Start the warehouse
            start_url = f"https://{instance}/api/2.0/sql/warehouses/{warehouse_id}/start"
            start_resp = await asyncio.to_thread(requests.post, start_url, headers=headers)

            if start_resp.status_code != 200:
                raise RuntimeError(f"Failed to start warehouse: {start_resp.text}")

            logger.info(f"Started warehouse '{warehouse_name}', waiting for RUNNING state...")

            # Poll for RUNNING state with progress reporting
            max_wait = 300  # 5 minutes
            poll_interval = 5  # seconds
            waited = 0

            while waited < max_wait:
                await asyncio.sleep(poll_interval)
                waited += poll_interval

                # Report progress
                if progress_callback:
                    try:
                        result = progress_callback(waited, max_wait, f"Pre-warming warehouse '{warehouse_name}'... ({waited}s)")
                        if asyncio.iscoroutine(result):
                            await result
                    except Exception as e:
                        logger.warning(f"Progress callback error: {e}")

                state_resp = await asyncio.to_thread(requests.get, warehouse_url, headers=headers)
                if state_resp.status_code == 200:
                    state = state_resp.json().get("state")
                    logger.info(f"Warehouse state after {waited}s: {state}")

                    if state == "RUNNING":
                        logger.info(f"Warehouse is RUNNING after {waited}s")
                        # Final progress update
                        if progress_callback:
                            try:
                                result = progress_callback(max_wait, max_wait, f"Warehouse '{warehouse_name}' ready")
                                if asyncio.iscoroutine(result):
                                    await result
                            except Exception as e:
                                logger.warning(f"Progress callback error: {e}")
                        self._is_running = True
                        return
                else:
                    logger.warning(f"Failed to get warehouse state: {state_resp.text}")

            # Timeout
            raise RuntimeError(f"Timed out waiting for warehouse to start after {max_wait}s")

        except requests.RequestException as e:
            logger.error(f"Network error during warehouse pre-warming: {e}")
            raise RuntimeError(f"Failed to pre-warm warehouse: {e}")

    async def _get_connection_info(self) -> tuple[str, str, str]:
        """
        Extract Databricks connection info from dbt profile.

        Returns:
            Tuple of (instance, token, warehouse_id)

        Raises:
            DatabricksProfileError: If required configuration is missing
        """
        if self._connection_info:
            return self._connection_info

        # Get dbt profile configuration
        profile = await self._get_dbt_profile()

        # Extract connection details
        instance = profile.get("host", "").replace("https://", "").replace("/", "")
        token = profile.get("token")
        warehouse_id = profile.get("http_path", "").replace("/sql/1.0/warehouses/", "")

        if not instance or not token:
            raise DatabricksProfileError("Could not find Databricks instance or token in dbt profile")

        if not warehouse_id:
            raise DatabricksProfileError("No warehouse ID found in dbt profile config (http_path)")

        logger.debug(f"Using Databricks instance: {instance}, warehouse: {warehouse_id}")

        # Cache for future calls
        self._connection_info = (instance, token, warehouse_id)
        return self._connection_info

    async def _get_dbt_profile(self) -> dict[str, Any]:
        """
        Load dbt profile configuration from profiles.yml.

        Searches for profiles.yml in:
        1. Project directory (profiles.yml)
        2. User home directory (~/.dbt/profiles.yml)

        Returns:
            Dictionary with dbt profile target configuration

        Raises:
            DatabricksProfileError: If profiles not found or invalid
        """
        # First check project directory for profiles.yml
        local_profiles_path = self.project_dir / "profiles.yml"
        if local_profiles_path.exists():
            profiles_path = local_profiles_path
            logger.debug(f"Using local profiles.yml at {profiles_path}")
        else:
            # Fall back to ~/.dbt/profiles.yml
            dbt_dir = Path.home() / ".dbt"
            profiles_path = dbt_dir / "profiles.yml"
            if not profiles_path.exists():
                raise DatabricksProfileError(f"Could not find profiles.yml at {profiles_path}")
            logger.debug(f"Using user profiles.yml at {profiles_path}")

        # Load profiles.yml
        try:
            with open(profiles_path) as f:
                profiles = yaml.safe_load(f)
        except Exception as e:
            raise DatabricksProfileError(f"Failed to parse profiles.yml: {e}")

        # Get profile name from dbt_project.yml
        project_yml_path = self.project_dir / "dbt_project.yml"
        if not project_yml_path.exists():
            raise DatabricksProfileError(f"Could not find dbt_project.yml at {project_yml_path}")

        try:
            with open(project_yml_path) as f:
                project = yaml.safe_load(f)
        except Exception as e:
            raise DatabricksProfileError(f"Failed to parse dbt_project.yml: {e}")

        profile_name = project.get("profile")
        if not profile_name:
            raise DatabricksProfileError("No 'profile' key found in dbt_project.yml")

        # Get profile
        profile = profiles.get(profile_name)
        if not profile:
            raise DatabricksProfileError(f"Profile '{profile_name}' not found in profiles.yml")

        # Get target
        target_name = profile.get("target", "default")
        target = profile.get("outputs", {}).get(target_name)

        if not target:
            raise DatabricksProfileError(f"Target '{target_name}' not found in profile '{profile_name}'")

        logger.debug(f"Using profile '{profile_name}', target '{target_name}'")
        return target

        # Get profile name from dbt_project.yml
        project_yml_path = self.project_dir / "dbt_project.yml"
        if not project_yml_path.exists():
            raise DatabricksProfileError(f"Could not find dbt_project.yml at {project_yml_path}")

        try:
            with open(project_yml_path) as f:
                project = yaml.safe_load(f)
        except Exception as e:
            raise DatabricksProfileError(f"Failed to parse dbt_project.yml: {e}")

        profile_name = project.get("profile")
        if not profile_name:
            raise DatabricksProfileError("No 'profile' key found in dbt_project.yml")

        # Get profile
        profile = profiles.get(profile_name)
        if not profile:
            raise DatabricksProfileError(f"Profile '{profile_name}' not found in profiles.yml")

        # Get target
        target_name = profile.get("target", "default")
        target = profile.get("outputs", {}).get(target_name)

        if not target:
            raise DatabricksProfileError(f"Target '{target_name}' not found in profile '{profile_name}'")

        logger.debug(f"Using profile '{profile_name}', target '{target_name}'")
        return target
