"""
Database Module
Handles all database connections with step-by-step query execution
"""

from pathlib import Path

import pandas as pd
import streamlit as st
from snowflake.snowpark.context import get_active_session


@st.cache_resource
def get_snowpark_session():  # type: ignore[misc]
    """
    Get cached Snowpark session for Container Runtime
    All timestamps are in UTC

    Returns:
        Active Snowpark session
    """
    try:
        session = get_active_session()
        return session
    except Exception as e:
        st.error(f'Failed to get Snowpark session: {str(e)}')
        raise


def load_query_from_file(query_file: str) -> str:
    """
    Load SQL query from file

    Args:
        query_file: Name of the query file in queries/ directory

    Returns:
        SQL query string
    """
    # Path is relative to project root (queries/ is at root level)
    query_path = Path(__file__).parent / query_file
    with open(query_path, 'r') as f:
        return f.read()


def resolve_email_to_uuid(email: str) -> str | None:
    """
    Resolve an email address to an identity UUID using Snowpark filters

    Args:
        email: Email address to resolve

    Returns:
        UUID string if found, None otherwise
    """
    try:
        from snowflake.snowpark.functions import col, lit, lower

        session = get_snowpark_session()

        # Use Snowpark DataFrame API instead of string substitution
        result = (
            session.table('FACTS.PROD.IDENTITY')
            .filter(lower(col('EMAIL')) == lower(lit(email)))
            .select('ID')
            .limit(1)
            .collect()
        )

        if not result:
            return None

        return str(result[0]['ID'])

    except Exception as e:
        st.error(f'Failed to resolve email to UUID: {str(e)}')
        return None


def get_identity(identity_id: str) -> pd.DataFrame | None:
    """
    Step 1: Load identity data from FACT.PROD.IDENTITY

    Args:
        identity_id: UUID of the identity

    Returns:
        DataFrame with identity information
    """
    try:
        import time

        session = get_snowpark_session()
        query = load_query_from_file('sql/01_get_identity.sql')

        # Add query tag as SQL comment for observability
        query = f'-- QUERY_TAG: plv:get_identity\n{query}'

        # Replace parameter
        query = query.replace(':identity_id', f"'{identity_id}'")

        # Track query execution time
        start_time = time.time()
        df = session.sql(query).to_pandas()
        elapsed_ms = (time.time() - start_time) * 1000

        # Normalize column names to lowercase for consistency
        df.columns = df.columns.str.lower()

        if len(df) > 0:
            st.success(f'✓ Found identity: {df.iloc[0]["email"]} ({elapsed_ms:.0f}ms)')

            # Display identity details
            identity = df.iloc[0]

            # Get name field, or construct from first/last name
            name = identity.get('name')
            if not name or pd.isna(name):
                first = identity.get('first_name') or ''
                last = identity.get('last_name') or ''
                name = f'{first} {last}'.strip() or 'N/A'

            display_df = pd.DataFrame(
                {
                    'Field': ['Name', 'Email', 'Created At', 'Last Modified At'],
                    'Value': [
                        name,
                        identity.get('email') or 'N/A',
                        str(identity.get('created_at') or 'N/A'),
                        str(identity.get('last_modified_at') or 'N/A'),
                    ],
                }
            )

            st.dataframe(
                display_df,
                use_container_width=True,
                hide_index=True,
                column_config={
                    'Field': st.column_config.TextColumn('Field', width='medium'),
                    'Value': st.column_config.TextColumn('Value', width='large'),
                },
            )

        return df

    except Exception as e:
        st.error(f'Failed to load identity: {str(e)}')
        return None


def create_canonical_profile_list(profiles_df: pd.DataFrame) -> pd.DataFrame:
    """
    Create a canonical list of unique profiles from CDC data

    The SQL query may return multiple rows for the same profile UUID if there are
    multiple CDC events. This function deduplicates by profile_uuid and preserves
    any available type/id information. Missing type/id will be filled in Step 2.5
    via hydration from FACT.PROD.PROFILE.

    Args:
        profiles_df: Raw DataFrame from get_profiles query

    Returns:
        Canonical DataFrame with one row per unique profile (may have incomplete data)
    """
    if profiles_df.empty:
        return profiles_df

    # Group by profile_uuid and take first non-null values for type and id
    canonical = profiles_df.groupby('profile_uuid', as_index=False).agg(
        {
            'profile_type': lambda x: x.dropna().iloc[0] if not x.dropna().empty else None,
            'profile_id': lambda x: x.dropna().iloc[0] if not x.dropna().empty else None,
        }
    )

    # Only require profile_uuid (type/id will be hydrated in Step 2.5)
    canonical = canonical.dropna(subset=['profile_uuid'])

    return canonical


