auto-setup.py

#!/usr/bin/env python3
"""MLflow Demo Automated Setup Script.

This script automates the entire end-to-end setup process for the MLflow demo application,
including creating Databricks resources, configuring environment, loading sample data,
and deploying the application.

Usage:
    python auto-setup.py [options]

Options:
    --dry-run          Show what would be created without actually creating resources
    --resume           Resume from previous failed/interrupted setup
    --reset            Reset all progress and start fresh
    --validate-only    Only run validation checks
    --help             Show this help message
"""

import argparse
import os
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Any, Dict, List

from databricks.sdk import WorkspaceClient
from databricks.sdk.errors import NotFound, PermissionDenied
from dotenv import load_dotenv

# Add automation directory to path
automation_dir = Path(__file__).parent / 'automation'
sys.path.insert(0, str(automation_dir))

# Import after sys.path modification
from environment_detector import EnvironmentDetector  # noqa: E402
from progress_tracker import ProgressTracker  # noqa: E402
from resource_manager import DatabricksResourceManager  # noqa: E402
from validation import SetupValidator  # noqa: E402


class Spinner:
  """Simple spinner to show progress during long operations."""

  def __init__(self, message: str):
    self.message = message
    self.spinning = False
    self.thread = None
    self.spinner_chars = '⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏'

  def start(self):
    """Start the spinner."""
    self.spinning = True
    self.thread = threading.Thread(target=self._spin)
    self.thread.start()

  def stop(self, success_message: str = None):
    """Stop the spinner."""
    self.spinning = False
    if self.thread:
      self.thread.join()
    # Clear the line and show completion
    print(f'\r{" " * (len(self.message) + 10)}', end='')  # Clear line
    if success_message:
      print(f'\r✅ {success_message}')
    else:
      print(f'\r✅ {self.message}')

  def _spin(self):
    """Internal spinner loop."""
    i = 0
    while self.spinning:
      print(
        f'\r{self.spinner_chars[i % len(self.spinner_chars)]} {self.message}', end='', flush=True
      )
      time.sleep(0.1)
      i += 1


