import streamlit as st
import yaml
import time
import os
from opcua_connector import opcua_connector

# Load config early for password if present
try:
    _APP_CFG = yaml.safe_load(open('./cfg.yaml')) or {}
except Exception:
    _APP_CFG = {}

# Page configuration
st.set_page_config(
    page_title="OPC UA Robot Control Dashboard",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS for better styling
st.markdown("""
<style>
    .main-header {
        text-align: center;
        color: #1f77b4;
        margin-bottom: 2rem;
    }
    .button-container {
        display: flex;
        justify-content: center;
        gap: 2rem;
        margin: 2rem 0;
        flex-wrap: wrap;
    }
    .control-button {
        display: flex;
        flex-direction: column;
        align-items: center;
        padding: 1.5rem;
        border-radius: 15px;
        background: linear-gradient(145deg, #f0f2f6, #ffffff);
        box-shadow: 5px 5px 15px #d1d9e6, -5px -5px 15px #ffffff;
        transition: all 0.3s ease;
        min-width: 150px;
        cursor: pointer;
    }
    .control-button:hover {
        transform: translateY(-2px);
        box-shadow: 8px 8px 20px #d1d9e6, -8px -8px 20px #ffffff;
    }
    .status-indicator {
        padding: 0.5rem 1rem;
        border-radius: 20px;
        font-weight: bold;
        text-align: center;
        margin: 1rem 0;
    }
    .status-connected {
        background-color: #d4edda;
        color: #155724;
        border: 1px solid #c3e6cb;
    }
    .status-disconnected {
        background-color: #f8d7da;
        color: #721c24;
        border: 1px solid #f5c6cb;
    }
    .log-container {
        background-color: #f8f9fa;
        border: 1px solid #dee2e6;
        border-radius: 5px;
        padding: 1rem;
        max-height: 300px;
        overflow-y: auto;
    }
    .stButton > button {
        height: 60px !important;
        font-size: 14px !important;
        font-weight: bold !important;
        border-radius: 10px !important;
        border: 2px solid #e0e0e0 !important;
        transition: all 0.3s ease !important;
    }
    .stButton > button:hover {
        transform: translateY(-2px) !important;
        box-shadow: 0 4px 8px rgba(0,0,0,0.2) !important;
        border-color: #1f77b4 !important;
    }
</style>
""", unsafe_allow_html=True)

# Initialize session state
if 'opcua_connector' not in st.session_state:
    st.session_state.opcua_connector = None
if 'connection_status' not in st.session_state:
    st.session_state.connection_status = "Disconnected"
if 'authenticated' not in st.session_state:
    st.session_state.authenticated = False

# Re-validate connection status each render (do this before rendering UI)
def is_opcua_connected(OPCcon) -> bool:
    try:
        if OPCcon is None:
            return False
        client = OPCcon.opcuaclient
        return hasattr(client, 'uaclient') and client.uaclient and getattr(client.uaclient, 'session', None) is not None
    except Exception:
        return False

if st.session_state.opcua_connector and is_opcua_connected(st.session_state.opcua_connector):
    if st.session_state.connection_status != "Connected":
        st.session_state.connection_status = "Connected"
else:
    if st.session_state.connection_status != "Disconnected":
        st.session_state.connection_status = "Disconnected"

def add_log(message):
    """No-op logger (logs disabled for minimal UI)."""
    return None

def check_opcua_server(OPCcon):
    """Check if OPC UA server is reachable"""
    try:
        # Check if already connected
        if hasattr(OPCcon.opcuaclient, 'uaclient') and OPCcon.opcuaclient.uaclient:
            add_log("✅ OPC UA server already connected")
            return True
        
        # Try to connect
        OPCcon.opcuaclient.connect()
        add_log("✅ OPC UA server is reachable")
        return True
    except Exception as e:
        add_log(f"❌ Error connecting to OPC UA server: {e}")
        return False

def establish_connection(OPCcon):
    """Establish connection to OPC UA server"""
    try:
        if OPCcon.opcuaclient.uaclient:
            add_log("✅ OPC UA connection established")
            return True
        else:
            add_log("❌ Error establishing connection")
            return False
    except Exception as e:
        add_log(f"❌ Error establishing connection: {e}")
        return False

def push_button(OPCcon, btn_name):
    """Push a button on the OPC UA server"""
    try:
        if btn_name == 'INIT':
            OPCcon.adapt_access_rights()
            add_log(f"🔧 INIT button pressed - Access rights adapted")
        else:
            OPCcon.press_btn(btn_name)
            add_log(f"🔘 {btn_name} button pressed")
        return True
    except Exception as e:
        add_log(f"❌ Error pressing {btn_name} button: {e}")
        return False

def initialize_connection():
    """Initialize OPC UA connection"""
    try:
        config = yaml.safe_load(open('./cfg.yaml'))
        OPCcon = opcua_connector(config)
        
        # Check server
        if check_opcua_server(OPCcon):
            # Establish connection
            if establish_connection(OPCcon):
                st.session_state.opcua_connector = OPCcon
                st.session_state.connection_status = "Connected"
                add_log("🎉 Connection initialized successfully")
                return True
        
        st.session_state.connection_status = "Disconnected"
        return False
    except Exception as e:
        add_log(f"❌ Failed to initialize connection: {e}")
        st.session_state.connection_status = "Disconnected"
        return False

def execute_button_action(button_name):
    """Execute button action with proper error handling"""
    if st.session_state.opcua_connector is None:
        add_log("❌ No connection available. Please initialize first.")
        return
    
    # Execute directly instead of threading to avoid race conditions
    try:
        push_button(st.session_state.opcua_connector, button_name)
    except Exception as e:
        add_log(f"❌ Critical error executing {button_name}: {e}")


# Sidebar
with st.sidebar:
    st.write(f"Status: {st.session_state.connection_status}")
    
    # Initialize button
    if st.button("🔌 Initialize Connection", use_container_width=True):
        with st.spinner("Initializing connection..."):
            initialize_connection()
            st.rerun()
    
    # Disconnect button
    if st.button("🔌 Disconnect", use_container_width=True):
        if st.session_state.opcua_connector:
            try:
                # Close InfluxDB connection
                if hasattr(st.session_state.opcua_connector, 'influxclient'):
                    st.session_state.opcua_connector.influxclient.close()
                    add_log("📊 InfluxDB connection closed")
                
                # Disconnect OPC UA
                if hasattr(st.session_state.opcua_connector, 'opcuaclient'):
                    st.session_state.opcua_connector.opcuaclient.disconnect()
                    add_log("🔌 Disconnected from OPC UA server")
                
            except Exception as e:
                add_log(f"⚠️ Warning during disconnect: {e}")
        
        st.session_state.opcua_connector = None
        st.session_state.connection_status = "Disconnected"
        add_log("🔌 Connection reset")
        st.rerun()
    
    st.divider()


    
# Password gate
# Prefer cfg.yaml cred.dashboard.password, else env DASHBOARD_PASSWORD, else 'admin'
expected_password = (
    (_APP_CFG.get('cred', {}).get('dashboard', {}) or {}).get('password')
    or os.getenv('DASHBOARD_PASSWORD')
    or 'admin'
)
ui_locked = not st.session_state.authenticated
if ui_locked:
    pwd = st.text_input("Enter Dashboard Password", type="password")
    if st.button("Login", use_container_width=True):
        if pwd == expected_password:
            st.session_state.authenticated = True
            st.rerun()
        else:
            st.error("Wrong password")

if not ui_locked:
    # Button container
    st.markdown('<div class="button-container">', unsafe_allow_html=True)

    # Create columns for buttons
    btn_col1, btn_col2, btn_col3, btn_col4 = st.columns(4)

    buttons_disabled = st.session_state.connection_status != "Connected"

    with btn_col1:
        if st.button("🔧\n\nINIT", key="init_btn", use_container_width=True, disabled=buttons_disabled):
            execute_button_action('INIT')

    with btn_col2:
        if st.button("▶️\n\nSTART", key="start_btn", use_container_width=True, disabled=buttons_disabled):
            execute_button_action('BTN_START')

    with btn_col3:
        if st.button("⏹️\n\nSTOP", key="stop_btn", use_container_width=True, disabled=buttons_disabled):
            execute_button_action('BTN_STOP')

    with btn_col4:
        if st.button("🔄\n\nRESET", key="reset_btn", use_container_width=True, disabled=buttons_disabled):
            execute_button_action('BTN_RESET')

    st.markdown('</div>', unsafe_allow_html=True)

# Footer removed for minimal UI