def get_profiles(identity_id: str) -> pd.DataFrame | None:
    """
    Step 2: Load all profiles for the identity from CDC_MUSICGRAPH_HASPROFILE
    Returns a canonical list of unique profiles (one per UUID)

    Args:
        identity_id: UUID of the identity

    Returns:
        DataFrame with canonical profile list (uuid, type, id)
    """
    try:
        session = get_snowpark_session()
        query = load_query_from_file('sql/v4/02a_get_profiles.sql')

        # Replace parameter
        query = query.replace(':identity_id', f"'{identity_id}'")

        df = session.sql(query).to_pandas()
        # Normalize column names to lowercase for consistency
        df.columns = df.columns.str.lower()

        # Create canonical list (deduplicate and ensure complete info)
        canonical_df = create_canonical_profile_list(df)

        st.info(f'✓ Found {len(canonical_df)} unique profile(s)')
        if len(df) != len(canonical_df):
            st.caption(
                f'Note: Processed {len(df)} distinct rows into '
                f'{len(canonical_df)} canonical profiles'
            )

        return canonical_df

    except Exception as e:
        st.error(f'Failed to load profiles: {str(e)}')
        return None


def get_deleted_profiles(identity_id: str) -> pd.DataFrame | None:
    """
    Step 2b: Load all deleted profiles for the identity from CDC_MUSICGRAPH_DELETEDHASPROFILE
    Returns a list of profiles that were deleted from this identity

    Args:
        identity_id: UUID of the identity

    Returns:
        DataFrame with deleted profile list (uuid, type, id)
    """
    try:
        session = get_snowpark_session()
        query = load_query_from_file('sql/v4/02b_get_deleted_profiles.sql')

        # Replace parameter
        query = query.replace(':identity_id', f"'{identity_id}'")

        df = session.sql(query).to_pandas()
        # Normalize column names to lowercase for consistency
        df.columns = df.columns.str.lower()

        # Create canonical list (deduplicate)
        canonical_df = create_canonical_profile_list(df)

        st.info(f'✓ Found {len(canonical_df)} deleted profile(s)')
        if len(df) != len(canonical_df):
            st.caption(
                f'Note: Processed {len(df)} distinct rows into '
                f'{len(canonical_df)} canonical deleted profiles'
            )

        return canonical_df

    except Exception as e:
        st.warning(f'Failed to load deleted profiles: {str(e)}')
        return None


def get_profiles_v5(identity_id: str) -> pd.DataFrame | None:
    """
    Step 2: Load all profiles for the identity from CDC V5 HAS_PROFILE
    Returns a canonical list of unique profiles (one per UUID)

    Args:
        identity_id: UUID of the identity

    Returns:
        DataFrame with canonical profile list (uuid, type, id)
    """
    try:
        session = get_snowpark_session()
        query = load_query_from_file('sql/v5/02a_get_profiles.sql')

        # Replace parameter
        query = query.replace(':identity_id', f"'{identity_id}'")

        df = session.sql(query).to_pandas()
        # Normalize column names to lowercase for consistency
        df.columns = df.columns.str.lower()

        # Create canonical list (deduplicate and ensure complete info)
        canonical_df = create_canonical_profile_list(df)

        st.info(f'✓ V5: Found {len(canonical_df)} unique profile(s)')
        if len(df) != len(canonical_df):
            st.caption(
                f'Note: Processed {len(df)} distinct rows into '
                f'{len(canonical_df)} canonical profiles'
            )

        return canonical_df

    except Exception as e:
        st.warning(f'Failed to load V5 profiles: {str(e)}')
        return None


def get_deleted_profiles_v5(identity_id: str) -> pd.DataFrame | None:
    """
    Step 2b: Load all deleted profiles for the identity from CDC V5 DELETED_HAS_PROFILE
    Returns a list of profiles that were deleted from this identity

    Args:
        identity_id: UUID of the identity

    Returns:
        DataFrame with deleted profile list (uuid, type, id)
    """
    try:
        session = get_snowpark_session()
        query = load_query_from_file('sql/v5/02b_get_deleted_profiles.sql')

        # Replace parameter
        query = query.replace(':identity_id', f"'{identity_id}'")

        df = session.sql(query).to_pandas()
        # Normalize column names to lowercase for consistency
        df.columns = df.columns.str.lower()

        # Create canonical list (deduplicate)
        canonical_df = create_canonical_profile_list(df)

        st.info(f'✓ V5: Found {len(canonical_df)} deleted profile(s)')
        if len(df) != len(canonical_df):
            st.caption(
                f'Note: Processed {len(df)} distinct rows into '
                f'{len(canonical_df)} canonical deleted profiles'
            )

        return canonical_df

    except Exception as e:
        st.warning(f'Failed to load V5 deleted profiles: {str(e)}')
        return None


