import streamlit as st
import pandas as pd
import os
import json
from google.cloud import bigquery
from google.oauth2 import service_account
from html_templates import download_button_styles
import re
import time
import io #added

# Set page configuration
st.set_page_config(
    page_title="Querypad",  # Changed from "BigQuery Terminal"
    page_icon="💻",
    layout="wide",
    initial_sidebar_state="collapsed"
)

st.markdown(
        """
        <style>
        div[data-testid="stSidebarHeader"] > img, div[data-testid="collapsedControl"] > img {
            height: 6rem; /* Increased height */
            width: 15rem; /* Adjust width proportionally */
        }
        
        div[data-testid="stSidebarHeader"], div[data-testid="stSidebarHeader"] > *,
        div[data-testid="collapsedControl"], div[data-testid="collapsedControl"] > * {
            display: flex;
            align-items: center;
        }
        </style>
        """,
        unsafe_allow_html=True
    )
st.logo("alerter_4.jpeg")


# Custom CSS for terminal-like UI with typewriter effects
st.markdown("""
<style>
    /* Main container */
    .main {
        background-color: #1e1e1e;
        color: #f0f0f0;
        font-family: 'Courier New', monospace;
        padding: 2rem;
    }
    
    /* Title styling */
    h1, h2, h3 {
        color: #f0f0f0 !important;
        font-family: 'Courier New', monospace !important;
    }
    
    /* Query input box with thicker cursor */
    .stTextArea textarea {
        background-color: #252525 !important;
        color: #f0f0f0 !important;
        border: 1px solid #444 !important;
        font-family: 'Courier New', monospace !important;
        padding: 10px !important;
        height: 500px !important;  /* Increase size */
        font-size: 18px !important; /* Bigger text */
        caret-color: white !important; /* White cursor */
        caret-shape: block !important; /* Try to use block cursor if supported */
        text-shadow: 0 0 0 #f0f0f0 !important; /* Text glow to match cursor */
    }
    
    
    .stTextArea textarea::after {
    content: "▋";
    position: relative;
    display: inline-block;
    color: white;
    font-weight: bold; /* Make it bolder */
    font-size: 30px; /* Slightly larger */
    animation: blink 1s step-end infinite;
    }
    
    /* Use the same blink animation you already have */
    @keyframes blink {
        0%, 100% { opacity: 1; }
        50% { opacity: 0; }
    }
    


    /* Button styling - simplified and more reliable */
    .stButton > button {
        background: linear-gradient(to right, #800000, #ff0000);
        color: white;
        padding: 0.4rem 0.8rem;
        border-radius: 4px;
        border: none;
        cursor: pointer;
        text-align: center;
        font-weight: bold;
        font-family: 'Courier New', monospace;
        display: inline-block;
        margin: 10px 0;
        width: auto;
        font-size: 0.9rem;
    }
    
    /* Success message with typewriter cursor effect */
    .success-msg {
        background-color: #0f2a0f;
        color: #4eff4e;
        padding: 0.75rem;
        border-radius: 4px;
        margin: 1rem 0;
        font-family: 'Courier New', monospace;
        border-left: 3px solid #4eff4e;
        position: relative;
        overflow: hidden;
    }
    
    .success-msg::after {
        content: "▋";
        position: relative;
        display: inline-block;
        color: #4eff4e;
        animation: blink 1s step-end infinite;
    }
    
    /* Error message with typewriter cursor effect */
    .error-msg {
        background-color: #2a0f0f;
        color: #ff4e4e;
        padding: 0.75rem;
        border-radius: 4px;
        margin: 1rem 0;
        font-family: 'Courier New', monospace;
        border-left: 3px solid #ff4e4e;
        position: relative;
        overflow: hidden;
    }
    
    .error-msg::after {
        content: "▋";
        position: relative;
        display: inline-block;
        color: #ff4e4e;
        animation: blink 1s step-end infinite;
    }
    
    /* Blinking cursor animation */
    @keyframes blink {
        0%, 100% { opacity: 1; }
        50% { opacity: 0; }
    }
    
    /* Dataframe styling */
    .dataframe-container {
        margin-top: 1.5rem;
        background-color: #252525;
        border-radius: 4px;
        padding: 0.5rem;
        border: 1px solid #444;
    }
    
    /* Page header */
    .header {
        display: flex;
        align-items: center;
        margin-bottom: 1.5rem;
        border-bottom: 1px solid #444;
        padding-bottom: 0.5rem;
    }
    
    /* Updated header styling */
    .header {
        display: flex;
        align-items: center;
        margin-bottom: 1.5rem;
        padding-bottom: 0.5rem;
    }
    
    /* Create a shiny white 3D button container for the header */
    .header-container {
        display: inline-flex;
        align-items: center;
        background: linear-gradient(to bottom, #ffffff, #e6e6e6);
        border-radius: 8px;
        padding: 0.5rem 1.5rem;
        box-shadow: 0 4px 6px rgba(0,0,0,0.1), 
                    0 1px 3px rgba(0,0,0,0.2), 
                    inset 0 1px 0 rgba(255,255,255,0.9);
        border: 1px solid #ccc;
    }
    
    .header-icon {
        margin-right: 0.5rem;
        font-size: 1.5rem;
    }
    
    /* Gradient text for the title */
    .gradient-title {
        background: linear-gradient(to right, #800000, #ff0000);
        -webkit-background-clip: text;
        background-clip: text;
        color: transparent;
        font-family: 'Courier New', monospace !important;
        font-size: 1.8rem;
        font-weight: bold;
        margin: 0;
    }
    
    /* Terminal prompt style */
    .prompt {
        color: #4eff4e;
        font-weight: bold;
        margin-right: 0.5rem;
    }
    
    /* Status indicator */
    .status-indicator {
        padding: 0.25rem 0.5rem;
        border-radius: 3px;
        font-size: 0.8rem;
        margin-left: 0.5rem;
    }
    
    /* Status indicator */
    .status-indicator {
        padding: 0.25rem 0.5rem;
        border-radius: 3px;
        font-size: 0.8rem;
        margin-left: 0.5rem;
        display: flex;
        align-items: center;
    }
    
    .status-connected {
        background-color: #0f2a0f;
        color: #4eff4e;
    }
    
    .status-disconnected {
        background-color: #2a0f0f;
        color: #ff4e4e;
    }
    
    /* Blinking dot */
    .status-dot {
        width: 8px;
        height: 8px;
        border-radius: 50%;
        margin-right: 6px;
        display: inline-block;
    }
    
    .green-dot {
        background-color: #4eff4e;
        animation: blink-dot 1.5s ease-in-out infinite;
    }
    
    .red-dot {
        background-color: #ff4e4e;
        animation: blink-dot 1.5s ease-in-out infinite;
    }
    
    /* Dot blinking animation */
    @keyframes blink-dot {
        0%, 100% { opacity: 1; }
        50% { opacity: 0.3; }
    }
""", unsafe_allow_html=True)


