"""Deployment orchestration for Tellr on Databricks Apps.

This module provides the main create/update/delete functions for deploying
the Tellr AI slide generator to Databricks Apps from a notebook.
"""

import logging
import shutil
import uuid
from contextlib import contextmanager
from importlib import resources
from pathlib import Path
from string import Template
from typing import Iterator, Optional

import psycopg2
import psycopg2.extensions
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.apps import (
    App,
    AppDeployment,
    AppResource,
    AppResourceDatabase,
    AppResourceDatabaseDatabasePermission,
    ComputeSize,
)
from databricks.sdk.service.database import DatabaseInstance
from databricks.sdk.service.workspace import ImportFormat

logger = logging.getLogger(__name__)


class DeploymentError(Exception):
    """Raised when deployment fails."""

    pass


def _get_workspace_client(
    workspace_client: Optional[WorkspaceClient] = None,
    profile: Optional[str] = None,
) -> WorkspaceClient:
    """Get WorkspaceClient with priority: external > profile > environment.

    Args:
        workspace_client: Externally created WorkspaceClient to use directly
        profile: Profile name from .databrickscfg file
        (If neither provided, uses environment variables)

    Returns:
        WorkspaceClient instance
    """
    if workspace_client is not None:
        return workspace_client
    if profile is not None:
        return WorkspaceClient(profile=profile)
    return WorkspaceClient()


def create(
    lakebase_name: str,
    schema_name: str,
    app_name: str,
    app_file_workspace_path: str,
    workspace_client: Optional[WorkspaceClient] = None,
    profile: Optional[str] = None,
    lakebase_compute: str = "CU_1",
    app_compute: str = "MEDIUM",
    app_version: Optional[str] = None,
    description: str = "Tellr AI Slide Generator",
) -> dict:
    """Create and deploy Tellr to Databricks Apps.

    This function creates all necessary infrastructure and deploys the app:
    1. Creates/gets Lakebase database instance
    2. Generates requirements.txt with pinned app version
    3. Generates app.yaml with environment variables
    4. Uploads files to workspace
    5. Creates Databricks App with database resource
    6. Sets up database schema, creates tables, and seeds default data

    Args:
        lakebase_name: Name for the Lakebase database instance
        schema_name: PostgreSQL schema name for app tables
        app_name: Name for the Databricks App
        app_file_workspace_path: Workspace path to upload app files
        workspace_client: Externally created WorkspaceClient (highest priority)
        profile: Profile name from .databrickscfg file (second priority)
        lakebase_compute: Lakebase capacity (CU_1, CU_2, CU_4, CU_8)
        app_compute: App compute size (MEDIUM, LARGE, LIQUID)
        app_version: Specific databricks-tellr-app version (default: latest)
        description: App description

    Returns:
        Dictionary with deployment info:
        - url: App URL
        - app_name: Created app name
        - lakebase_name: Database instance name
        - schema_name: Schema name
        - status: "created"

    Raises:
        DeploymentError: If deployment fails
    """
    print(f"🚀 Deploying Tellr to Databricks Apps...")
    print(f"   App name: {app_name}")
    print(f"   Workspace path: {app_file_workspace_path}")
    print(f"   Lakebase: {lakebase_name} (capacity: {lakebase_compute})")
    print(f"   Schema: {schema_name}")
    print()

    ws = _get_workspace_client(workspace_client, profile)

    try:
        # Step 1: Create/get Lakebase instance
        print("📊 Setting up Lakebase database...")
        lakebase_result = _get_or_create_lakebase(ws, lakebase_name, lakebase_compute)
        print(f"   ✅ Lakebase: {lakebase_result['name']} ({lakebase_result['status']})")
        print()

        # Step 2: Generate and upload files
        print("📁 Preparing deployment files...")
        with _staging_dir(app_file_workspace_path) as staging:
            # Generate requirements.txt
            _write_requirements(staging, app_version)
            print("   ✓ Generated requirements.txt")

            # Generate app.yaml
            _write_app_yaml(staging, lakebase_name, schema_name)
            print("   ✓ Generated app.yaml")

            # Upload to workspace
            print(f"☁️  Uploading to: {app_file_workspace_path}")
            _upload_files(ws, staging, app_file_workspace_path)
            print("   ✅ Files uploaded")
        print()

        # Step 3: Create app
        print(f"🔧 Creating Databricks App: {app_name}")
        app = _create_app(
            ws,
            app_name=app_name,
            description=description,
            workspace_path=app_file_workspace_path,
            compute_size=app_compute,
            lakebase_name=lakebase_name,
        )
        print(f"   ✅ App created")
        if app.url:
            print(f"   🌐 URL: {app.url}")
        print()

        # Step 4: Set up database schema and permissions
        print("📊 Setting up database schema...")
        _setup_database_schema(ws, app, lakebase_name, schema_name)
        print(f"   ✅ Schema '{schema_name}' configured")
        print()

        # Note: Database tables are created by the app on first startup
        print("✅ Deployment complete!")
        return {
            "url": app.url,
            "app_name": app_name,
            "lakebase_name": lakebase_name,
            "schema_name": schema_name,
            "status": "created",
        }

    except Exception as e:
        raise DeploymentError(f"Deployment failed: {e}") from e


