"""Simplified Snowflake connection utilities using secrets.toml only."""

import streamlit as st
from snowflake.snowpark import Session

def get_snowflake_connection():
    """Get Snowflake connection parameters from secrets.toml."""
    try:
        secrets = st.secrets["snowflake"]
        return {
            'account': secrets["account"],
            'user': secrets["user"],
            'password': secrets["password"],
            'warehouse': secrets["warehouse"],
            'database': secrets["database"],
            'schema': secrets["schema"],
            'role': secrets["role"]
        }
    except Exception as e:
        print(f"❌ Failed to load Snowflake credentials from secrets.toml: {str(e)}")
        return None


@st.cache_resource
def init_snowflake_session():
    """Initialize Snowflake Snowpark session."""
    try:
        conn_params = get_snowflake_connection()
        if not conn_params:
            return None
        
        # Create Snowpark session
        session = Session.builder.configs(conn_params).create()
        
        # Test the connection
        test_query = "SELECT CURRENT_VERSION() as version"
        result = session.sql(test_query).collect()

        return session
        
    except Exception as e:
        print(f"❌ Failed to connect to Snowflake: {str(e)}")
        return None


def test_snowflake_connection():
    """Test Snowflake connection and display connection info."""
    session = init_snowflake_session()
    
    if session:
        try:
            # Get current context info
            context_info = session.sql("""
                SELECT 
                    CURRENT_ACCOUNT() as account,
                    CURRENT_USER() as user,
                    CURRENT_WAREHOUSE() as warehouse,
                    CURRENT_DATABASE() as database,
                    CURRENT_SCHEMA() as schema,
                    CURRENT_ROLE() as role
            """).collect()
            
            if context_info:
                info = context_info[0]
                st.success("🎉 Snowflake Connection Successful!")
                
                col1, col2 = st.columns(2)
                with col1:
                    st.info(f"**Account:** {info['ACCOUNT']}")
                    st.info(f"**User:** {info['USER']}")
                    st.info(f"**Role:** {info['ROLE']}")
                
                with col2:
                    st.info(f"**Warehouse:** {info['WAREHOUSE']}")
                    st.info(f"**Database:** {info['DATABASE']}")
                    st.info(f"**Schema:** {info['SCHEMA']}")
                
                return True
        except Exception as e:
            print(f"Connection test failed: {str(e)}")
            return False
    
    return False