def load_credentials():
    """Load GCP credentials from environment variable or file."""
    try:
        # First try environment variable
        credentials_json = os.getenv("GCP_CREDENTIALS")
        
        # Check if credentials file exists as fallback
        credentials_file = os.getenv("GCP_CREDENTIALS_FILE")
        if not credentials_json and credentials_file and os.path.exists(credentials_file):
            with open(credentials_file, 'r') as f:
                credentials_json = f.read()
        
        if not credentials_json:
            return None, "GCP credentials not found. Set GCP_CREDENTIALS environment variable or GCP_CREDENTIALS_FILE."
        
        credentials_info = json.loads(credentials_json)
        credentials = service_account.Credentials.from_service_account_info(credentials_info)
        return credentials, None
    except json.JSONDecodeError:
        return None, "Invalid JSON in GCP credentials. Check file format."
    except Exception as e:
        return None, f"Authentication error: {str(e)}"

def init_bigquery_client():
    """Initialize the BigQuery client with GCP credentials."""
    credentials, error = load_credentials()
    if credentials:
        try:
            client = bigquery.Client(credentials=credentials)
            # Test connection with a simple query
            client.query("SELECT 1").result()
            return client, None
        except Exception as e:
            return None, f"Failed to connect to BigQuery: {str(e)}"
    return None, error

