import streamlit as st
import logging
import numpy as np
import os
import base64

# Configure Logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def get_base64_of_bin_file(bin_file):
    with open(bin_file, 'rb') as f:
        data = f.read()
    return base64.b64encode(data).decode()
    
# Page configuration
sidebar_state = "expanded"
st.set_page_config(
    page_title="spMetaTME-Atlas",
    page_icon=":material/hub:",
    layout="wide",
    initial_sidebar_state=sidebar_state,
)


from src.ui.components.header import render_header, load_css
from src.ui.components.footer import render_footer

from src.ui.pages.home import show_overview
from src.ui.pages.analyze import (
    page_domain_statistics,
    page_spatial_flux,
    page_umap_analysis,
    page_differential_analysis,
    page_metabolic_interactions,
    page_metabolite_balance,
    page_reset
)

def init_session_state():
    """Initialise global session state."""
    if "adata" not in st.session_state:
        st.session_state.adata = None
    if "metabolic_adata" not in st.session_state:
        st.session_state.metabolic_adata = None
    if "data_type" not in st.session_state:
        st.session_state.data_type = None
    if "preprocessing_done" not in st.session_state:
        st.session_state.preprocessing_done = False
    if "flux_analysis_done" not in st.session_state:
        st.session_state.flux_analysis_done = False
    if "interaction_scores" not in st.session_state:
        st.session_state.interaction_scores = None
    if "interaction_type" not in st.session_state:
        st.session_state.interaction_type = None
    
    # Pagination States
    if "dataset_page" not in st.session_state:
        st.session_state.dataset_page = 1
    if "umap_page" not in st.session_state:
        st.session_state.umap_page = 1
    if "spatial_flux_page" not in st.session_state:
        st.session_state.spatial_flux_page = 1
    
    # Developer Mode
    if "dev_mode" not in st.session_state:
        st.session_state.dev_mode = True

def render_branding_header():
    """Display a clean, professional branding header using the network background."""
    try:
        bg_img_path = "assets/network_bg_red.png"
        bg_base64 = get_base64_of_bin_file(bg_img_path) if os.path.exists(bg_img_path) else ""
        subtitle = "A spatial atlas of tumour microenvironment metabolism and metabolic interactions inferred by a pretrained self-supervised metabolic hypergraph"

        # Load brand logo from npy
        brand_text = "spMetaTME-Atlas"
        brand_path = "assets/spMetaTME_brand.npy"
        if os.path.exists(brand_path):
            try:
                brand_data = np.load(brand_path, allow_pickle=True)
                brand_text = brand_data.item() if hasattr(brand_data, 'item') else str(brand_data)
            except Exception:
                pass
        html_code = f"""
<div style="position: relative; width: 100%; background: transparent; padding: 0.5rem 0; margin-bottom: 1rem; overflow: hidden; border-bottom: 1px solid #f0f0f0; text-align: center; border-radius: 12px;">
    <div style="position: absolute; top: 0; left: 0; width: 100%; height: 100%; background-image: url('data:image/png;base64,{bg_base64}'); background-size: cover; background-position: center; opacity: 0.05; z-index: 1; pointer-events: none;"></div>
    <div style="position: relative; z-index: 10; display: flex; flex-direction: column; align-items: center;">
        <pre class="brand-logo-ascii" style="margin-bottom: 5px; zoom: 0.85;">{brand_text}</pre>
        <p style="font-size: 0.85rem; color: #666; max-width: 900px; margin: 0 auto; font-weight: 400; line-height: 1.2;">{subtitle}</p>
    </div>
</div>
"""
        st.html(html_code)
    except Exception as e:
        try:
            st.markdown(html_code, unsafe_allow_html=True)
        except:
            logger.error(f"Failed to load branding header: {e}")

def render_sidebar_dev():
    """Developer shortcuts in sidebar."""
    with st.sidebar:
        st.markdown("---")
        st.session_state.dev_mode = st.checkbox("Developer Mode", value=st.session_state.dev_mode)
        
        if st.session_state.dev_mode:
            st.info("Dev Shortcuts Active")
            if st.button("Load Breast Cancer Block A", use_container_width=True):
                with st.spinner("Loading example data..."):
                    for key in ['interaction_scores', 'interaction_type']:
                        if key in st.session_state:
                            del st.session_state[key]
                    import scanpy as sc
                    adata = sc.read_h5ad(r"example_data/metabolic_Breast_cancer_Block_A.h5ad")
                    if adata is not None:
                        st.session_state.metabolic_adata = adata
                        st.session_state.data_type = "metabolic"
                        st.session_state.just_loaded = True
                        if 'domain' not in adata.obs.columns and 'domain_id' in adata.obs.columns:
                            adata.obs['domain'] = adata.obs['domain_id']
                        st.success("Loaded Breast Cancer Block A")
                        st.rerun()

def main():

    init_session_state()
    load_css()

    home_page = st.Page(show_overview, title="Home", icon=":material/home:", url_path="home")
    domain_stats_page = st.Page(page_domain_statistics, title="Domain Statistics", icon=":material/pie_chart:", url_path="domain_stats")
    
    analysis_pages = [
        st.Page(page_reset, title="Back to Home", icon=":material/home:", url_path="reset_to_home"),
        domain_stats_page,
        st.Page(page_spatial_flux, title="Spatial Flux Distribution", icon=":material/image:", url_path="spatial_flux"),
        st.Page(page_umap_analysis, title="UMAP Analysis", icon=":material/palette:", url_path="umap_analysis"),
        st.Page(page_differential_analysis, title="Differential Reactions", icon=":material/bar_chart:", url_path="differential_analysis"),
        st.Page(page_metabolic_interactions, title="Spatial Metabolic Interactions", icon=":material/link:", url_path="metabolic_interactions"),
        st.Page(page_metabolite_balance, title="Metabolite Balance Analysis", icon=":material/opacity:", url_path="metabolite_balance"),
    ]
    
    if st.session_state.metabolic_adata is None:
        pg = st.navigation([home_page], position="hidden")
        render_header()
        # render_sidebar_dev()
        pg.run()
    else:
        pg = st.navigation({"Metabolic Analysis": analysis_pages}, position="sidebar")
        render_branding_header()
        
        # Developer Shortcuts
        # render_sidebar_dev()
        
        if st.session_state.get('just_loaded', False):
            st.session_state.just_loaded = False
            st.switch_page(domain_stats_page)
        
        pg.run()

    render_footer()

if __name__ == "__main__":
    main()