def get_access_records_snowpark(profile_uuids: list[str], table_type: str, query_file: str):
    """
    Query access tables and return Snowpark DataFrame (for efficient union operations)

    Args:
        profile_uuids: List of profile UUIDs to query
        table_type: Type of table (for logging)
        query_file: Query file to use

    Returns:
        Tuple of (Snowpark DataFrame, row count, elapsed_ms)
    """
    try:
        import re
        import time

        session = get_snowpark_session()
        query = load_query_from_file(query_file)

        # Add query tag for observability
        query_tag = f'plv:query_{table_type.lower().replace("_", "")}'
        query = f'-- QUERY_TAG: {query_tag}\n{query}'

        # Validate UUIDs to prevent any injection (defense in depth)
        uuid_pattern = re.compile(
            r'^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$', re.I
        )
        validated_uuids = [uuid for uuid in profile_uuids if uuid_pattern.match(uuid)]

        if len(validated_uuids) != len(profile_uuids):
            st.warning(
                f'⚠️ Filtered out {len(profile_uuids) - len(validated_uuids)} invalid UUID(s)'
            )

        if not validated_uuids:
            st.warning(f'No valid UUIDs to query for {table_type}')
            return None, 0, 0

        # Create SQL IN clause with validated UUIDs
        uuids_str = "','".join(validated_uuids)
        query = query.replace(':profile_uuids', f"'{uuids_str}'")

        # Remove date filter placeholder (we always load all events)
        query = query.replace(':date_filter', '')

        # Track query execution time
        start_time = time.time()

        # Return Snowpark DataFrame (NOT pandas)
        snowpark_df = session.sql(query)

        # Get count for display (this executes the query)
        row_count = snowpark_df.count()
        elapsed_ms = (time.time() - start_time) * 1000

        # Cache result in Snowflake to avoid re-executing during union_all
        snowpark_df = snowpark_df.cache_result()

        st.info(f'✓ Found {row_count} record(s) from {table_type} ({elapsed_ms:.0f}ms)')
        return snowpark_df, row_count, elapsed_ms

    except Exception as e:
        st.warning(f'Query {table_type} returned no results or failed: {str(e)}')
        return None, 0, 0


def resolve_uuids_to_profiles(profile_uuids: list[str]) -> pd.DataFrame | None:
    """
    Resolve UUIDs to profileId/profileType from FACT.PROD.PROFILE
    Use this for hydrating profiles with complete metadata

    Args:
        profile_uuids: List of profile UUIDs to resolve

    Returns:
        DataFrame with profile resolution data (uuid, type, id)
    """
    try:
        import re

        session = get_snowpark_session()
        query = load_query_from_file('sql/02c_resolve_uuids_to_profiles.sql')

        # Validate UUIDs (defense in depth)
        uuid_pattern = re.compile(
            r'^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$', re.I
        )
        validated_uuids = [uuid for uuid in profile_uuids if uuid_pattern.match(uuid)]

        if not validated_uuids:
            st.warning('No valid UUIDs to resolve')
            return pd.DataFrame()

        # Create SQL array literal with validated UUIDs
        uuids_str = "','".join(validated_uuids)
        query = query.replace(':profile_uuids', f"'{uuids_str}'")

        df = session.sql(query).to_pandas()
        # Normalize column names to lowercase for consistency
        if not df.empty:
            df.columns = df.columns.str.lower()
        st.info(f'✓ Resolved {len(df)} profile(s) from FACT.PROD.PROFILE')
        return df

    except Exception as e:
        st.warning(f'Failed to resolve profiles: {str(e)}')
        return pd.DataFrame()


def get_profiles_from_has_profile(identity_id: str) -> pd.DataFrame | None:
    """
    Step 2c: Get profiles from FACTS.PROD.HAS_PROFILE + FACTS.PROD.PROFILE
    This finds profiles that may not be in the CDC table

    Args:
        identity_id: UUID of the identity

    Returns:
        DataFrame with profile data (uuid, type, id)
    """
    try:
        session = get_snowpark_session()
        query = load_query_from_file('sql/02c_get_profiles_from_has_profile.sql')

        # Replace parameter
        query = query.replace(':identity_id', f"'{identity_id}'")

        df = session.sql(query).to_pandas()
        # Normalize column names to lowercase for consistency
        if not df.empty:
            df.columns = df.columns.str.lower()
        st.info(f'✓ Found {len(df)} profile(s) from FACTS.PROD.HAS_PROFILE')
        return df

    except Exception as e:
        st.warning(f'Failed to load profiles from HAS_PROFILE: {str(e)}')
        return pd.DataFrame()