def parse_operation_details(query, result):
    """Parse query to extract operation type and affected records."""
    # Determine operation type
    if re.search(r'^\s*SELECT', query, re.IGNORECASE):
        count = len(result) if result is not None else 0
        return f"Query complete. {count} records retrieved."
    elif re.search(r'^\s*INSERT', query, re.IGNORECASE):
        return f"INSERT operation successful."
    elif re.search(r'^\s*UPDATE', query, re.IGNORECASE):
        return f"UPDATE operation successful."
    elif re.search(r'^\s*DELETE', query, re.IGNORECASE):
        return f"DELETE operation successful."
    else:
        return f"Query executed successfully."

def execute_query(client, query):
    """Execute a BigQuery query and return results."""
    if not query.strip():
        return None, None, "Query is empty. Please enter a valid SQL query."
        
    try:
        # Start timer
        start_time = time.time()
        
        # For SELECT queries
        if re.search(r'^\s*SELECT', query, re.IGNORECASE):
            query_job = client.query(query)
            result = query_job.result()
            df = result.to_dataframe()
            
            # Get execution time
            elapsed_time = time.time() - start_time
            
            # Create success message
            success_msg = f"{parse_operation_details(query, df)} [{elapsed_time:.2f}s]"
            
            return df, success_msg, None
        
        # For DML queries (INSERT, UPDATE, DELETE)
        else:
            query_job = client.query(query)
            result = query_job.result()
            
            # Get execution time
            elapsed_time = time.time() - start_time
            
            # Create success message with affected rows
            success_msg = f"{parse_operation_details(query, None)} {query_job.num_dml_affected_rows} rows affected. [{elapsed_time:.2f}s]"
            
            return pd.DataFrame(), success_msg, None
            
    except Exception as e:
        # Create concise, actionable error message
        if "Syntax error" in str(e):
            error_msg = "Syntax error. Check SQL syntax near: " + str(e).split("Syntax error")[-1].strip()
        elif "Not found: Table" in str(e):
            table_match = re.search(r'Table\s+([^\s:]+)', str(e))
            table = table_match.group(1) if table_match else "specified table"
            error_msg = f"Table '{table}' not found. Verify table name and permissions."
        elif "Permission denied" in str(e):
            error_msg = "Access denied. Check service account permissions."
        elif "exceeded" in str(e).lower():
            error_msg = "Query quota exceeded. Add LIMIT clause or refine query."
        elif "Cannot query over table" in str(e):
            error_msg = "Schema error. Table structure changed or field doesn't exist."
        else:
            # Simplify error message for better readability
            error_msg = str(e).split("\n")[0] if "\n" in str(e) else str(e)
            error_msg = error_msg[:100] + "..." if len(error_msg) > 100 else error_msg
        
        return None, None, error_msg

