import streamlit as st
import snowflake.connector
import pandas as pd
import os
import json

st.set_page_config(
    page_title="Fermentation Status Predictor",
    page_icon="🧪",
    layout="wide"
)

@st.cache_resource
def get_snowflake_connection():
    try:
        user = os.environ.get("SNOWFLAKE_USER")
        password = os.environ.get("SNOWFLAKE_PASSWORD")
        account = os.environ.get("SNOWFLAKE_ACCOUNT") 
        warehouse = os.environ.get("SNOWFLAKE_WAREHOUSE")
        
        st.write("🔍 Debug Info:")
        st.write(f"- User found: {'✅' if user else '❌'}")
        st.write(f"- Password found: {'✅' if password else '❌'}")
        st.write(f"- Account found: {'✅' if account else '❌'} (using: {account[:10]}...)")
        st.write(f"- Warehouse found: {'✅' if warehouse else '❌'} (using: {warehouse})")
        
        if not all([user, password, account, warehouse]):
            missing = []
            if not user: missing.append("SNOWFLAKE_USER")
            if not password: missing.append("SNOWFLAKE_PASSWORD")
            if not account: missing.append("SNOWFLAKE_ACCOUNT")
            if not warehouse: missing.append("SNOWFLAKE_WAREHOUSE")
            
            st.error(f"❌ Missing environment variables: {', '.join(missing)}")
            st.info("💡 Please add these as Repository Secrets in your Hugging Face Space settings.")
            return None
        
        st.write("🔑 Attempting Snowflake connection...")
        st.write(f"Trying to connect to: {account}.snowflakecomputing.com")
        
        conn = snowflake.connector.connect(
            user=user,
            password=password,
            account=account,
            warehouse=warehouse,
            database="sennos_interview_prep",
            schema="public",
            login_timeout=60,
            network_timeout=60
        )
        
        st.success("✅ Successfully connected to Snowflake!")
        return conn
        
    except Exception as e:
        error_msg = str(e)
        st.error(f"❌ Failed to connect to Snowflake: {error_msg}")
        
        st.write("**Troubleshooting steps:**")
        
        if "250001" in error_msg or "Could not connect" in error_msg:
            st.warning("🔴 **Network/Account Issue Detected**")
            st.write("Your account identifier format might be wrong. Try these:")
            st.code(f"""
# If your Snowflake URL is like:
https://abc12345.snowflakecomputing.com
👉 Use: abc12345

# If your Snowflake URL is like:
https://abc12345.us-east-1.snowflakecomputing.com  
👉 Use: abc12345.us-east-1

# If you have a new org-based account:
https://orgname-accountname.snowflakecomputing.com
👉 Use: orgname-accountname
            """)
            st.write("**Current account value (first 15 chars):**", account[:15] if len(account) > 15 else account)
            
        st.write("- Check your Snowflake login URL and extract the exact account identifier")
        st.write("- Make sure you can log into Snowflake from your browser")
        st.write("- Verify the warehouse name is spelled correctly")
        st.write("- Check if your Snowflake account allows external connections")
        
        return None

def make_prediction(batch_id, temp_mean, temp_std, ph_mean, ph_std, 
                   pressure_mean, pressure_std, conductivity_mean, conductivity_std):
    try:
        user = os.environ.get("SNOWFLAKE_USER")
        password = os.environ.get("SNOWFLAKE_PASSWORD")
        account = os.environ.get("SNOWFLAKE_ACCOUNT") 
        warehouse = os.environ.get("SNOWFLAKE_WAREHOUSE")
        
        if not all([user, password, account, warehouse]):
            return None, "Missing credentials"
        
        conn = snowflake.connector.connect(
            user=user,
            password=password,
            account=account,
            warehouse=warehouse,
            database="sennos_interview_prep",
            schema="public"
        )
        
        cur = conn.cursor()
        
        create_temp_sql = f"""
        CREATE OR REPLACE TEMPORARY TABLE temp_prediction_input AS
        SELECT 
            {temp_mean} as TEMP_MEAN,
            {temp_std} as TEMP_STD,
            {ph_mean} as PH_MEAN,
            {ph_std} as PH_STD,
            {pressure_mean} as PRESSURE_MEAN,
            {pressure_std} as PRESSURE_STD,
            {conductivity_mean} as CONDUCTIVITY_MEAN,
            {conductivity_std} as CONDUCTIVITY_STD
        """
        
        cur.execute(create_temp_sql)
        
        prediction_sql = """
        SELECT 
            sennos_interview_prep.public.fermentation_status_predictor!PREDICT(
                TEMP_MEAN, TEMP_STD, PH_MEAN, PH_STD,
                PRESSURE_MEAN, PRESSURE_STD, CONDUCTIVITY_MEAN, CONDUCTIVITY_STD
            ) as prediction_result
        FROM temp_prediction_input
        """
        
        cur.execute(prediction_sql)
        result = cur.fetchone()
        
        if result:
            prediction_json = result[0]
            
            if isinstance(prediction_json, str):
                prediction_data = json.loads(prediction_json)
            else:
                prediction_data = prediction_json
            
            prediction_int = int(prediction_data.get('output_feature_0', prediction_data))
            
            if prediction_int == 0:
                status_label = "Good"
            elif prediction_int == 1:
                status_label = "Infection"
            else:
                status_label = "Stalled"
            
            cur.close()
            conn.close()
            return prediction_int, status_label
        else:
            cur.close()
            conn.close()
            return None, "No result returned"
            
    except Exception as e:
        return None, f"Prediction error: {str(e)}"