def hydrate_canonical_profiles(canonical_df: pd.DataFrame) -> tuple[pd.DataFrame, list[str]]:
    """
    Hydrate canonical profile list with complete data from FACT.PROD.PROFILE

    Some CDC events only have UUID without profileType/profileId. This function
    ensures all profiles have complete information (uuid, profileType, profileId)
    by querying FACT.PROD.PROFILE.

    Args:
        canonical_df: Canonical profile list from Step 2 (may have incomplete data)

    Returns:
        Tuple of (Hydrated DataFrame with all 3 identifiers, List of missing profile UUIDs)
    """
    if canonical_df.empty:
        return canonical_df, []

    # Get all profile UUIDs
    profile_uuids = canonical_df['profile_uuid'].tolist()

    # Resolve all UUIDs to get complete profile data from FACT.PROD.PROFILE
    resolved_df = resolve_uuids_to_profiles(profile_uuids)

    if resolved_df.empty:
        st.warning('⚠️ Could not resolve profiles from FACT.PROD.PROFILE')
        return canonical_df, profile_uuids

    # Drop profile_type and profile_id from canonical if they exist
    # We'll use the data from FACT.PROD.PROFILE as the source of truth
    # Column slicing creates a view, no need for .copy()
    canonical_clean = canonical_df[['profile_uuid']]

    # Merge with resolved data
    hydrated = canonical_clean.merge(
        resolved_df[['profile_uuid', 'profile_type', 'profile_id']], on='profile_uuid', how='left'
    )

    # Identify missing profiles before filtering
    missing_mask = hydrated['profile_type'].isna() | hydrated['profile_id'].isna()
    missing_profiles = hydrated.loc[missing_mask, 'profile_uuid'].tolist()

    # Keep only profiles with all 3 identifiers
    hydrated = hydrated.dropna(subset=['profile_uuid', 'profile_type', 'profile_id'])

    # Count how many profiles were successfully hydrated
    missing_count = len(canonical_df) - len(hydrated)
    if missing_count > 0:
        st.warning(
            f'⚠️ {missing_count} profile(s) could not be hydrated (not found in FACT.PROD.PROFILE)'
        )
        for uuid in missing_profiles:
            st.caption(f'  - Missing: {uuid}')

    st.success(f'✓ Hydrated {len(hydrated)} profile(s) with identifiers')

    return hydrated, missing_profiles


def enrich_with_profile_data(df: pd.DataFrame, profiles_df: pd.DataFrame) -> pd.DataFrame:
    """
    Enrich access records with profile metadata

    Args:
        df: Access records DataFrame
        profiles_df: Profiles DataFrame from get_profiles()

    Returns:
        Enriched DataFrame
    """
    if df.empty or profiles_df.empty:
        return df

    # Merge on profile_uuid to fill in missing profile_type and profile_id
    enriched = df.merge(
        profiles_df[['profile_uuid', 'profile_type', 'profile_id']],
        on='profile_uuid',
        how='left',
        suffixes=('', '_enriched'),
    )

    # Fill in missing values from enriched data
    if 'profile_type_enriched' in enriched.columns:
        enriched['profile_type'] = enriched['profile_type'].fillna(
            enriched['profile_type_enriched']
        )
        enriched.drop('profile_type_enriched', axis=1, inplace=True)

    if 'profile_id_enriched' in enriched.columns:
        enriched['profile_id'] = enriched['profile_id'].fillna(enriched['profile_id_enriched'])
        enriched.drop('profile_id_enriched', axis=1, inplace=True)

    return enriched


@st.cache_data(ttl=600, show_spinner=False)  # type: ignore[misc]
def hydrate_collaborators(tenant_ids: list[str]) -> pd.DataFrame:
    """
    Hydrate Collaborator tenant metadata
    Cached for 10 minutes

    Args:
        tenant_ids: List of tenant IDs to hydrate

    Returns:
        DataFrame with tenant_id, tenant_name columns
    """
    if not tenant_ids:
        return pd.DataFrame(columns=['tenant_id', 'tenant_name'])

    try:
        session = get_snowpark_session()
        query = load_query_from_file('sql/04a_hydrate_collaborators.sql')

        # Replace placeholder with SQL IN clause
        ids_str = "','".join(tenant_ids)
        query = query.replace(':tenant_ids', f"'{ids_str}'")

        df = session.sql(query).to_pandas()
        if not df.empty:
            df.columns = df.columns.str.lower()
            # Rename to standard column names
            df = df.rename(columns={'collaborator_id': 'tenant_id', 'name': 'tenant_name'})
            df['tenant_id'] = df['tenant_id'].astype(str)
        return df

    except Exception as e:
        st.warning(f'Failed to hydrate Collaborators: {str(e)}')
        return pd.DataFrame(columns=['tenant_id', 'tenant_name'])


@st.cache_data(ttl=600, show_spinner=False)  # type: ignore[misc]
def hydrate_vendors(tenant_ids: list[str], target_uuids: list[str]) -> pd.DataFrame:
    """
    Hydrate Vendor tenant metadata
    Cached for 10 minutes

    Args:
        tenant_ids: List of tenant IDs to hydrate
        target_uuids: List of tenant UUIDs to hydrate

    Returns:
        DataFrame with tenant_id, tenant_name columns
    """
    # Filter out wildcard '*' vendor IDs (they're not numeric)
    tenant_ids_filtered = [tid for tid in tenant_ids if tid != '*']

    if not tenant_ids_filtered and not target_uuids:
        return pd.DataFrame(columns=['tenant_id', 'tenant_name'])

    try:
        session = get_snowpark_session()
        query = load_query_from_file('sql/04b_hydrate_vendors.sql')

        # Replace placeholders
        ids_str = "','".join(tenant_ids_filtered) if tenant_ids_filtered else ''
        uuids_str = "','".join(target_uuids) if target_uuids else ''

        query = query.replace(':vendor_ids', f"'{ids_str}'")
        query = query.replace(':target_uuids', f"'{uuids_str}'")

        df = session.sql(query).to_pandas()
        if not df.empty:
            df.columns = df.columns.str.lower()
            # Rename to standard column names
            df = df.rename(columns={'vendor_id': 'tenant_id', 'vendor_name': 'tenant_name'})
            df['tenant_id'] = df['tenant_id'].astype(str)
        return df

    except Exception as e:
        st.warning(f'Failed to hydrate Vendors: {str(e)}')
        return pd.DataFrame(columns=['tenant_id', 'tenant_name'])