def update(
    app_name: str,
    app_file_workspace_path: str,
    lakebase_name: str,
    schema_name: str,
    workspace_client: Optional[WorkspaceClient] = None,
    profile: Optional[str] = None,
    app_version: Optional[str] = None,
    update_database: bool = False,
) -> dict:
    """Deploy a new version of an existing Tellr app.

    Updates the app files and triggers a new deployment. Optionally resets
    the database schema (drops all tables and recreates).

    Args:
        app_name: Name of the existing Databricks App
        app_file_workspace_path: Workspace path with app files
        lakebase_name: Lakebase instance name
        schema_name: Schema name
        workspace_client: Externally created WorkspaceClient (highest priority)
        profile: Profile name from .databrickscfg file (second priority)
        app_version: Specific databricks-tellr-app version (default: latest)
        update_database: If True, drops schema and recreates tables (hard reset)

    Returns:
        Dictionary with deployment info

    Raises:
        DeploymentError: If update fails
    """
    print(f"🔄 Updating Tellr app: {app_name}")

    ws = _get_workspace_client(workspace_client, profile)

    try:
        # Generate and upload updated files
        with _staging_dir(app_file_workspace_path) as staging:
            _write_requirements(staging, app_version)
            _write_app_yaml(staging, lakebase_name, schema_name)
            _upload_files(ws, staging, app_file_workspace_path)
            print("   ✅ Files updated")

        # Optionally reset database
        if update_database:
            print("📊 Resetting database schema...")
            app = ws.apps.get(name=app_name)

            # Drop existing schema
            print(f"   ⏳ Dropping schema '{schema_name}'...")
            _drop_schema(ws, lakebase_name, schema_name)
            print(f"   ✓ Schema dropped")

            # Recreate schema with permissions
            print(f"   ⏳ Recreating schema...")
            _setup_database_schema(ws, app, lakebase_name, schema_name)
            print(f"   ✅ Schema recreated (tables will be created on app startup)")
            print()

        # Trigger new deployment
        print("   ⏳ Deploying...")
        deployment = AppDeployment(source_code_path=app_file_workspace_path)
        result = ws.apps.deploy_and_wait(app_name=app_name, app_deployment=deployment)
        print(f"   ✅ Deployment completed: {result.deployment_id}")

        app = ws.apps.get(name=app_name)
        if app.url:
            print(f"   🌐 URL: {app.url}")

        return {
            "url": app.url,
            "app_name": app_name,
            "deployment_id": result.deployment_id,
            "status": "updated",
            "database_reset": update_database,
        }

    except Exception as e:
        raise DeploymentError(f"Update failed: {e}") from e


def delete(
    app_name: str,
    workspace_client: Optional[WorkspaceClient] = None,
    profile: Optional[str] = None,
    lakebase_name: Optional[str] = None,
    schema_name: Optional[str] = None,
    delete_schema: bool = False,
    truncate: bool = False,
) -> dict:
    """Delete a Tellr app with optional schema cleanup.

    By default, only deletes the app. Use delete_schema or truncate to also
    handle the database schema.

    Args:
        app_name: Name of the app to delete
        workspace_client: Externally created WorkspaceClient (highest priority)
        profile: Profile name from .databrickscfg file (second priority)
        lakebase_name: Lakebase instance name (required if delete_schema or truncate)
        schema_name: Schema name (required if delete_schema or truncate)
        delete_schema: If True, drops schema permanently (full cleanup)
        truncate: If True, drops and recreates schema (data reset for future use)

    Returns:
        Dictionary with deletion status

    Raises:
        DeploymentError: If deletion fails or invalid options provided
    """
    # Validate options
    if delete_schema and truncate:
        raise DeploymentError(
            "Cannot specify both delete_schema and truncate. Choose one."
        )

    if (delete_schema or truncate) and (not lakebase_name or not schema_name):
        raise DeploymentError(
            "lakebase_name and schema_name are required when using "
            "delete_schema or truncate options."
        )

    print(f"🗑️  Deleting app: {app_name}")

    ws = _get_workspace_client(workspace_client, profile)

    try:
        result = {"app_name": app_name, "status": "deleted"}

        # Handle schema operations before deleting app
        if delete_schema:
            print(f"   ⏳ Dropping schema '{schema_name}'...")
            _drop_schema(ws, lakebase_name, schema_name)
            print(f"   ✅ Schema dropped")
            result["schema_deleted"] = True

        elif truncate:
            print(f"   ⏳ Resetting schema '{schema_name}'...")
            # Get app info before deletion for permissions
            app = ws.apps.get(name=app_name)

            # Drop schema
            _drop_schema(ws, lakebase_name, schema_name)
            print(f"   ✓ Schema dropped")

            # Recreate schema with permissions
            _setup_database_schema(ws, app, lakebase_name, schema_name)
            print(f"   ✅ Schema reset (tables will be created on app startup)")
            result["schema_truncated"] = True

        # Delete the app
        ws.apps.delete(name=app_name)
        print("   ✅ App deleted")

        return result

    except Exception as e:
        raise DeploymentError(f"Deletion failed: {e}") from e