def main():
    # Terminal-like header
    st.markdown('<div class="header"><div class="header-container"><span class="header-icon">💻</span><span class="gradient-title">Querypad</span></div></div>', unsafe_allow_html=True)
    
    # Initialize session state
    if "client" not in st.session_state:
        st.session_state.client = None
        st.session_state.connection_error = None
    
    if "last_result" not in st.session_state:
        st.session_state.last_result = None
    if "last_success" not in st.session_state:
        st.session_state.last_success = None
    if "last_error" not in st.session_state:
        st.session_state.last_error = None
    
    # Initialize BigQuery client (if not already done)
    if st.session_state.client is None:
        client, error = init_bigquery_client()
        st.session_state.client = client
        st.session_state.connection_error = error
    
    # Display connection status
    if st.session_state.client:
        st.markdown('<span class="status-indicator status-connected"><span class="status-dot green-dot"></span>CONNECTED</span>', unsafe_allow_html=True)
    else:
        st.markdown('<span class="status-indicator status-disconnected"><span class="status-dot red-dot"></span>DISCONNECTED</span>', unsafe_allow_html=True)
        if st.session_state.connection_error:
            st.markdown(f'<div class="error-msg">[ERROR] {st.session_state.connection_error}</div>', unsafe_allow_html=True)
            st.write("Please set the GCP_CREDENTIALS environment variable and restart the application.")
            return  # Exit function if not connected
    
    # Query input with terminal-like prompt
    st.markdown('<span class="prompt">sql></span> Enter your query:', unsafe_allow_html=True)
    query = st.text_area("", height=500, 
                         placeholder="SELECT * FROM your_dataset.your_table LIMIT 10",
                         label_visibility="collapsed")
    
    # Execute button (simplified)
    execute_button = st.button("EXECUTE")
    
    # Execute query when button is clicked
    if execute_button:
        if not query.strip():
            st.session_state.last_error = "Query is empty. Please enter a valid SQL query."
            st.session_state.last_success = None
            st.session_state.last_result = None
        else:
            with st.spinner("Executing query..."):
                # Clear previous results
                st.session_state.last_error = None
                st.session_state.last_success = None
                st.session_state.last_result = None
                
                # Execute the query
                df, success_msg, error_msg = execute_query(st.session_state.client, query)
                
                # Store results
                st.session_state.last_result = df
                st.session_state.last_success = success_msg
                st.session_state.last_error = error_msg
                
                # Use st.rerun() instead of experimental_rerun()
                st.rerun()
    
    # Display error if any - with typewriter style
    if st.session_state.last_error:
        st.markdown(f'<div class="error-msg">[ERROR] {st.session_state.last_error}</div>', unsafe_allow_html=True)
    
    # Display success message and results - with typewriter style
    elif st.session_state.last_success:
        st.markdown(f'<div class="success-msg">[SUCCESS] {st.session_state.last_success}</div>', unsafe_allow_html=True)
        
        # Display dataframe if not empty
        if st.session_state.last_result is not None and not st.session_state.last_result.empty:
            st.markdown('<div class="dataframe-container">', unsafe_allow_html=True)
            # st.dataframe(st.session_state.last_result, use_container_width=True)
            st.markdown('</div>', unsafe_allow_html=True)

            
            #Added
            if st.session_state.last_result is not None and not st.session_state.last_result.empty:
                st.markdown('<div class="dataframe-container">', unsafe_allow_html=True)
                st.dataframe(st.session_state.last_result, use_container_width=True)
                st.markdown('</div>', unsafe_allow_html=True)
    
                # CSV download
                csv_buffer = io.StringIO()
                st.session_state.last_result.to_csv(csv_buffer, index=False)
                st.markdown(download_button_styles, unsafe_allow_html = True)
                st.download_button(
                    label="📥 Download CSV",
                    data=csv_buffer.getvalue(),
                    file_name="query_results.csv",
                    mime="text/csv"
                )
    
                # XLSX download
                xlsx_buffer = io.BytesIO()
                with pd.ExcelWriter(xlsx_buffer, engine='xlsxwriter') as writer:
                    # Convert any timezone-aware datetime columns to naive
                    df_to_export = st.session_state.last_result.copy()
                    for col in df_to_export.select_dtypes(include=["datetimetz"]).columns:
                        df_to_export[col] = df_to_export[col].dt.tz_localize(None)
                    
                    df_to_export.to_excel(writer, index=False, sheet_name='Results')

                xlsx_buffer.seek(0)
                st.markdown(download_button_styles, unsafe_allow_html = True)
                st.download_button(
                    label="📥 Download Excel",
                    data=xlsx_buffer,
                    file_name="query_results.xlsx",
                    mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
                )

if __name__ == "__main__":
    main()