@st.cache_data(ttl=600, show_spinner=False)  # type: ignore[misc]
def hydrate_subaccounts(tenant_ids: list[str], target_uuids: list[str]) -> pd.DataFrame:
    """
    Hydrate Subaccount tenant metadata
    Cached for 10 minutes

    Args:
        tenant_ids: List of tenant IDs to hydrate
        target_uuids: List of tenant UUIDs to hydrate

    Returns:
        DataFrame with tenant_id, tenant_name columns
    """
    if not tenant_ids and not target_uuids:
        return pd.DataFrame(columns=['tenant_id', 'tenant_name'])

    try:
        session = get_snowpark_session()
        query = load_query_from_file('sql/04c_hydrate_subaccounts.sql')

        # Replace placeholders
        ids_str = "','".join(tenant_ids) if tenant_ids else ''
        uuids_str = "','".join(target_uuids) if target_uuids else ''

        query = query.replace(':tenant_ids', f"'{ids_str}'")
        query = query.replace(':target_uuids', f"'{uuids_str}'")

        df = session.sql(query).to_pandas()
        if not df.empty:
            df.columns = df.columns.str.lower()
            # Rename to standard column names
            df = df.rename(columns={'subaccount_id': 'tenant_id', 'subaccount_name': 'tenant_name'})
            df['tenant_id'] = df['tenant_id'].astype(str)
        return df

    except Exception as e:
        st.warning(f'Failed to hydrate Subaccounts: {str(e)}')
        return pd.DataFrame(columns=['tenant_id', 'tenant_name'])


@st.cache_data(ttl=600, show_spinner=False)  # type: ignore[misc]
def hydrate_label_participants(tenant_ids: list[str], target_uuids: list[str]) -> pd.DataFrame:
    """
    Hydrate LabelParticipant tenant metadata
    Cached for 10 minutes

    Args:
        tenant_ids: List of tenant IDs to hydrate
        target_uuids: List of tenant UUIDs to hydrate

    Returns:
        DataFrame with tenant_id, tenant_name columns
    """
    if not tenant_ids and not target_uuids:
        return pd.DataFrame(columns=['tenant_id', 'tenant_name'])

    try:
        session = get_snowpark_session()
        query = load_query_from_file('sql/04d_hydrate_label_participants.sql')

        # Replace placeholders
        ids_str = "','".join(tenant_ids) if tenant_ids else ''
        uuids_str = "','".join(target_uuids) if target_uuids else ''

        query = query.replace(':tenant_ids', f"'{ids_str}'")
        query = query.replace(':target_uuids', f"'{uuids_str}'")

        df = session.sql(query).to_pandas()
        if not df.empty:
            df.columns = df.columns.str.lower()
            # Rename to standard column names
            df = df.rename(columns={'label_participant_id': 'tenant_id', 'name': 'tenant_name'})
            df['tenant_id'] = df['tenant_id'].astype(str)
        return df

    except Exception as e:
        st.warning(f'Failed to hydrate LabelParticipants: {str(e)}')
        return pd.DataFrame(columns=['tenant_id', 'tenant_name'])


def enrich_events_with_tenants(events_df: pd.DataFrame, tenant_metadata: dict) -> pd.DataFrame:
    """
    Enrich events DataFrame with tenant names from metadata DataFrames

    Args:
        events_df: Events DataFrame with tenant_type and tenant_id columns
        tenant_metadata: Dict of {tenant_type: DataFrame with tenant_id, tenant_name columns}

    Returns:
        Enriched DataFrame with tenant_name column
    """
    if events_df.empty or not tenant_metadata:
        if 'tenant_name' not in events_df.columns:
            events_df['tenant_name'] = ''
        return events_df

    enriched = events_df.copy()

    # Initialize tenant_name column
    if 'tenant_name' not in enriched.columns:
        enriched['tenant_name'] = ''

    # Handle wildcard vendor IDs specially
    wildcard_mask = (enriched['tenant_type'] == 'Vendor') & (
        enriched['tenant_id'].astype(str) == '*'
    )
    enriched.loc[wildcard_mask, 'tenant_name'] = 'All Orchard Labels (*)'

    # Merge metadata for each tenant type
    for tenant_type, metadata_df in tenant_metadata.items():
        if metadata_df.empty:
            continue

        # Get rows for this tenant type
        type_mask = enriched['tenant_type'] == tenant_type

        # Create a lookup dict for faster access
        tenant_lookup = dict(zip(metadata_df['tenant_id'].astype(str), metadata_df['tenant_name']))

        # Apply lookup
        enriched.loc[type_mask, 'tenant_name'] = (
            enriched.loc[type_mask, 'tenant_id'].astype(str).map(tenant_lookup).fillna('')
        )

    return enriched