# -----------------------------------------------------------------------------
# Internal functions
# -----------------------------------------------------------------------------


def _get_or_create_lakebase(
    ws: WorkspaceClient, database_name: str, capacity: str
) -> dict:
    """Get or create a Lakebase database instance."""
    try:
        # Check if exists
        existing = ws.database.get_database_instance(name=database_name)
        return {
            "name": existing.name,
            "status": "exists",
            "state": existing.state.value if existing.state else "UNKNOWN",
        }
    except Exception as e:
        error_str = str(e).lower()
        if "not found" not in error_str and "does not exist" not in error_str:
            raise

    # Create new instance
    instance = ws.database.create_database_instance_and_wait(
        DatabaseInstance(name=database_name, capacity=capacity)
    )
    return {
        "name": instance.name,
        "status": "created",
        "state": instance.state.value if instance.state else "RUNNING",
    }


def _write_requirements(staging_dir: Path, app_version: Optional[str]) -> None:
    """Generate requirements.txt with app package."""
    if app_version and not _is_valid_version(app_version):
        raise DeploymentError(
            f"Invalid app version '{app_version}'. Expected a PEP 440 version."
        )

    if app_version:
        package_line = f"databricks-tellr-app=={app_version}"
    else:
        package_line = "databricks-tellr-app"

    content = "\n".join(
        [
            "# Generated by databricks-tellr setup",
            package_line,
        ]
    )
    (staging_dir / "requirements.txt").write_text(content)


def _write_app_yaml(staging_dir: Path, lakebase_name: str, schema_name: str) -> None:
    """Generate app.yaml with environment variables."""
    template_content = _load_template("app.yaml.template")
    content = Template(template_content).substitute(
        LAKEBASE_INSTANCE=lakebase_name,
        LAKEBASE_SCHEMA=schema_name,
    )
    (staging_dir / "app.yaml").write_text(content)


def _load_template(template_name: str) -> str:
    """Load a template file from package resources."""
    try:
        # Python 3.9+ style
        files = resources.files("databricks_tellr") / "_templates" / template_name
        return files.read_text()
    except (TypeError, AttributeError):
        # Fallback for older Python
        with resources.open_text(
            "databricks_tellr._templates", template_name
        ) as f:
            return f.read()


def _is_valid_version(version: str) -> bool:
    try:
        from packaging.version import Version
    except ImportError:
        return True
    try:
        Version(version)
        return True
    except Exception:
        return False


@contextmanager
def _staging_dir(app_file_workspace_path: str) -> Iterator[Path]:
    base_path = Path(app_file_workspace_path)
    staging_dir = base_path / f".tellr_staging_{uuid.uuid4().hex}"
    try:
        staging_dir.mkdir(parents=True, exist_ok=False)
    except Exception as exc:
        raise DeploymentError(
            f"Failed to create staging directory at {staging_dir}: {exc}"
        ) from exc

    try:
        yield staging_dir
    finally:
        shutil.rmtree(staging_dir, ignore_errors=True)