class AutoSetup:
  """Main orchestrator for automated MLflow demo setup."""

  def __init__(self, dry_run: bool = False):
    """Initialize the auto setup.

    Args:
        dry_run: If True, only show what would be done without executing
    """
    self.dry_run = dry_run
    self.project_root = Path(__file__).parent

    # Initialize components (defer authentication until actually running setup)
    self.client = None
    self.resource_manager = None
    self.env_detector = None
    self.validator = None
    self.progress = ProgressTracker(self.project_root)

    # Store configuration and created resources
    self.config = {}
    self.created_resources = {}
    self.detected_settings = {}

  def _initialize_databricks_components(self, skip_auth_prompts: bool = False) -> bool:
    """Initialize Databricks components with authentication."""
    if not self.dry_run:
      # Handle authentication first
      if not self._ensure_databricks_auth(skip_prompts=skip_auth_prompts):
        return False

      try:
        self.client = WorkspaceClient()
        self.resource_manager = DatabricksResourceManager(self.client)
        self.env_detector = EnvironmentDetector(self.client)
        self.validator = SetupValidator(self.client)
        return True
      except Exception as e:
        print(f'❌ Failed to initialize Databricks SDK: {e}')
        return False
    else:
      # For dry run, keep placeholder components
      return True

  def _test_create_schema_permission(self, catalog_name: str) -> bool:
    """Test if user can actually create schemas in this catalog."""
    import os

    from databricks.sdk.errors import PermissionDenied

    # Try to create a test schema with a unique name
    test_schema_name = f'test_perms_{int(os.urandom(4).hex(), 16)}'

    try:
      # Try to create the test schema
      self.client.schemas.create(
        name=test_schema_name,
        catalog_name=catalog_name,
        comment='Temporary test schema for permission verification',
      )

      # If successful, clean up immediately
      try:
        self.client.schemas.delete(f'{catalog_name}.{test_schema_name}')
      except Exception:
        pass  # Cleanup failed but permission test succeeded

      return True

    except PermissionDenied:
      return False
    except Exception:
      return False  # Any other error means we can't create schemas

  def _get_available_catalogs_with_permissions(self) -> Dict[str, str]:
    """Get catalogs where user has required permissions (USE CATALOG + CREATE SCHEMA)."""
    available_catalogs = {}
    try:
      spinner = Spinner('Loading available catalogs...')
      spinner.start()
      try:
        catalogs = list(self.client.catalogs.list())
        spinner.stop('Found catalogs')
      except Exception as e:
        spinner.stop()
        raise e

      # Limit to first 50 catalogs to avoid timeout issues
      # Most users won't need to check thousands of catalogs
      catalog_sample = catalogs[:50] if len(catalogs) > 50 else catalogs
      print(
        f'✅ Found {len(catalogs)} catalogs, verifying CREATE SCHEMA permissions on '
        f'first {len(catalog_sample)}'
      )
      print('🔍 This may take a moment as we test actual CREATE SCHEMA permissions...')

      verified_count = 0
      for i, catalog in enumerate(catalog_sample):
        catalog_name = catalog.name

        # Show progress for long operations
        if i % 10 == 0 and i > 0:
          print(
            f'   Verified {verified_count} usable catalogs after checking '
            f'{i}/{len(catalog_sample)}...'
          )

        try:
          # First check if we can list schemas (USE CATALOG permission)
          list(self.client.schemas.list(catalog_name=catalog_name))

          # Then test CREATE SCHEMA permission by actually trying to create a schema
          if self._test_create_schema_permission(catalog_name):
            available_catalogs[catalog_name] = 'VERIFIED: USE CATALOG + CREATE SCHEMA'
            verified_count += 1
            print(f'   ✅ {catalog_name} - CREATE SCHEMA verified')

            # Stop after finding a reasonable number to avoid excessive API calls
            if verified_count >= 10:
              print(
                f'   Found {verified_count} verified catalogs, stopping search to avoid timeouts'
              )
              break

        except Exception:
          # Can't access this catalog
          continue

      print(f'✅ Found {len(available_catalogs)} catalogs with verified CREATE SCHEMA permissions')

    except Exception as e:
      print(f'⚠️  Could not verify catalog permissions: {e}')

    return available_catalogs

  def _get_available_schemas_in_catalog(self, catalog_name: str) -> Dict[str, str]:
    """Get schemas in a catalog where user has required permissions."""
    available_schemas = {}
    try:
      spinner = Spinner(f"Loading schemas in catalog '{catalog_name}'...")
      spinner.start()
      try:
        schemas = list(self.client.schemas.list(catalog_name=catalog_name))
        spinner.stop(f"Found {len(schemas)} schemas in '{catalog_name}'")
      except Exception as e:
        spinner.stop()
        raise e

      if not schemas:
        print(f"   No schemas found in '{catalog_name}'")
        return available_schemas

      print(f'🔍 Checking permissions on {len(schemas)} schemas...')

      for i, schema in enumerate(schemas):
        schema_name = schema.name
        full_schema_name = f'{catalog_name}.{schema_name}'

        # Show progress for long operations
        if len(schemas) > 10 and i % 5 == 0 and i > 0:
          print(f'   Checked {i}/{len(schemas)} schemas...')

        # Check if user has required permissions: MANAGE and CREATE TABLE
        has_manage = False
        has_create_table = False

        try:
          # Get schema info to check basic access
          schema_info = self.client.schemas.get(full_schema_name)

          # Check for MANAGE permission by trying to get effective permissions
          try:
            # Try to get grants/permissions on the schema
            # This is an approximation - checking if we can read schema grants
            grants = list(
              self.client.grants.get_effective(securable_type='schema', full_name=full_schema_name)
            )

            # Look for our permissions in the grants
            current_user = self.client.current_user.me()
            user_email = current_user.user_name if hasattr(current_user, 'user_name') else None

            for grant in grants:
              # Check if this grant applies to current user
              principal = getattr(grant, 'principal', '')
              privileges = getattr(grant, 'privileges', [])

              if user_email and user_email in principal:
                for privilege in privileges:
                  if privilege == 'ALL_PRIVILEGES' or privilege == 'OWNER':
                    has_manage = True
                    has_create_table = True
                    break
                  elif privilege == 'CREATE_TABLE':
                    has_create_table = True
                  elif privilege == 'MANAGE':
                    has_manage = True

            # If we can't determine from grants, try alternative check
            if not (has_manage and has_create_table):
              # Try to list tables as a proxy for having reasonable access
              try:
                list(self.client.tables.list(catalog_name=catalog_name, schema_name=schema_name))
                # If we can list tables, assume we have at least some useful access
                has_create_table = True
              except Exception:
                pass

          except Exception:
            # If we can't check grants, try a simpler approach
            # Try to list tables as a proxy for having reasonable access
            try:
              list(self.client.tables.list(catalog_name=catalog_name, schema_name=schema_name))
              # If we can list tables, assume we have reasonable access
              has_create_table = True
              has_manage = True  # Assume if we can list tables, we have good access
            except Exception:
              pass

          # Determine permission level
          if has_manage and has_create_table:
            available_schemas[schema_name] = 'MANAGE + CREATE TABLE'
          elif has_create_table:
            available_schemas[schema_name] = 'CREATE TABLE only'
          elif schema_info:
            # We can see the schema but don't have required permissions
            continue  # Don't include schemas without required permissions

        except Exception:
          # Can't access schema at all
          continue

      print(f'✅ Found {len(available_schemas)} accessible schemas with permissions')

    except Exception as e:
      print(f'⚠️  Could not list schemas in {catalog_name}: {e}')

    return available_schemas

  def _prompt_for_catalog_selection(self, suggested_catalog: str = None) -> str:
    """Interactive catalog selection with permission checking."""
    print('\n📁 Unity Catalog Selection')

    # Get available catalogs
    available_catalogs = self._get_available_catalogs_with_permissions()

    if not available_catalogs:
      print('❌ No catalogs found with verified CREATE SCHEMA permissions in the sample checked.')
      print('')
      print('💡 This means you need CREATE SCHEMA permission on a catalog. Options:')
      print('   1. Ask your workspace admin to grant CREATE SCHEMA permission')
      print('   2. Use a catalog where you already have permissions')
      print('   3. Create your own catalog (if you have CREATE CATALOG permission)')
      print('')

      # Allow manual entry as fallback
      while True:
        try:
          manual_catalog = input('Enter catalog name manually (or press Enter to skip): ').strip()
          if not manual_catalog:
            return None

          # Test the manually entered catalog with CREATE SCHEMA verification
          print(f"🔍 Verifying CREATE SCHEMA permission on '{manual_catalog}'...")
          try:
            # First check if we can list schemas (USE CATALOG)
            list(self.client.schemas.list(catalog_name=manual_catalog))
            print(f"✅ Can list schemas in '{manual_catalog}' - USE CATALOG confirmed")

            # Then test CREATE SCHEMA permission
            if self._test_create_schema_permission(manual_catalog):
              print(f"✅ CREATE SCHEMA permission verified on '{manual_catalog}'")
              return manual_catalog
            else:
              print(f"❌ No CREATE SCHEMA permission on '{manual_catalog}'")
              print('   You need CREATE SCHEMA permission to use this catalog')
              continue

          except Exception as e:
            print(f"❌ Cannot access catalog '{manual_catalog}': {e}")
            continue

        except KeyboardInterrupt:
          return None

    print(
      f'Available catalogs (showing {len(available_catalogs)} with VERIFIED '
      f'CREATE SCHEMA permissions):'
    )
    catalog_list = list(available_catalogs.keys())

    # Show suggested catalog first if it exists
    if suggested_catalog and suggested_catalog in available_catalogs:
      print(f'   0. {suggested_catalog} (suggested) - {available_catalogs[suggested_catalog]}')
      start_idx = 1
    else:
      start_idx = 0

    # Show other catalogs
    for i, (catalog_name, access_level) in enumerate(available_catalogs.items()):
      if catalog_name != suggested_catalog:
        print(f'   {start_idx + i}. {catalog_name} - {access_level}')

    # Add option to manually enter catalog name
    manual_entry_idx = len(catalog_list) + (1 if suggested_catalog in available_catalogs else 0)
    print(f'   {manual_entry_idx}. Enter catalog name manually')

    max_choice = manual_entry_idx

    while True:
      try:
        choice = input(f'\nSelect catalog (0-{max_choice}) or type catalog name: ').strip()

        # Check if it's a number
        try:
          choice_num = int(choice)
          if choice_num == 0 and suggested_catalog and suggested_catalog in available_catalogs:
            return suggested_catalog
          elif 1 <= choice_num <= len(catalog_list):
            # Adjust index based on whether suggested catalog is shown
            if suggested_catalog and suggested_catalog in available_catalogs:
              selected_catalogs = [cat for cat in catalog_list if cat != suggested_catalog]
              return selected_catalogs[choice_num - 1]
            else:
              return catalog_list[choice_num - 1]
          elif choice_num == manual_entry_idx:
            # Manual entry option
            while True:
              manual_catalog = input('Enter catalog name: ').strip()
              if not manual_catalog:
                print('❌ Catalog name cannot be empty')
                continue
              try:
                list(self.client.schemas.list(catalog_name=manual_catalog))
                print(f"✅ Catalog '{manual_catalog}' is accessible")
                return manual_catalog
              except Exception as e:
                print(f"❌ Cannot access catalog '{manual_catalog}': {e}")
                continue
          else:
            print(f'❌ Please enter a number between 0 and {max_choice}')
            continue
        except ValueError:
          # User typed a catalog name directly
          if choice in available_catalogs:
            return choice
          else:
            # Test the typed catalog name
            try:
              list(self.client.schemas.list(catalog_name=choice))
              print(f"✅ Catalog '{choice}' is accessible")
              return choice
            except Exception as e:
              print(f"❌ Cannot access catalog '{choice}': {e}")
              continue

      except KeyboardInterrupt:
        return None

  def _prompt_for_schema_selection(self, catalog_name: str, suggested_schema: str = None) -> str:
    """Interactive schema selection with permission checking."""
    print(f"\n📂 Schema Selection in '{catalog_name}'")

    # Get available schemas
    available_schemas = self._get_available_schemas_in_catalog(catalog_name)

    if not available_schemas:
      print(f"No accessible schemas found in '{catalog_name}'. Will create new schema.")
      new_schema = input('Enter new schema name [default]: ').strip()
      return new_schema or 'default'

    print('Available schemas:')
    schema_list = list(available_schemas.keys())

    # Show suggested schema first if it exists
    if suggested_schema and suggested_schema in available_schemas:
      print(f'   0. {suggested_schema} (suggested) - {available_schemas[suggested_schema]}')
      start_idx = 1
    else:
      start_idx = 0

    # Show other schemas
    for i, (schema_name, access_level) in enumerate(available_schemas.items()):
      if schema_name != suggested_schema:
        print(f'   {start_idx + i}. {schema_name} - {access_level}')

    create_option_num = len(schema_list) + (1 if suggested_schema in available_schemas else 0)
    print(f'   {create_option_num}. Create new schema')

    while True:
      try:
        max_choice = len(schema_list) + (1 if suggested_schema in available_schemas else 0)
        choice = input(f'\nSelect schema (0-{max_choice}) or type schema name: ').strip()

        # Check if it's a number
        try:
          choice_num = int(choice)
          if choice_num == 0 and suggested_schema and suggested_schema in available_schemas:
            return suggested_schema
          elif 1 <= choice_num <= len(schema_list):
            # Adjust index based on whether suggested schema is shown
            if suggested_schema and suggested_schema in available_schemas:
              selected_schemas = [sch for sch in schema_list if sch != suggested_schema]
              return selected_schemas[choice_num - 1]
            else:
              return schema_list[choice_num - 1]
          elif choice_num == len(schema_list) + (1 if suggested_schema in available_schemas else 0):
            # Create new schema
            new_schema = input('Enter new schema name: ').strip()
            if new_schema:
              print(f'💡 Will create new schema: {new_schema}')
              return new_schema
          else:
            max_choice = len(schema_list) + (1 if suggested_schema in available_schemas else 0)
            print(f'❌ Please enter a number between 0 and {max_choice}')
            continue
        except ValueError:
          # User typed a schema name directly
          if choice in available_schemas:
            return choice
          else:
            print(f'💡 Will create new schema: {choice}')
            return choice

      except KeyboardInterrupt:
        return None

  def _validate_app_name(self, app_name: str) -> bool:
    """Validate app name format for Databricks Apps."""
    import re

    if not app_name:
      return False
    # App name must contain only lowercase letters, numbers, and dashes
    pattern = r'^[a-z0-9-]+$'
    return bool(re.match(pattern, app_name))

  def _get_available_chat_models(self) -> List[str]:
    """Get list of available chat completion models from Databricks."""
    try:
      # Use the model discovery logic to find chat models

      spinner = Spinner('Discovering available chat models...')
      spinner.start()
      try:
        endpoints = self.client.serving_endpoints.list()
        spinner.stop('Found serving endpoints')
      except Exception as e:
        spinner.stop()
        raise e

      # Common chat model patterns
      chat_model_patterns = [
        'gpt-',
        'claude-',
        'gemini-',
        'llama',
        'mistral',
        'databricks-',
        'chat',
        'instruct',
        'turbo',
      ]

      potential_chat_models = []

      for endpoint in endpoints:
        model_name = endpoint.name.lower()

        # Check if it matches chat patterns
        is_likely_chat = any(pattern in model_name for pattern in chat_model_patterns)

        # Exclude obvious non-chat models
        is_not_chat = any(
          exclude in model_name
          for exclude in ['embedding', 'vision', 'audio', 'whisper', 'imageai']
        )

        if is_likely_chat and not is_not_chat:
          # Check if endpoint has chat task capability
          try:
            endpoint_details = self.client.serving_endpoints.get(name=endpoint.name)
            if hasattr(endpoint_details, 'config') and endpoint_details.config:
              # Check served entities for task type
              if (
                hasattr(endpoint_details.config, 'served_entities')
                and endpoint_details.config.served_entities
              ):
                for entity in endpoint_details.config.served_entities:
                  if (
                    hasattr(entity, 'external_model')
                    and entity.external_model
                    and hasattr(entity.external_model, 'task')
                    and entity.external_model.task == 'llm/v1/chat'
                  ):
                    potential_chat_models.append(endpoint.name)
                    break
                  elif (
                    hasattr(entity, 'foundation_model')
                    and entity.foundation_model
                    and hasattr(entity.foundation_model, 'name')
                  ):
                    # Foundation models typically support chat
                    potential_chat_models.append(endpoint.name)
                    break
          except Exception:
            # If we can't get details, include it based on name pattern
            potential_chat_models.append(endpoint.name)

      # Remove duplicates and sort
      chat_models = sorted(list(set(potential_chat_models)))

      # Prioritize certain models at the top
      priority_models = ['databricks-claude-3-7-sonnet', 'databricks-claude-sonnet-4', 'gpt-4o']
      prioritized_models = []

      for priority in priority_models:
        if priority in chat_models:
          prioritized_models.append(priority)
          chat_models.remove(priority)

      return prioritized_models + chat_models

    except Exception as e:
      print(f'⚠️  Could not discover chat models: {e}')
      # Return default options
      return [
        'databricks-claude-3-7-sonnet',
        'databricks-claude-sonnet-4',
        'databricks-meta-llama-3-3-70b-instruct',
        'gpt-4o',
      ]

  def _prompt_for_deployment_mode(self) -> str:
    """Interactive deployment mode selection."""
    print('\n🚀 Choose Your Interface')
    print('')
    print('Both options include the complete MLflow evaluation setup:')
    print(
      '• MLflow Experiment with sample traces, evaluation runs, prompts, and production monitoring'
    )
    print('• Sales email generation code with sample data')
    print('• Interactive Notebooks that walk you through using MLflow to improve GenAI quality')
    print('')
    print('Choose how you want to interact with the demo:')
    print('')
    print('   1. 📱 Databricks App (Recommended if Databricks Apps are enabled in your workspace)')
    print('      • Adds a user-friendly web UI for exploring the workflows')
    print('      • Interactive demo UI to try the sample app')
    print('      • Include all the notebooks for deeper exploration')
    print('')
    print('   2. 📓 Notebooks Only')
    print("      • Best option if you can't deploy Databricks Apps in your workspace")
    print('      • Use the notebooks to understand how MLflow helps you improve GenAI quality')
    print('')

    while True:
      try:
        choice = input('Select experience (1 or 2) [default: 1]: ').strip()

        # Default to full app deployment
        if not choice:
          choice = '1'

        if choice == '1':
          print('✅ Selected: Full App Deployment')
          return 'full_deployment'
        elif choice == '2':
          print('✅ Selected: Notebook-Only Experience')
          return 'notebook_only'
        else:
          print('❌ Please enter 1 or 2')
          continue

      except KeyboardInterrupt:
        return 'notebook_only'

  def _prompt_for_llm_model(self, suggested_model: str = None) -> str:
    """Interactive LLM model selection with available chat models."""
    print('\n🤖 LLM Model Selection')

    # Get available chat models
    available_models = self._get_available_chat_models()

    if not available_models:
      print('❌ No chat models found. Using default.')
      return suggested_model or 'databricks-claude-3-7-sonnet'

    print('Available chat completion models:')

    # Show suggested model first if it exists
    if suggested_model and suggested_model in available_models:
      print(f'   0. {suggested_model} (suggested)')
      start_idx = 1
    else:
      start_idx = 0

    # Show other models
    for i, model_name in enumerate(available_models):
      if model_name != suggested_model:
        print(f'   {start_idx + i}. {model_name}')

    max_choice = len(available_models) - 1 + (1 if suggested_model in available_models else 0)

    while True:
      try:
        choice = input(f'\nSelect model (0-{max_choice}) or press ENTER for default: ').strip()

        # Use default if empty
        if not choice:
          return suggested_model or available_models[0]

        # Check if it's a number
        try:
          choice_num = int(choice)
          if choice_num == 0 and suggested_model and suggested_model in available_models:
            return suggested_model
          elif 1 <= choice_num <= len(available_models):
            # Adjust index based on whether suggested model is shown
            if suggested_model and suggested_model in available_models:
              selected_models = [model for model in available_models if model != suggested_model]
              return selected_models[choice_num - 1]
            else:
              return available_models[choice_num - 1]
          else:
            print(f'❌ Please enter a number between 0 and {max_choice}')
            continue
        except ValueError:
          # User typed a model name directly
          if choice in available_models:
            return choice
          else:
            print(f"❌ Model '{choice}' not found in available models")
            continue

      except KeyboardInterrupt:
        return suggested_model or available_models[0]

  def _generate_default_app_name(self) -> str:
    """Generate a default app name with 4 random characters."""
    import os

    # Generate 4 random hex characters
    random_chars = os.urandom(2).hex()  # 2 bytes = 4 hex chars
    return f'mlflow-demo-app-{random_chars}'

  def _prompt_for_app_name(self, suggested_app_name: str = None) -> str:
    """Interactive app name selection with default suggestion."""
    print('\n📱 Databricks App Name Selection')

    # Generate a default name if no suggestion provided
    if not suggested_app_name:
      suggested_app_name = self._generate_default_app_name()

    while True:
      app_name = input(f'App name [{suggested_app_name}]: ').strip()
      if not app_name:
        app_name = suggested_app_name

      if not self._validate_app_name(app_name):
        print('❌ App name must contain only lowercase letters, numbers, and dashes')
        continue

      print(f'💡 Will create app: {app_name}')
      return app_name

  def _restore_config_from_progress(self):
    """Restore configuration from saved progress."""
    try:
      # Look for saved config in completed steps
      completed_step_ids = self.progress.get_completed_steps()
      for step_id in completed_step_ids:
        step_data = self.progress.steps.get(step_id)
        if step_data and hasattr(step_data, 'result_data') and step_data.result_data:
          if 'config' in step_data.result_data:
            saved_config = step_data.result_data['config']
            self.config.update(saved_config)
            print(f'🔄 Restored configuration from progress: {list(saved_config.keys())}')

          # Also restore experiment ID if available
          if 'experiment_id' in step_data.result_data:
            self.config['MLFLOW_EXPERIMENT_ID'] = step_data.result_data['experiment_id']
            print(f'🔄 Restored experiment ID: {step_data.result_data["experiment_id"]}')

          # Restore other important config values
          if 'app_name' in step_data.result_data:
            self.config['DATABRICKS_APP_NAME'] = step_data.result_data['app_name']
    except Exception as e:
      print(f'⚠️  Could not restore config from progress: {e}')

  def _load_config_from_env_file(self):
    """Load configuration from .env.local file."""
    try:
      env_file = self.project_root / '.env.local'
      if env_file.exists():
        # Load environment variables from .env.local
        load_dotenv(env_file)

        # Update self.config with values from environment
        env_mappings = {
          'DATABRICKS_HOST': 'DATABRICKS_HOST',
          'UC_CATALOG': 'UC_CATALOG',
          'UC_SCHEMA': 'UC_SCHEMA',
          'DATABRICKS_APP_NAME': 'DATABRICKS_APP_NAME',
          'LLM_MODEL': 'LLM_MODEL',
          'DEPLOYMENT_MODE': 'DEPLOYMENT_MODE',
          'MLFLOW_EXPERIMENT_ID': 'MLFLOW_EXPERIMENT_ID',
          'LHA_SOURCE_CODE_PATH': 'LHA_SOURCE_CODE_PATH',
        }

        loaded_keys = []
        for config_key, env_key in env_mappings.items():
          value = os.getenv(env_key)
          if value:
            self.config[config_key] = value
            loaded_keys.append(config_key)

        if loaded_keys:
          print(f'🔄 Loaded configuration from .env.local: {loaded_keys}')
      else:
        print('⚠️  .env.local file not found')
    except Exception as e:
      print(f'⚠️  Could not load config from .env.local: {e}')

  def run_setup(self, resume: bool = False) -> bool:
    """Run the complete setup process.

    Args:
        resume: If True, resume from previous progress

    Returns:
        True if setup completed successfully
    """
    print('🚀 MLflow Demo Automated Setup')
    print('=' * 50)

    if not resume:
      print(f'🔧 Setup mode: {"DRY RUN" if self.dry_run else "LIVE"}')
    else:
      print('🔄 Resuming previous setup...')
      self.progress.show_detailed_progress()

    # Restore configuration from previous progress if resuming
    if resume:
      self._restore_config_from_progress()
      # Also load configuration from .env.local to ensure all values are present
      self._load_config_from_env_file()

    # Initialize Databricks components (skip auth prompts on resume if already working)
    if not self._initialize_databricks_components(skip_auth_prompts=resume):
      return False

    try:
      success = True

      # Execute setup steps in order
      while True:
        next_step = self.progress.get_next_step()
        if not next_step:
          break

        if not self.progress.start_step(next_step):
          continue

        try:
          # Check if we should skip app-related steps for notebook-only mode
          deployment_mode = self.config.get('DEPLOYMENT_MODE', 'full_deployment')
          app_steps = [
            'create_app',
            'setup_permissions',
            'validate_deployment',
            'run_integration_tests',
          ]

          if deployment_mode == 'notebook_only' and next_step in app_steps:
            print(f'📓 Skipping {next_step} for notebook-only mode')
            success = True  # Skip successfully
          elif next_step == 'validate_prerequisites':
            success = self._validate_prerequisites()
          elif next_step == 'detect_environment':
            success = self._detect_environment()
          elif next_step == 'collect_user_input':
            success = self._collect_user_input()
          elif next_step == 'validate_config':
            success = self._validate_config()
          elif next_step == 'show_installation_preview':
            success = self._show_installation_preview()
          elif next_step == 'create_catalog_schema':
            success = self._create_catalog_schema()
          elif next_step == 'create_experiment':
            success = self._create_experiment()
          elif next_step == 'create_app':
            success = self._create_app()
          elif next_step == 'setup_permissions':
            success = self._setup_permissions()
          elif next_step == 'generate_env_file':
            success = self._generate_env_file()
          elif next_step == 'install_dependencies':
            success = self._install_dependencies()
          elif next_step == 'load_sample_data':
            success = self._load_sample_data()
          elif next_step == 'validate_local_setup':
            success = self._validate_local_setup()
          elif next_step == 'deploy_app':
            success = self._deploy_app()
          elif next_step == 'validate_deployment':
            success = self._validate_deployment()
          elif next_step == 'run_integration_tests':
            success = self._run_integration_tests()
          else:
            success = False
            raise ValueError(f'Unknown step: {next_step}')

          if success:
            self.progress.complete_step(next_step, self._get_step_result(next_step))
          else:
            self.progress.fail_step(next_step, 'Step execution failed')
            break

        except Exception as e:
          error_msg = f'Error in step {next_step}: {str(e)}'
          print(f'❌ {error_msg}')
          self.progress.fail_step(next_step, error_msg)
          success = False
          break

      # Show final results
      self._show_final_results(success)
      return success

    except KeyboardInterrupt:
      print('\n⚠️  Setup interrupted by user')
      self._show_final_results(False)
      return False
    except Exception as e:
      print(f'\n❌ Unexpected error during setup: {e}')
      self._show_final_results(False)
      return False

  def _validate_prerequisites(self) -> bool:
    """Validate prerequisites before setup."""
    print('🔍 Validating prerequisites...')

    if self.dry_run:
      print('   [DRY RUN] Would validate prerequisites')
      return True

    valid, issues = self.validator.validate_prerequisites()

    if not valid:
      print('❌ Prerequisites validation failed:')
      for issue in issues:
        print(f'   • {issue}')
      return False

    return True

  def _detect_environment(self) -> bool:
    """Detect environment settings."""
    print('🔍 Detecting environment settings...')

    if self.dry_run:
      print('   [DRY RUN] Would detect environment settings')
      # Set dummy values for dry run
      self.detected_settings = {
        'workspace_url': 'https://demo.cloud.databricks.com',
        'suggested_catalog': 'workspace',
        'suggested_schema': 'default',
      }
      return True

    # Detect workspace URL
    workspace_url = self.env_detector.detect_workspace_url()

    # Suggest catalog/schema
    catalog, schema = self.env_detector.suggest_catalog_schema()

    # Store detected settings
    self.detected_settings = {
      'workspace_url': workspace_url,
      'suggested_catalog': catalog,
      'suggested_schema': schema,
    }

    return True

  def _collect_user_input(self) -> bool:
    """Collect required user input."""
    print('📝 Collecting configuration...')

    if self.dry_run:
      print('   [DRY RUN] Would collect user input')
      # Use dummy values for dry run
      self.config = {
        'DATABRICKS_HOST': 'https://demo.cloud.databricks.com',
        'UC_CATALOG': 'workspace',
        'UC_SCHEMA': 'default',
        'DATABRICKS_APP_NAME': 'mlflow_demo_app',
        'MLFLOW_EXPERIMENT_ID': '123456789',
        'DEPLOYMENT_MODE': 'notebook_only',  # Default to notebook-only for dry run
      }
      return True

    print('\n📝 Configuration Setup')
    print('Please provide the following information:')
    print('(Press Enter to use suggested values where available)\n')

    # Get workspace URL automatically from the authenticated profile
    workspace_url = self.detected_settings.get('workspace_url')
    if not workspace_url:
      try:
        workspace_url = self.client.config.host
      except Exception:
        workspace_url = 'https://unknown-workspace.cloud.databricks.com'

    print(f'✅ Using workspace URL from profile: {workspace_url}')

    # Choose deployment mode first
    deployment_mode = self._prompt_for_deployment_mode()

    # Interactive catalog selection
    suggested_catalog = self.detected_settings.get('suggested_catalog')
    catalog = self._prompt_for_catalog_selection(suggested_catalog)
    if not catalog:
      print('❌ Catalog selection is required')
      return False

    # Interactive schema selection
    suggested_schema = self.detected_settings.get('suggested_schema')
    schema = self._prompt_for_schema_selection(catalog, suggested_schema)
    if not schema:
      print('❌ Schema selection is required')
      return False

    # App name selection - only required for full deployment
    app_name = None
    if deployment_mode == 'full_deployment':
      app_name = self._prompt_for_app_name()
      if not app_name:
        print('❌ App name is required for full deployment')
        return False
    else:
      # For notebook-only mode, use a default app name for workspace path generation
      app_name = self._generate_default_app_name()
      print(f'📓 Using default name for workspace sync: {app_name}')

    # LLM model selection
    llm_model = self._prompt_for_llm_model('databricks-claude-3-7-sonnet')

    # Store configuration
    self.config = {
      'DATABRICKS_HOST': workspace_url,
      'UC_CATALOG': catalog,
      'UC_SCHEMA': schema,
      'DATABRICKS_APP_NAME': app_name,
      'LLM_MODEL': llm_model,
      'DEPLOYMENT_MODE': deployment_mode,
    }

    return True

  def _validate_config(self) -> bool:
    """Validate user configuration."""
    print('✅ Validating configuration...')

    if self.dry_run:
      print('   [DRY RUN] Would validate configuration')
      return True

    # Only validate the configuration we have at this point
    issues = []

    # Check required variables that should exist at this stage
    required_at_this_stage = [
      'DATABRICKS_HOST',
      'UC_CATALOG',
      'UC_SCHEMA',
      'DATABRICKS_APP_NAME',
      'LLM_MODEL',
    ]

    for var in required_at_this_stage:
      if not self.config.get(var):
        issues.append(f'Missing required configuration: {var}')

    # Validate workspace URL format
    if 'DATABRICKS_HOST' in self.config:
      host = self.config['DATABRICKS_HOST']
      if not host.startswith('https://'):
        issues.append('DATABRICKS_HOST must start with https://')
      if '.cloud.databricks.com' not in host and '.azuredatabricks.net' not in host:
        issues.append("DATABRICKS_HOST doesn't appear to be a valid Databricks URL")

    # Validate catalog.schema format
    if 'UC_CATALOG' in self.config and 'UC_SCHEMA' in self.config:
      catalog = self.config['UC_CATALOG']
      schema = self.config['UC_SCHEMA']
      if '.' in catalog or '.' in schema:
        issues.append('UC_CATALOG and UC_SCHEMA should not contain dots')

    if issues:
      print('❌ Configuration validation failed:')
      for issue in issues:
        print(f'   • {issue}')
      return False

    return True

  def _show_installation_preview(self) -> bool:
    """Show preview of what will be created and get user confirmation."""
    print('\n' + '=' * 60)
    print('📋 INSTALLATION PREVIEW')
    print('=' * 60)
    print('\nThe following resources will be created/configured:')

    deployment_mode = self.config.get('DEPLOYMENT_MODE', 'full_deployment')

    # Workspace and authentication
    print(f'\n🏢 Workspace: {self.config.get("DATABRICKS_HOST", "Unknown")}')

    # Unity Catalog resources
    catalog = self.config.get('UC_CATALOG', 'Unknown')
    schema = self.config.get('UC_SCHEMA', 'Unknown')
    print(f'📁 Unity Catalog: {catalog}.{schema}')
    print(f"   • Will create catalog '{catalog}' if it doesn't exist")
    print(f"   • Will create schema '{schema}' if it doesn't exist")

    # MLflow experiment
    app_name = self.config.get('DATABRICKS_APP_NAME', 'mlflow_demo_app')
    experiment_path = f'/Shared/{app_name}'
    print(f'🧪 MLflow Experiment: {experiment_path}')

    # Deployment mode specific resources
    if deployment_mode == 'full_deployment':
      print(f'📱 Databricks App: {app_name}')
      print(f"   • Will create/update app '{app_name}'")
      print('   • Will deploy web application to the app')
      print('   • Will configure app permissions for service principal')
    else:
      print(f'📓 Notebook-Only Mode: {app_name}')
      # Try to get the actual current user
      try:
        current_user = self.env_detector.get_current_user() if self.env_detector else None
        if current_user:
          user_path = f'/Workspace/Users/{current_user}/{app_name}'
        else:
          user_path = f'/Workspace/Users/[user]/{app_name}'
      except Exception:
        user_path = f'/Workspace/Users/[user]/{app_name}'
      print(f'   • Will sync notebooks to workspace at {user_path}')

    # LLM model
    llm_model = self.config.get('LLM_MODEL', 'Unknown')
    print(f'🤖 LLM Model: {llm_model}')

    # Sample data
    print('\n📊 Sample Data Setup:')
    print('   • Load prompt templates into MLflow')
    print('   • Generate sample traces, evaluations, and labeling sessions')
    print('   • Configure production monitoring')

    # Permissions (only for full deployment mode)
    if deployment_mode == 'full_deployment':
      print('\n🔐 Permissions (for app service principal):')
      print(f"   • USE CATALOG on '{catalog}'")
      print(f"   • ALL_PRIVILEGES + MANAGE on '{catalog}.{schema}'")
      print('   • CAN_MANAGE on MLflow experiment')
      print(f"   • CAN_QUERY on model serving endpoint '{llm_model}'")

    if self.dry_run:
      print('\n🏃 DRY RUN MODE: No actual resources will be created')
      return True

    print('\n' + '=' * 60)

    # Get user confirmation
    while True:
      try:
        choice = input('\n❓ Proceed with installation? (Y/n/details): ').strip().lower()

        if choice in ['', 'y', 'yes']:
          print('✅ Starting installation...')
          return True
        elif choice in ['n', 'no']:
          print('❌ Installation cancelled by user')
          return False
        elif choice in ['d', 'details']:
          self._show_detailed_preview()
          continue
        else:
          print('❌ Please enter Y (yes), N (no), or D (details)')
          continue

      except KeyboardInterrupt:
        print('\n❌ Installation cancelled by user')
        return False

  def _show_detailed_preview(self):
    """Show detailed information about what will be created."""
    print('\n' + '=' * 60)
    print('📋 DETAILED PREVIEW')
    print('=' * 60)

    print('\n🔧 Configuration Details:')
    for key, value in sorted(self.config.items()):
      if key in [
        'DATABRICKS_HOST',
        'UC_CATALOG',
        'UC_SCHEMA',
        'DATABRICKS_APP_NAME',
        'LLM_MODEL',
        'DEPLOYMENT_MODE',
        'CUSTOM_EXPERIMENT_PATH',
      ]:
        print(f'   {key}: {value}')

    print('\n📂 Directory Structure (will be created):')
    app_name = self.config.get('DATABRICKS_APP_NAME', 'mlflow_demo_app')
    print(f'   /Workspace/Users/[your-username]/{app_name}/')
    print('   ├── mlflow_demo/')
    print('   │   ├── notebooks/')
    print('   │   │   ├── 0_demo_overview.ipynb')
    print('   │   │   ├── 1_observe_with_traces.ipynb')
    print('   │   │   └── [other demo notebooks]')
    print('   │   └── [application source code]')

    deployment_mode = self.config.get('DEPLOYMENT_MODE', 'full_deployment')
    if deployment_mode == 'full_deployment':
      print('\n🚀 Deployment Process:')
      print(f'   1. Create Databricks App: {app_name}')
      print('   2. Upload source code to workspace')
      print('   3. Deploy application using ./deploy.sh')
      print('   4. Configure service principal permissions')
      print('   5. Start and validate deployment')

    print('\n💡 After installation, you will receive:')
    if deployment_mode == 'full_deployment':
      print('   • Direct link to deployed application')
      print('   • Access to interactive demo interface')
    else:
      print('   • Link to demo overview notebook in your workspace')
      print('   • Path to all notebooks for learning MLflow evaluation')
    print('   • MLflow experiment with sample data for exploration')

  def _create_catalog_schema(self) -> bool:
    """Create catalog and schema if needed."""
    catalog = self.config['UC_CATALOG']
    schema = self.config['UC_SCHEMA']

    print(f'📁 Setting up Unity Catalog: {catalog}.{schema}')

    if self.dry_run:
      print(f"   [DRY RUN] Would create catalog '{catalog}' and schema '{schema}'")
      return True

    try:
      # Check if catalog exists first
      catalog_exists = False
      try:
        self.client.catalogs.get(catalog)
        catalog_exists = True
        print(f"✅ Catalog '{catalog}' already exists")
      except NotFound:
        print(f"📁 Catalog '{catalog}' does not exist, attempting to create...")

      # Create catalog if needed
      if not catalog_exists:
        try:
          catalog_info = self.resource_manager.create_catalog_if_not_exists(catalog)
          self.created_resources['catalog'] = catalog_info.name
          print(f"✅ Created catalog '{catalog}'")
        except PermissionDenied:
          print(f"❌ Permission denied: Cannot create catalog '{catalog}'")
          print('   Please ask your workspace admin to create the catalog or grant you permissions')
          return False
        except Exception as e:
          print(f"❌ Failed to create catalog '{catalog}': {e}")
          return False

      # Check if schema exists
      schema_exists = False
      try:
        self.client.schemas.get(f'{catalog}.{schema}')
        schema_exists = True
        print(f"✅ Schema '{catalog}.{schema}' already exists")
      except NotFound:
        print(f"📂 Schema '{schema}' does not exist in '{catalog}', attempting to create...")

      # Create schema if needed
      if not schema_exists:
        try:
          self.resource_manager.create_schema_if_not_exists(catalog, schema)
          self.created_resources['schema'] = f'{catalog}.{schema}'
          print(f"✅ Created schema '{catalog}.{schema}'")
        except Exception as e:
          print(f"❌ Failed to create schema '{catalog}.{schema}': {e}")
          print(f"   You may need permissions to create schemas in catalog '{catalog}'")
          return False
      else:
        self.created_resources['schema'] = f'{catalog}.{schema}'

      return True
    except Exception as e:
      print(f'❌ Failed to set up catalog/schema: {e}')
      return False

  def _create_experiment(self) -> bool:
    """Create MLflow experiment automatically."""
    # Get app name from config for experiment naming
    app_name = self.config.get('DATABRICKS_APP_NAME', 'mlflow_demo_app')

    # Use default experiment path: /Shared/{app_name}
    experiment_name = f'/Shared/{app_name}'
    print(f'🧪 Creating MLflow experiment: {experiment_name}')

    if self.dry_run:
      print(f"   [DRY RUN] Would create experiment '{experiment_name}'")
      self.config['MLFLOW_EXPERIMENT_ID'] = '123456789'
      return True

    try:
      experiment_id = self.resource_manager.create_mlflow_experiment(name=experiment_name)

      self.config['MLFLOW_EXPERIMENT_ID'] = experiment_id
      self.created_resources['experiment_id'] = experiment_id

      return True
    except Exception as e:
      print(f'❌ Failed to create experiment: {e}')
      return False

  def _create_app(self) -> bool:
    """Create Databricks App."""
    # Use app name from config (should be set by user input collection)
    app_name = self.config.get('DATABRICKS_APP_NAME')
    if not app_name:
      app_name = self._generate_default_app_name()
      self.config['DATABRICKS_APP_NAME'] = app_name

    print(f'📱 Creating Databricks App: {app_name}')

    if self.dry_run:
      print(f"   [DRY RUN] Would create app '{app_name}'")
      return True

    try:
      # Generate workspace path for the app
      current_user = self.env_detector.get_current_user()
      if current_user:
        workspace_path = f'/Workspace/Users/{current_user}/{app_name}'
        self.config['LHA_SOURCE_CODE_PATH'] = workspace_path
      else:
        # Fallback to shared workspace
        workspace_path = f'/Workspace/Shared/{app_name}'
        self.config['LHA_SOURCE_CODE_PATH'] = workspace_path
        print('⚠️  Could not determine current user - using shared workspace path')

      print(f'📁 App source code path: {workspace_path}')

      self.resource_manager.create_databricks_app(
        name=app_name,
        description='MLflow demo application - automated setup',
        source_code_path=workspace_path,
      )

      self.created_resources['app_name'] = app_name

      # Automatically start the app after creation
      print(f"\n📱 App '{app_name}' has been created successfully!")
      print(f"🚀 Starting app '{app_name}'...")
      if not self.resource_manager.start_app(app_name, timeout_minutes=10):
        print(f"⚠️  Failed to start app '{app_name}' - it may need to be started manually")
        print('💡 You can start it later from the Databricks UI or after deployment')
        # Don't fail the setup if app start fails - it can be started later

      return True
    except Exception as e:
      print(f'❌ Failed to create app: {e}')
      return False

  def _setup_permissions(self) -> bool:
    """Setup permissions for app service principal."""
    print('🔐 Setting up permissions...')

    if self.dry_run:
      print('   [DRY RUN] Would setup permissions')
      return True

    try:
      app_name = self.config['DATABRICKS_APP_NAME']

      # Get app service principal (this should work after deployment)
      service_principal = self.resource_manager.get_app_service_principal(app_name)

      if service_principal:
        print(f'✅ Found app service principal: {service_principal}')

        # Grant catalog permissions first (USE CATALOG)
        catalog_name = self.config['UC_CATALOG']
        print(f'🔐 Granting catalog permissions on {catalog_name}...')
        self.resource_manager.grant_catalog_permissions(
          catalog_name, service_principal, permissions=['USE CATALOG']
        )

        # Grant schema permissions (ALL PERMISSIONS + MANAGE)
        schema_name = f'{self.config["UC_CATALOG"]}.{self.config["UC_SCHEMA"]}'
        print(f'🔐 Granting schema permissions on {schema_name}...')
        self.resource_manager.grant_schema_permissions(
          schema_name, service_principal, permissions=['ALL_PRIVILEGES', 'MANAGE']
        )

        # Grant experiment permissions (CAN MANAGE)
        experiment_id = self.config['MLFLOW_EXPERIMENT_ID']
        print(f'🔐 Granting experiment permissions on {experiment_id}...')
        self.resource_manager.grant_experiment_permissions(
          experiment_id, service_principal, permissions=['CAN_MANAGE']
        )

        # Grant model serving endpoint access
        llm_model = self.config.get('LLM_MODEL', 'databricks-claude-3-7-sonnet')
        print(f'🔐 Granting model serving access to {llm_model}...')
        self.resource_manager.grant_model_serving_permissions(app_name, llm_model)

        print('✅ Permissions set successfully')
      else:
        print(
          '⚠️  App service principal not available yet - permissions may need to be set manually'
        )
        print("   This is normal if the app hasn't been deployed yet")

      return True
    except Exception as e:
      print(f'⚠️  Permission setup had issues: {e}')
      print('   You may need to configure permissions manually via the UI')
      return True  # Don't fail setup for permission issues

  def _generate_env_file(self) -> bool:
    """Generate .env.local file."""
    print('📄 Generating environment file...')

    if self.dry_run:
      print('   [DRY RUN] Would generate .env.local file')
      return True

    try:
      # Generate workspace path for both deployment modes (moved from _create_app)
      app_name = self.config.get('DATABRICKS_APP_NAME')
      if app_name and 'LHA_SOURCE_CODE_PATH' not in self.config:
        current_user = self.env_detector.get_current_user()
        if current_user:
          workspace_path = f'/Workspace/Users/{current_user}/{app_name}'
          self.config['LHA_SOURCE_CODE_PATH'] = workspace_path
        else:
          # Fallback to shared workspace
          workspace_path = f'/Workspace/Shared/{app_name}'
          self.config['LHA_SOURCE_CODE_PATH'] = workspace_path
          print('⚠️  Could not determine current user - using shared workspace path')

        print(f'📁 App source code path: {workspace_path}')

      # Debug: print current config
      print(f'📋 Current config: {self.config}')

      # Complete configuration with defaults
      complete_config = self.env_detector.generate_environment_config(self.config)

      # Debug: print complete config
      print(f'📋 Complete config: {complete_config}')

      # Write .env.local file
      env_file = self.project_root / '.env.local'
      with open(env_file, 'w') as f:
        f.write(f'# Generated by auto-setup.py on {self._get_timestamp()}\n\n')

        for key, value in complete_config.items():
          f.write(f'{key}="{value}"\n')

      print(f'✅ Created {env_file}')
      return True
    except Exception as e:
      print(f'❌ Failed to create environment file: {e}')
      return False

  def _install_dependencies(self) -> bool:
    """Install Python and frontend dependencies."""
    print('📦 Installing dependencies...')

    if self.dry_run:
      print('   [DRY RUN] Would install dependencies')
      return True

    try:
      # Install Python dependencies with uv
      print('   Installing Python dependencies...')
      result = subprocess.run(['uv', 'sync'], cwd=self.project_root, text=True)
      if result.returncode != 0:
        print('❌ Failed to install Python dependencies')
        return False

      # Install frontend dependencies with bun
      print('   Installing frontend dependencies...')
      client_dir = self.project_root / 'client'
      result = subprocess.run(['bun', 'install'], cwd=client_dir, text=True)
      if result.returncode != 0:
        print('❌ Failed to install frontend dependencies')
        return False

      print('✅ Dependencies installed successfully')
      return True
    except Exception as e:
      print(f'❌ Failed to install dependencies: {e}')
      return False

  def _load_sample_data(self) -> bool:
    """Load sample data using setup scripts in order."""
    print('📊 Loading sample data...')

    if self.dry_run:
      print('   [DRY RUN] Would load sample data')
      return True

    try:
      # Use the existing load_sample_data.sh script
      script_path = self.project_root / 'load_sample_data.sh'
      if not script_path.exists():
        print(f'❌ Script {script_path} not found')
        return False

      # Run the shell script and stream output
      result = subprocess.run(
        ['./load_sample_data.sh'], cwd=self.project_root, env=os.environ.copy(), text=True
      )

      return result.returncode == 0

    except Exception as e:
      print(f'❌ Failed to load sample data: {e}')
      return False

  def _validate_local_setup(self) -> bool:
    """Validate local setup by running development server briefly."""
    print('✅ Validating local setup...')

    if self.dry_run:
      print('   [DRY RUN] Would validate local setup')
      return True

    # For now, just check that the setup files are in place
    # A full test would involve starting the server and testing endpoints

    env_file = self.project_root / '.env.local'
    if not env_file.exists():
      print('❌ .env.local file not found')
      return False

    print('✅ Local setup validation passed')
    return True

  def _deploy_app(self) -> bool:
    """Deploy the application."""
    deployment_mode = self.config.get('DEPLOYMENT_MODE', 'full_deployment')

    if deployment_mode == 'notebook_only':
      print('📓 Syncing notebooks to workspace...')
    else:
      print('🚀 Deploying application...')

    if self.dry_run:
      print('   [DRY RUN] Would deploy application')
      return True

    try:
      # Choose deploy command based on deployment mode
      if deployment_mode == 'notebook_only':
        # Only sync notebooks, skip app deployment
        result = subprocess.run(['./deploy.sh', '--sync-only'], cwd=self.project_root, text=True)
        success_message = '✅ Notebooks synced to workspace successfully'
      else:
        # Full deployment including app
        result = subprocess.run(['./deploy.sh'], cwd=self.project_root, text=True)
        success_message = '✅ Application deployed successfully'

      if result.returncode != 0:
        print('❌ Deployment failed')
        return False

      print(success_message)
      return True
    except Exception as e:
      print(f'❌ Failed to deploy application: {e}')
      return False

  def _validate_deployment(self) -> bool:
    """Validate deployment."""
    print('✅ Validating deployment...')

    if self.dry_run:
      print('   [DRY RUN] Would validate deployment')
      return True

    app_name = self.config['DATABRICKS_APP_NAME']

    # Wait for app to be ready
    if not self.validator.wait_for_app_ready(app_name, timeout_minutes=5):
      print('⚠️  App deployment validation timed out')
      return False

    # Run deployment validation
    valid, issues = self.validator.validate_deployment(app_name)

    if not valid:
      print('❌ Deployment validation failed:')
      for issue in issues:
        print(f'   • {issue}')
      return False

    print('✅ Deployment validation passed')
    return True

  def _run_integration_tests(self) -> bool:
    """Run integration tests."""
    print('🧪 Running integration tests...')

    if self.dry_run:
      print('   [DRY RUN] Would run integration tests')
      return True

    valid, issues = self.validator.run_integration_tests(self.config)

    if not valid:
      print('❌ Integration tests failed:')
      for issue in issues:
        print(f'   • {issue}')
      return False

    print('✅ Integration tests passed')
    return True

  def _get_step_result(self, step_id: str) -> Dict[str, Any]:
    """Get result data for a completed step."""
    if step_id == 'collect_user_input':
      return {'config': self.config.copy()}
    elif step_id == 'create_experiment':
      return {'experiment_id': self.config.get('MLFLOW_EXPERIMENT_ID')}
    elif step_id == 'create_app':
      return {'app_name': self.config.get('DATABRICKS_APP_NAME')}
    elif step_id == 'create_catalog_schema':
      return {'schema': self.created_resources.get('schema')}
    return {}

  def _get_app_url(self, app_name: str) -> str:
    """Get the URL of the deployed Databricks App."""
    try:
      app = self.client.apps.get(app_name)
      if hasattr(app, 'url') and app.url:
        return app.url
    except Exception as e:
      print(f'⚠️  Could not get app URL from API: {e}')

    # Fallback to constructed URL
    workspace_host = self.config.get('DATABRICKS_HOST', '').rstrip('/')
    return f'{workspace_host}/apps/{app_name}'

  def _get_notebook_url(self, notebook_name: str) -> str:
    """Generate the direct URL to a notebook in the Databricks workspace."""
    workspace_host = self._ensure_https_protocol(self.config.get('DATABRICKS_HOST', '')).rstrip('/')
    lha_source_code_path = self.config.get('LHA_SOURCE_CODE_PATH')

    for i in self.client.workspace.list(
      f'{lha_source_code_path}/mlflow_demo/notebooks', recursive=True
    ):
      if i.path and i.path.endswith(notebook_name):
        return f'{workspace_host}/editor/notebooks/{i.resource_id}'
    return 'NOT FOUND'

  def _ensure_https_protocol(self, host: str | None) -> str:
    """Ensure the host URL has HTTPS protocol."""
    if not host:
      return ''

    if host.startswith('https://') or host.startswith('http://'):
      return host

    return f'https://{host}'

  def _get_experiment_url(self, experiment_id: str) -> str:
    """Generate the direct URL to the MLflow experiment."""
    workspace_host = self.config.get('DATABRICKS_HOST', '').rstrip('/')
    return f'{workspace_host}/#mlflow/experiments/{experiment_id}'

  def _show_final_results(self, success: bool):
    """Show final setup results."""
    print('\n' + '=' * 60)
    print('🎉 SETUP COMPLETE!')
    print('=' * 60)

    if success:
      print('\n✅ Your MLflow demo environment is ready to use!')

      deployment_mode = self.config.get('DEPLOYMENT_MODE', 'full_deployment')
      workspace_host = self.config.get('DATABRICKS_HOST', '').rstrip('/')

      # Show the primary access URL prominently
      print('\n🔗 YOUR PRIMARY ACCESS LINK:')
      print('=' * 40)

      if deployment_mode == 'full_deployment':
        app_name = self.config.get('DATABRICKS_APP_NAME', 'mlflow_demo_app')
        app_url = self._get_app_url(app_name)
        print(f'📱 Databricks App: {app_url}')
        print('   ↳ Interactive demo application ready to use')
      else:
        workspace_path = self.config.get('LHA_SOURCE_CODE_PATH', '/Workspace/...')
        notebook_url = self._get_notebook_url('0_demo_overview')
        print(f'📓 Demo Overview Notebook: {notebook_url}')
        print('   ↳ Start here for interactive learning experience')

      print('\n📋 Resources Created:')
      print('-' * 30)

      if 'MLFLOW_EXPERIMENT_ID' in self.config:
        experiment_id = self.config['MLFLOW_EXPERIMENT_ID']
        experiment_url = self._get_experiment_url(experiment_id)
        print(f'🧪 MLflow Experiment: {experiment_url}')
        print(f'   ↳ ID: {experiment_id}')

      if 'UC_CATALOG' in self.config and 'UC_SCHEMA' in self.config:
        catalog = self.config['UC_CATALOG']
        schema = self.config['UC_SCHEMA']
        catalog_url = f'{workspace_host}/#unitycatalog/catalogs/{catalog}/schemas/{schema}'
        print(f'📁 Unity Catalog Schema: {catalog_url}')
        print(f'   ↳ {catalog}.{schema}')

      if deployment_mode == 'full_deployment':
        app_name = self.config.get('DATABRICKS_APP_NAME', 'mlflow_demo_app')
        apps_url = f'{workspace_host}/#apps'
        print(f'📱 Databricks Apps Console: {apps_url}')
        print(f'   ↳ Manage app: {app_name}')

      # # Workspace notebooks (available in both modes)
      # workspace_path = self.config.get('LHA_SOURCE_CODE_PATH', '/Workspace/...')
      # workspace_url = f"{workspace_host}/#workspace{workspace_path}"
      # print(f'📂 Workspace Files: {workspace_url}')
      # print(f'   ↳ All demo notebooks and source code')

      print('\n🚀 Quick Start Guide:')
      print('-' * 30)
      if deployment_mode == 'notebook_only':
        print('1. 📖 Follow the step-by-step interactive guide in the notebook')
        print('2. 🔬 Explore notebooks in mlflow_demo/notebooks/')
        print('3. 📊 Review your experiment data and traces')
        print('4. 🧪 Run your own experiments using the provided examples')
      else:
        print('1. 🎯 Click the Databricks App link above')
        print('2. 🧪 Try the email generation demo')
        print('3. 📝 Submit feedback to see MLflow tracing in action')
        print('4. 📊 Explore the MLflow experiment for evaluation data')
        print('5. 📂 Check out the workspace notebooks for deeper learning')

      # Show detailed progress
      print('\n📊 Setup Progress:')
      self.progress.show_detailed_progress()

      # For notebook-only mode, show the primary access link at the very bottom in a super obvious way
      if deployment_mode == 'notebook_only':
        print('\n\n' + '=' * 80)
        print('🚨 🚨 🚨  YOUR NOTEBOOK IS READY - CLICK HERE TO START  🚨 🚨 🚨')
        print('=' * 80)
        workspace_path = self.config.get('LHA_SOURCE_CODE_PATH', '/Workspace/...')
        notebook_url = self._get_notebook_url('0_demo_overview')
        print(f'\n🎯 👉 START HERE: {notebook_url}')
        print('\n   ↳ This opens the Demo Overview Notebook - your starting point!')
        print('   ↳ Follow the step-by-step interactive guide')
        print('   ↳ Learn how to use MLflow to improve GenAI quality')
        print('\n' + '=' * 80)

      # For full deployment mode, show the app URL prominently
      elif deployment_mode == 'full_deployment':
        print('\n\n' + '=' * 80)
        print('🚀 🚀 🚀  YOUR APP IS DEPLOYED - CLICK HERE TO START  🚀 🚀 🚀')
        print('=' * 80)
        app_name = self.config.get('DATABRICKS_APP_NAME', 'mlflow_demo_app')
        app_url = self._get_app_url(app_name)
        print(f'\n🎯 👉 START HERE: {app_url}')
        print('\n   ↳ This opens your deployed MLflow Demo App')
        print('   ↳ Interactive web interface with all features')
        print('   ↳ Learn how to use MLflow to improve GenAI quality')
        print('\n' + '=' * 80)

    else:
      print('❌ Setup failed or was interrupted')
      print('\n🔧 Troubleshooting:')
      print("   • Run './auto-setup.sh --resume' to continue from where you left off")
      print("   • Check the progress with './auto-setup.sh --status'")
      print('   • Review any error messages above')

      # Show failed steps
      failed_steps = self.progress.get_failed_steps()
      if failed_steps:
        print(f'\n❌ Failed steps: {", ".join(failed_steps)}')

  def _get_timestamp(self) -> str:
    """Get current timestamp string."""
    from datetime import datetime

    return datetime.now().strftime('%Y-%m-%d %H:%M:%S')

  def _ensure_databricks_auth(self, skip_prompts: bool = False) -> bool:
    """Ensure Databricks authentication is configured."""
    if skip_prompts:
      # For resume, first check if current auth is working
      try:
        from databricks.sdk import WorkspaceClient

        test_client = WorkspaceClient()
        test_client.current_user.me()
        print('🔐 Using existing Databricks authentication')
        return True
      except Exception:
        pass  # Fall through to profile selection

    print('🔐 Databricks Profile Selection')

    # Get profiles
    profiles = self._get_databricks_profiles()

    if not profiles:
      print("❌ No Databricks profiles found. Please run 'databricks auth login' first.")
      return False

    # Always show profile selection menu
    return self._handle_auth_selection(profiles)

  def _get_databricks_profiles(self) -> Dict[str, Dict[str, Any]]:
    """Get available Databricks authentication profiles."""
    try:
      spinner = Spinner('Loading Databricks authentication profiles...')
      spinner.start()
      try:
        result = subprocess.run(
          ['databricks', 'auth', 'profiles'], capture_output=True, text=True, timeout=30
        )
        spinner.stop('Loaded auth profiles')
      except Exception as e:
        spinner.stop()
        raise e

      if result.returncode != 0:
        print(f'❌ Failed to get profiles: {result.stderr}')
        return {}

      # Parse the profiles output (table format)
      profiles = {}
      lines = result.stdout.strip().split('\n')

      # Skip header line and parse each profile
      for line in lines[1:]:  # Skip the "Name Host Valid" header
        if line.strip():
          parts = line.split(None, 2)  # Split on whitespace, max 3 parts
          if len(parts) >= 2:
            profile_name = parts[0].strip()
            host = parts[1].strip()
            valid = parts[2].strip() if len(parts) > 2 else 'UNKNOWN'

            profiles[profile_name] = {'host': host, 'valid': valid}

      return profiles

    except subprocess.TimeoutExpired:
      print('❌ Timeout getting profiles')
      return {}
    except FileNotFoundError:
      print("❌ 'databricks' command not found")
      return {}
    except Exception as e:
      print(f'❌ Error getting profiles: {e}')
      return {}

  def _handle_auth_selection(self, profiles: Dict[str, Dict[str, Any]]) -> bool:
    """Handle profile selection and authentication."""
    print('\n🔧 Databricks Profile Configuration')
    print('⚠️  This script only works with the DEFAULT databricks-cli profile.')

    profile_list = list(profiles.keys())
    if not profile_list:
      print("❌ No profiles available. Please run 'databricks auth login' to create one.")
      return False

    # Check if DEFAULT profile exists and is valid
    if 'DEFAULT' not in profiles:
      print('❌ DEFAULT profile not found.')
      print('   Please create a DEFAULT profile by running:')
      print('   databricks auth login --profile DEFAULT')
      return False

    default_profile = profiles['DEFAULT']
    if default_profile.get('valid') != 'YES':
      print('❌ DEFAULT profile exists but is not valid.')
      print('   Please re-authenticate your DEFAULT profile by running:')
      print('   databricks auth login --profile DEFAULT')
      return False

    # Check if current auth is working and using DEFAULT
    current_auth_works = False
    current_profile = None
    try:
      from databricks.sdk import WorkspaceClient

      test_client = WorkspaceClient()
      test_client.current_user.me()
      current_auth_works = True
      # Try to determine current profile from workspace URL
      current_host = test_client.config.host
      for profile_name, profile_info in profiles.items():
        if profile_info.get('host') == current_host:
          current_profile = profile_name
          break
    except Exception:
      pass

    print('\n🔧 Available option:')
    default_host = default_profile.get('host', 'Unknown host')

    if current_auth_works and current_profile == 'DEFAULT':
      print(f'   0. Keep current DEFAULT profile authentication ({default_host}) ✅')
      choice = input("\nPress ENTER to continue with DEFAULT profile or 'q' to quit: ").strip()
      if choice.lower() == 'q':
        return False
      print('✅ Using DEFAULT profile')
      return True
    else:
      print(f'   1. Use DEFAULT profile ({default_host}) ✅')
      choice = input("\nPress ENTER to use DEFAULT profile or 'q' to quit: ").strip()
      if choice.lower() == 'q':
        return False
      selected_profile = 'DEFAULT'

    # Authenticate with selected profile
    print(f"\n🔐 Authenticating with profile '{selected_profile}'...")

    profile_info = profiles[selected_profile]
    host = profile_info.get('host')

    if not host:
      print(f"❌ No host found for profile '{selected_profile}'")
      return False

    try:
      # Run databricks auth login with the specific host and profile
      result = subprocess.run(
        ['databricks', 'auth', 'login', '--host', host, '--profile', selected_profile],
        timeout=120,  # 2 minutes timeout for auth
      )

      if result.returncode == 0:
        print(f"✅ Successfully authenticated with profile '{selected_profile}'")

        # Test the authentication
        try:
          from databricks.sdk import WorkspaceClient

          test_client = WorkspaceClient(profile=selected_profile)
          test_client.current_user.me()
          print('✅ Authentication test successful')
          return True
        except Exception as e:
          print(f'❌ Authentication test failed: {e}')
          return False
      else:
        print(f"❌ Authentication failed for profile '{selected_profile}'")
        return False

    except subprocess.TimeoutExpired:
      print('❌ Authentication timed out')
      return False
    except Exception as e:
      print(f'❌ Authentication error: {e}')
      return False

  def cleanup_resources(self):
    """Clean up created resources (for rollback)."""
    print('🧹 Cleaning up created resources...')
    if hasattr(self, 'resource_manager'):
      self.resource_manager.cleanup_created_resources()