@st.cache_data(ttl=300, show_spinner=False)  # type: ignore[misc]
def query_audit_log(identity_id: str) -> pd.DataFrame:
    """
    Query audit log for all events (no date filtering at SQL level)
    Results are cached for 5 minutes to avoid redundant queries
    Date filtering happens in the UI after data is loaded

    This function breaks down the query into steps:
    1. Load identity from FACT.PROD.IDENTITY
    2a. Load profiles from CDC_MUSICGRAPH_HASPROFILE
    2b. Load deleted profiles from CDC_MUSICGRAPH_DELETEDHASPROFILE
    2c. Hydrate profiles with identifiers from FACT.PROD.PROFILE
    2d. Load additional profiles from FACTS.PROD.HAS_PROFILE (not in CDC)
    3. Query all access tables (unified UNION ALL query)
    4. Enrich all records with hydrated profile data and mark deleted profiles

    Args:
        identity_id: UUID of the identity to query

    Returns:
        Combined DataFrame with all audit log entries (includes is_deleted column)
    """
    try:
        st.subheader('⚙️ Gathering Logs...')

        # Initialize dict to store step results
        step_results = {}

        # Step 1: Get identity
        with st.expander('Step 1: Load Identity', expanded=True):
            identity_df = get_identity(identity_id)
            if identity_df is None or len(identity_df) == 0:
                st.error(f'Identity {identity_id} not found')
                return pd.DataFrame()
            step_results['identity_df'] = identity_df

        # Step 2a: Get profiles from CDC
        with st.expander('Step 2a: Load Profiles from CDC', expanded=True):
            cdc_profiles_df = get_profiles(identity_id)
            if cdc_profiles_df is None:
                st.error(f'Failed to load profiles for identity {identity_id}')
                return pd.DataFrame()
            if len(cdc_profiles_df) == 0:
                st.warning('No profiles found in CDC table')
                cdc_profiles_df = pd.DataFrame(
                    columns=['profile_uuid', 'profile_type', 'profile_id']
                )

        # Step 2b: Get deleted profiles from CDC
        with st.expander('Step 2b: Load Deleted Profiles from CDC', expanded=True):
            deleted_profiles_df = get_deleted_profiles(identity_id)
            if deleted_profiles_df is None:
                st.warning('Failed to load deleted profiles')
                deleted_profiles_df = pd.DataFrame(
                    columns=['profile_uuid', 'profile_type', 'profile_id']
                )
            if len(deleted_profiles_df) == 0:
                st.info('No deleted profiles found')
            # Store deleted profile UUIDs for later marking
            deleted_profile_uuids = set(deleted_profiles_df['profile_uuid'].tolist()) if not deleted_profiles_df.empty else set()
            step_results['deleted_profile_uuids'] = deleted_profile_uuids

        # Step 2c: Hydrate CDC profiles with data from FACT.PROD.PROFILE
        with st.expander('Step 2c: Hydrate CDC Profile Identifiers', expanded=True):
            if not cdc_profiles_df.empty:
                cdc_profiles_df, missing_profiles = hydrate_canonical_profiles(cdc_profiles_df)
            else:
                st.info('No CDC profiles to hydrate')
                missing_profiles = []

        # Step 2d: Get additional profiles from FACTS.PROD.HAS_PROFILE
        with st.expander('Step 2d: Load Profiles from HAS_PROFILE', expanded=True):
            has_profile_df = get_profiles_from_has_profile(identity_id)
            if has_profile_df is None:
                has_profile_df = pd.DataFrame(
                    columns=['profile_uuid', 'profile_type', 'profile_id']
                )

            # Find profiles that are in HAS_PROFILE but not in CDC
            if not has_profile_df.empty and not cdc_profiles_df.empty:
                # Get profile UUIDs already found in CDC
                cdc_uuids = set(cdc_profiles_df['profile_uuid'].tolist())
                # Filter HAS_PROFILE results to only new profiles
                new_profiles_df = has_profile_df[~has_profile_df['profile_uuid'].isin(cdc_uuids)]

                if not new_profiles_df.empty:
                    st.success(f'✓ Found {len(new_profiles_df)} additional profile(s) not in CDC')
                    # Combine CDC profiles with new profiles from HAS_PROFILE
                    profiles_df = pd.concat([cdc_profiles_df, new_profiles_df], ignore_index=True)
                else:
                    st.info('No additional profiles found (all profiles already in CDC)')
                    profiles_df = cdc_profiles_df
            elif has_profile_df.empty:
                st.info('No profiles found in HAS_PROFILE')
                profiles_df = cdc_profiles_df
            else:
                # CDC is empty, use all profiles from HAS_PROFILE
                st.success(
                    f'✓ Using {len(has_profile_df)} profile(s) from HAS_PROFILE (CDC was empty)'
                )
                profiles_df = has_profile_df

            # Final check
            if profiles_df.empty:
                st.error('No profiles found in either CDC or HAS_PROFILE')
                return pd.DataFrame()

            # Display combined profile list
            st.markdown('**Combined Profile List (All 3 Identifiers):**')

            # Create a formatted table for better display (column slicing creates a view)
            display_df = profiles_df[['profile_type', 'profile_id', 'profile_uuid']]
            display_df['profile_id'] = display_df['profile_id'].astype(int)

            st.dataframe(
                display_df,
                use_container_width=True,
                hide_index=True,
                column_config={
                    'profile_type': st.column_config.TextColumn('Type', width='small'),
                    'profile_id': st.column_config.NumberColumn('ID', width='small'),
                    'profile_uuid': st.column_config.TextColumn('UUID', width='large'),
                },
            )

            # Store profiles and missing profiles in step results
            step_results['profiles_df'] = profiles_df
            step_results['missing_profiles'] = missing_profiles

            # Extract UUIDs for Step 3 queries
            profile_uuids = profiles_df['profile_uuid'].tolist()

        # Step 3: Query access tables individually with 4 expanders
        st.info('🔍 Loading events from 4 CDC tables (filtering happens in UI)')

        # Step 3a: HAS_ACCESS_TO
        with st.expander('Step 3a: Query HAS_ACCESS_TO', expanded=False):
            df_3a, count_3a, elapsed_3a = get_access_records_snowpark(
                profile_uuids, 'HAS_ACCESS_TO', 'sql/v4/03a_get_has_access_to.sql'
            )

        # Step 3b: HAS_ADMIN_ACCESS_TO
        with st.expander('Step 3b: Query HAS_ADMIN_ACCESS_TO', expanded=False):
            df_3b, count_3b, elapsed_3b = get_access_records_snowpark(
                profile_uuids, 'HAS_ADMIN_ACCESS_TO', 'sql/v4/03b_get_has_admin_access_to.sql'
            )

        # Step 3c: DELETED_HAS_ACCESS_TO
        with st.expander('Step 3c: Query DELETED_HAS_ACCESS_TO', expanded=False):
            df_3c, count_3c, elapsed_3c = get_access_records_snowpark(
                profile_uuids, 'DELETED_HAS_ACCESS_TO', 'sql/v4/03c_get_deleted_has_access_to.sql'
            )

        # Step 3d: DELETED_HAS_ADMIN_ACCESS_TO
        with st.expander('Step 3d: Query DELETED_HAS_ADMIN_ACCESS_TO', expanded=False):
            df_3d, count_3d, elapsed_3d = get_access_records_snowpark(
                profile_uuids,
                'DELETED_HAS_ADMIN_ACCESS_TO',
                'sql/v4/03d_get_deleted_has_admin_access_to.sql',
            )

        # Combine using Snowpark union_all (computation stays in Snowflake)
        snowpark_dfs = [df for df in [df_3a, df_3b, df_3c, df_3d] if df is not None]

        if not snowpark_dfs:
            st.warning('No access records found in any table')
            st.session_state.query_step_results = step_results
            return pd.DataFrame()

        # Union all DataFrames in Snowpark (no pandas concat)
        combined_snowpark_df = snowpark_dfs[0]
        for df in snowpark_dfs[1:]:
            combined_snowpark_df = combined_snowpark_df.union_all(df)

        # Single conversion to pandas at the end
        combined_df = combined_snowpark_df.to_pandas()

        # Normalize column names to lowercase for consistency
        if not combined_df.empty:
            combined_df.columns = combined_df.columns.str.lower()

        # Display summary
        total_count = count_3a + count_3b + count_3c + count_3d
        total_elapsed = elapsed_3a + elapsed_3b + elapsed_3c + elapsed_3d
        st.success(f'✅ Found {total_count} total record(s) from 4 tables ({total_elapsed:.0f}ms)')

        # Store access counts by relationship type
        if not combined_df.empty:
            access_counts = combined_df.groupby('relationship_type').size().to_dict()
            step_results['access_counts'] = access_counts
        else:
            step_results['access_counts'] = {}

        # Check if we got any results
        if combined_df.empty:
            st.warning('No access records found')
            st.session_state.query_step_results = step_results
            return pd.DataFrame()

        # Enrich with hydrated profile data (single operation on combined DataFrame)
        combined_df = enrich_with_profile_data(combined_df, profiles_df)

        # Mark deleted profiles with is_deleted column
        deleted_profile_uuids = step_results.get('deleted_profile_uuids', set())
        if deleted_profile_uuids:
            combined_df['is_deleted'] = combined_df['profile_uuid'].isin(deleted_profile_uuids)
            deleted_count = combined_df['is_deleted'].sum()
            st.info(f'✓ Marked {deleted_count} event(s) with deleted profiles')
        else:
            combined_df['is_deleted'] = False

        # Sort by event_timestamp ascending (oldest to newest)
        combined_df = combined_df.sort_values('event_timestamp', ascending=True)

        st.success(f'✅ Total records found: {len(combined_df)}')

        # Store all step results in session state
        st.session_state.query_step_results = step_results

        return combined_df

    except Exception as e:
        st.error(f'Database query failed: {str(e)}')
        raise