st.title("🧪 Fermentation Status Predictor")
st.write("Predict fermentation status using ML model hosted in Snowflake")
st.write("*This is a public demo - no login required!*")

with st.expander("🔗 Connection Status"):
    if get_snowflake_connection():
        st.success("✅ Connected to Snowflake")
    else:
        st.error("❌ Failed to connect to Snowflake")

with st.form("prediction_form"):
    st.subheader("Enter Fermentation Parameters:")
    
    col1, col2 = st.columns(2)
    
    with col1:
        batch_id = st.number_input("Batch ID", 
                                 min_value=1000, 
                                 max_value=9999, 
                                 value=1001,
                                 step=1)
        
        temp_mean = st.number_input("Temperature Mean (°C)", 
                                  min_value=0.0, 
                                  max_value=50.0, 
                                  value=33.5,
                                  step=0.1)
        
        temp_std = st.number_input("Temperature Std Dev", 
                                 min_value=0.0, 
                                 max_value=5.0, 
                                 value=0.5,
                                 step=0.1)
        
        ph_mean = st.number_input("pH Mean", 
                                min_value=0.0, 
                                max_value=14.0, 
                                value=5.0,
                                step=0.1)
        
        ph_std = st.number_input("pH Std Dev", 
                               min_value=0.0, 
                               max_value=2.0, 
                               value=0.2,
                               step=0.01)
    
    with col2:
        pressure_mean = st.number_input("Pressure Mean", 
                                      min_value=0.0, 
                                      max_value=30.0, 
                                      value=15.0,
                                      step=0.1)
        
        pressure_std = st.number_input("Pressure Std Dev", 
                                     min_value=0.0, 
                                     max_value=5.0, 
                                     value=0.1,
                                     step=0.01)
        
        conductivity_mean = st.number_input("Conductivity Mean", 
                                          min_value=0, 
                                          max_value=3000, 
                                          value=1600,
                                          step=10)
        
        conductivity_std = st.number_input("Conductivity Std Dev", 
                                         min_value=0, 
                                         max_value=200, 
                                         value=50,
                                         step=1)
    
    submitted = st.form_submit_button("🔮 Make Prediction")
    
    if submitted:
        with st.spinner('🔄 Connecting to Snowflake and making prediction... Please wait...'):
            prediction_int, status_or_error = make_prediction(
                batch_id, temp_mean, temp_std, ph_mean, ph_std,
                pressure_mean, pressure_std, conductivity_mean, conductivity_std
            )
            
            if prediction_int is not None:
                st.success("✅ Prediction Complete!")
                
                col1, col2, col3, col4 = st.columns(4)
                with col1:
                    st.metric("🆔 Batch ID", batch_id)
                with col2:
                    st.metric("🌡️ Temp Mean", f"{temp_mean}°C")
                with col3:
                    st.metric("🧪 pH Mean", f"{ph_mean}")
                with col4:
                    st.metric("📊 Conductivity", conductivity_mean)
                
                st.subheader("📊 Prediction Result:")
                
                if status_or_error == "Good":
                    st.success(f"🟢 **Status: {status_or_error}** (Code: {prediction_int})")
                    st.info("✅ Fermentation is proceeding normally!")
                elif status_or_error == "Infection":
                    st.error(f"🔴 **Status: {status_or_error}** (Code: {prediction_int})")
                    st.warning("⚠️ Possible bacterial infection detected!")
                else:
                    st.warning(f"🟡 **Status: {status_or_error}** (Code: {prediction_int})")
                    st.info("⏸️ Fermentation appears to be stalled.")
                    
            else:
                st.error(f"❌ {status_or_error}")

with st.expander("ℹ️ About This Model"):
    st.write("""
    This fermentation status prediction model analyzes fermentation batch data to predict the current status:
    
    - **🟢 Good (0)**: Fermentation is proceeding normally
    - **🔴 Infection (1)**: Bacterial infection likely detected  
    - **🟡 Stalled (2)**: Fermentation has stalled or stopped
    
    **Input Parameters:**
    - **Temperature Stats**: Mean and standard deviation of batch temperature
    - **pH Stats**: Mean and standard deviation of pH levels
    - **Pressure Stats**: Mean and standard deviation of pressure readings
    - **Conductivity Stats**: Mean and standard deviation of conductivity measurements
    
    The model uses statistical features (mean + std dev) to capture both the central tendency 
    and variability of each measurement across the fermentation batch.
    
    **Model Location**: Hosted in Snowflake Cloud Data Platform
    """)

st.subheader("📋 Batch Predictions (CSV Upload)")
uploaded_file = st.file_uploader("Upload CSV with columns: BATCH_ID, TEMP_MEAN, TEMP_STD, PH_MEAN, PH_STD, PRESSURE_MEAN, PRESSURE_STD, CONDUCTIVITY_MEAN, CONDUCTIVITY_STD", type=['csv'])

if uploaded_file is not None:
    try:
        df = pd.read_csv(uploaded_file)
        st.write("Preview of uploaded data:")
        st.dataframe(df.head())
        
        if st.button("🚀 Run Batch Predictions"):
            st.info("Batch predictions feature coming soon! For now, use individual predictions above.")
            
    except Exception as e:
        st.error(f"Error reading file: {str(e)}")

st.markdown("---")
st.write("🚀 Powered by Snowflake ML | Hosted on Hugging Face Spaces")
st.write("*No authentication required - anyone can use this predictor!*")