def main():
  """Main entry point."""
  parser = argparse.ArgumentParser(
    description='MLflow Demo Automated Setup',
    formatter_class=argparse.RawDescriptionHelpFormatter,
    epilog=__doc__,
  )

  parser.add_argument(
    '--dry-run',
    action='store_true',
    help='Show what would be created without actually creating resources',
  )
  parser.add_argument(
    '--resume', action='store_true', help='Resume from previous failed/interrupted setup'
  )
  parser.add_argument('--validate-only', action='store_true', help='Only run validation checks')
  parser.add_argument('--status', action='store_true', help='Show current setup status')
  parser.add_argument(
    '--cleanup', action='store_true', help='Clean up progress file and start fresh'
  )

  args = parser.parse_args()

  if args.status:
    # Show status and exit
    progress = ProgressTracker()
    progress.show_detailed_progress()
    return

  if args.cleanup:
    # Clean up and exit
    progress = ProgressTracker()
    progress.cleanup_progress_file()
    print('🧹 Progress file cleaned up. Run auto-setup.py to start fresh.')
    return

  # Initialize setup
  auto_setup = AutoSetup(dry_run=args.dry_run)

  if args.validate_only:
    # Only run validation
    print('🔍 Running validation checks only...')
    valid, issues = auto_setup.validator.validate_prerequisites()
    if valid:
      print('✅ All validation checks passed')
      sys.exit(0)
    else:
      print('❌ Validation failed')
      for issue in issues:
        print(f'   • {issue}')
      sys.exit(1)

  # Run the setup
  try:
    success = auto_setup.run_setup(resume=args.resume)
    sys.exit(0 if success else 1)
  except KeyboardInterrupt:
    print('\n⚠️  Setup interrupted. Run with --resume to continue.')
    sys.exit(1)
  except Exception as e:
    print(f'\n❌ Unexpected error: {e}')
    print('Run with --resume to try continuing from the last successful step.')
    sys.exit(1)


if __name__ == '__main__':
  main()