@st.cache_data(ttl=300)  # type: ignore[misc]
def execute_query_cached(query: str) -> pd.DataFrame:
    """
    Execute a query with caching for performance

    Args:
        query: SQL query to execute

    Returns:
        DataFrame with query results
    """
    try:
        session = get_snowpark_session()
        df = session.sql(query).to_pandas()
        # Normalize column names to lowercase for consistency
        if not df.empty:
            df.columns = df.columns.str.lower()
        return df
    except Exception as e:
        st.error(f'Cached query failed: {str(e)}')
        return pd.DataFrame()


def hydrate_and_store_tenants(df: pd.DataFrame) -> pd.DataFrame:
    """
    Hydrate tenants using 4 separate queries (one per tenant type) with expanders
    Call this AFTER process_audit_data() to ensure tenant_type column exists

    Args:
        df: Processed DataFrame with tenant_type column

    Returns:
        DataFrame enriched with tenant_name column
    """
    if df.empty:
        return df

    st.info('🔍 Hydrating tenant metadata from 4 tables')

    # Initialize metadata dict
    tenant_metadata = {}

    # Extract unique tenants by type (IDs and UUIDs)
    for tenant_type in ['Collaborator', 'Vendor', 'Subaccount', 'LabelParticipant']:
        type_df = df[df['tenant_type'] == tenant_type]
        if type_df.empty:
            continue

        # Get unique IDs and UUIDs
        tenant_ids = type_df['tenant_id'].dropna().unique().tolist()
        target_uuids = type_df['target_uuid'].dropna().unique().tolist()

        # Filter out empty strings (but keep wildcards for Vendor)
        if tenant_type == 'Vendor':
            tenant_ids = [str(id) for id in tenant_ids if id and str(id) != '']
        else:
            tenant_ids = [str(id) for id in tenant_ids if id and str(id) != '' and str(id) != '*']
        target_uuids = [str(uuid) for uuid in target_uuids if uuid and str(uuid) != '']

        # Hydrate based on tenant type
        if tenant_type == 'Collaborator':
            with st.expander('Step 4a: Hydrate Collaborators', expanded=False):
                if tenant_ids:
                    metadata_df = hydrate_collaborators(tenant_ids)
                    tenant_metadata['Collaborator'] = metadata_df
                    st.info(f'✓ Hydrated {len(metadata_df)} Collaborator(s)')
                else:
                    tenant_metadata['Collaborator'] = pd.DataFrame(
                        columns=['tenant_id', 'tenant_name']
                    )
                    st.info('No Collaborators to hydrate')

        elif tenant_type == 'Vendor':
            with st.expander('Step 4b: Hydrate Vendors', expanded=False):
                if tenant_ids or target_uuids:
                    metadata_df = hydrate_vendors(tenant_ids, target_uuids)
                    tenant_metadata['Vendor'] = metadata_df
                    st.info(f'✓ Hydrated {len(metadata_df)} Vendor(s)')
                else:
                    tenant_metadata['Vendor'] = pd.DataFrame(columns=['tenant_id', 'tenant_name'])
                    st.info('No Vendors to hydrate')

        elif tenant_type == 'Subaccount':
            with st.expander('Step 4c: Hydrate Subaccounts', expanded=False):
                if tenant_ids or target_uuids:
                    metadata_df = hydrate_subaccounts(tenant_ids, target_uuids)
                    tenant_metadata['Subaccount'] = metadata_df
                    st.info(f'✓ Hydrated {len(metadata_df)} Subaccount(s)')
                else:
                    tenant_metadata['Subaccount'] = pd.DataFrame(
                        columns=['tenant_id', 'tenant_name']
                    )
                    st.info('No Subaccounts to hydrate')

        elif tenant_type == 'LabelParticipant':
            with st.expander('Step 4d: Hydrate LabelParticipants', expanded=False):
                if tenant_ids or target_uuids:
                    metadata_df = hydrate_label_participants(tenant_ids, target_uuids)
                    tenant_metadata['LabelParticipant'] = metadata_df
                    st.info(f'✓ Hydrated {len(metadata_df)} LabelParticipant(s)')
                else:
                    tenant_metadata['LabelParticipant'] = pd.DataFrame(
                        columns=['tenant_id', 'tenant_name']
                    )
                    st.info('No LabelParticipants to hydrate')

    # Enrich events with tenant names
    enriched_df = enrich_events_with_tenants(df, tenant_metadata)

    # Store metadata in session state for UI filters (already in DataFrame format)
    if 'query_step_results' not in st.session_state:
        st.session_state.query_step_results = {}

    st.session_state.query_step_results['tenant_metadata'] = tenant_metadata

    # Display summary
    total_hydrated = sum(len(v) for v in tenant_metadata.values() if not v.empty)
    st.success(f'✅ Hydrated {total_hydrated} total tenant(s) from 4 tables')

    return enriched_df


def test_connection() -> bool:
    """
    Test the database connection

    Returns:
        True if connection successful, False otherwise
    """
    try:
        session = get_snowpark_session()
        session.sql('SELECT CURRENT_VERSION()').collect()
        return True
    except Exception as e:
        st.error(f'Connection test failed: {str(e)}')
        return False