@contextmanager
def _get_db_connection(
    ws: WorkspaceClient, lakebase_name: str
) -> Iterator[psycopg2.extensions.connection]:
    """Get PostgreSQL connection to Lakebase instance.

    This context manager handles credential generation, connection setup,
    and cleanup. The connection is set to autocommit mode.

    Args:
        ws: WorkspaceClient for database access
        lakebase_name: Lakebase instance name

    Yields:
        psycopg2 connection with autocommit enabled
    """
    instance = ws.database.get_database_instance(name=lakebase_name)
    user = ws.current_user.me().user_name

    cred = ws.database.generate_database_credential(
        request_id=str(uuid.uuid4()),
        instance_names=[lakebase_name],
    )

    conn = psycopg2.connect(
        host=instance.read_write_dns,
        port=5432,
        user=user,
        password=cred.token,
        dbname="databricks_postgres",
        sslmode="require",
    )
    conn.autocommit = True

    try:
        yield conn
    finally:
        conn.close()


def _drop_schema(ws: WorkspaceClient, lakebase_name: str, schema_name: str) -> None:
    """Drop schema and all its contents (CASCADE).

    This is a destructive operation that removes the schema and all tables,
    sequences, and other objects within it.

    Args:
        ws: WorkspaceClient for database access
        lakebase_name: Lakebase instance name
        schema_name: Schema name to drop
    """
    with _get_db_connection(ws, lakebase_name) as conn:
        with conn.cursor() as cur:
            cur.execute(f'DROP SCHEMA IF EXISTS "{schema_name}" CASCADE')


def _upload_files(
    ws: WorkspaceClient, staging_dir: Path, workspace_path: str
) -> None:
    """Upload files from staging directory to workspace."""
    # Ensure directory exists
    try:
        ws.workspace.mkdirs(workspace_path)
    except Exception:
        pass  # May already exist

    # Upload each file
    for file_path in staging_dir.iterdir():
        if file_path.is_file():
            workspace_file_path = f"{workspace_path}/{file_path.name}"
            with open(file_path, "rb") as f:
                ws.workspace.upload(
                    workspace_file_path,
                    f,
                    format=ImportFormat.AUTO,
                    overwrite=True,
                )


def _create_app(
    ws: WorkspaceClient,
    app_name: str,
    description: str,
    workspace_path: str,
    compute_size: str,
    lakebase_name: str,
) -> App:
    """Create Databricks App with database resource."""
    compute_size_enum = ComputeSize(compute_size)

    # Database resource
    resources = [
        AppResource(
            name="app_database",
            database=AppResourceDatabase(
                instance_name=lakebase_name,
                database_name="databricks_postgres",
                permission=AppResourceDatabaseDatabasePermission.CAN_CONNECT_AND_CREATE,
            ),
        )
    ]

    # Create app
    app = App(
        name=app_name,
        description=description,
        compute_size=compute_size_enum,
        default_source_code_path=workspace_path,
        resources=resources,
        user_api_scopes=[
            "sql",
            "dashboards.genie",
            "catalog.tables:read",
            "catalog.schemas:read",
            "catalog.catalogs:read",
            "serving.serving-endpoints",
        ],
    )

    result = ws.apps.create_and_wait(app)

    # Trigger initial deployment
    deployment = AppDeployment(source_code_path=workspace_path)
    ws.apps.deploy_and_wait(app_name=app_name, app_deployment=deployment)

    # Refresh to get URL
    return ws.apps.get(name=app_name)


def _setup_database_schema(
    ws: WorkspaceClient, app: App, lakebase_name: str, schema_name: str
) -> None:
    """Set up database schema and grant permissions to app."""
    # Get app's service principal client ID
    client_id = None
    if hasattr(app, "service_principal_client_id") and app.service_principal_client_id:
        client_id = app.service_principal_client_id
    elif hasattr(app, "service_principal_id") and app.service_principal_id:
        client_id = str(app.service_principal_id)

    if not client_id:
        print("   ⚠️  Could not get app client ID - schema setup skipped")
        return

    with _get_db_connection(ws, lakebase_name) as conn:
        with conn.cursor() as cur:
            cur.execute(f'CREATE SCHEMA IF NOT EXISTS "{schema_name}"')
            cur.execute(f'GRANT USAGE ON SCHEMA "{schema_name}" TO "{client_id}"')
            cur.execute(f'GRANT CREATE ON SCHEMA "{schema_name}" TO "{client_id}"')
            cur.execute(
                f'GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA "{schema_name}" TO "{client_id}"'
            )
            cur.execute(
                f'GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA "{schema_name}" TO "{client_id}"'
            )
            cur.execute(
                f'ALTER DEFAULT PRIVILEGES IN SCHEMA "{schema_name}" '
                f'GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO "{client_id}"'
            )
            cur.execute(
                f'ALTER DEFAULT PRIVILEGES IN SCHEMA "{schema_name}" '
                f'GRANT USAGE, SELECT ON SEQUENCES TO "{client_id}"'
            )
