import streamlit as st
import pandas as pd
from graphviz import Digraph
import os
import base64
from io import StringIO

# Configure system path for Hugging Face Spaces
os.environ["PATH"] += os.pathsep + '/usr/bin/graphviz'

# Enhanced 3D styling CSS
st.markdown("""
<style>
    .stApp {
        background: #f8f9fa;
        font-family: 'Segoe UI', system-ui;
    }
    .st-emotion-cache-1kyxreq {
        justify-content: center;
    }
    .stButton>button {
        background: linear-gradient(135deg, #4CAF50, #45a049);
        color: white;
        border-radius: 8px;
        padding: 12px 24px;
        border: none;
        font-weight: 600;
        transition: transform 0.3s;
        box-shadow: 0 4px 12px rgba(76,175,80,0.3);
    }
    .stDownloadButton>button {
        background: linear-gradient(135deg, #2196F3, #1976D2);
    }
    .st-eb {
        border-radius: 10px;
        box-shadow: 0 2px 8px rgba(0,0,0,0.1);
    }
    .node rect {
        filter: drop-shadow(2px 2px 2px rgba(0,0,0,0.2));
    }
</style>
""", unsafe_allow_html=True)

def create_org_chart(df, title):
    """Create professional org chart with 3D effects"""
    dot = Digraph(comment=title)
    dot.attr(rankdir='TB', labelloc='t', label=title, 
            fontsize='24', fontname='Arial', margin='0.5',
            splines='ortho', bgcolor='transparent')
    
    # 3D Node styling with gradient and shadow
    dot.attr('node', 
            shape='box',
            style='filled,rounded',
            fillcolor='linear-gradient(45deg, #45a049 0%, #4CAF50 100%)',
            gradientangle='270',
            fontname='Arial', 
            fontcolor='white',
            margin='0.4', 
            width='2', 
            height='0.9',
            fixedsize='false', 
            fontsize='14',
            penwidth='0',
            color='#2c3e50')
    
    # Edge styling for 3D effect
    dot.attr('edge', 
            color='#666666', 
            arrowsize='0.8', 
            penwidth='1.5', 
            dir='forward')

    # Process relationships
    added_nodes = set()
    for _, row in df.iterrows():
        parent = row['Parent'].strip()
        child = row['Child'].strip()
        if parent and child:
            if parent not in added_nodes:
                dot.node(parent, _attributes={'label': f'<<B>{parent}</B>>'})
                added_nodes.add(parent)
            if child not in added_nodes:
                dot.node(child, _attributes={'label': f'<<B>{child}</B>>'})
                added_nodes.add(child)
            dot.edge(parent, child)
    
    return dot

def main():
    st.title("🏢 3D Organization Chart Generator")
    
    with st.sidebar:
        st.header("⚙️ Configuration")
        chart_title = st.text_input("Chart Title", "Company Structure")
        st.markdown("---")
        st.markdown("### 📋 CSV Template")
        
        # CSV template download with 3D styling
        template = "Parent,Child\nCEO,CTO\nCEO,CFO\nCTO,Engineering Manager"
        b64 = base64.b64encode(template.encode()).decode()
        href = f'''
        <a href="data:file/csv;base64,{b64}" download="template.csv" 
           style="display: inline-block;
                  padding: 12px 24px;
                  border-radius: 8px;
                  background: linear-gradient(135deg, #4CAF50, #45a049);
                  color: white;
                  text-decoration: none;
                  font-weight: bold;
                  box-shadow: 0 4px 12px rgba(76,175,80,0.3);
                  transition: transform 0.3s;">
           📥 Download Template
        </a>
        '''
        st.markdown(href, unsafe_allow_html=True)

    uploaded_file = st.file_uploader("📤 Upload CSV File", type=["csv"])
    
    if uploaded_file:
        try:
            df = pd.read_csv(uploaded_file)
            if {'Parent', 'Child'}.issubset(df.columns):
                with st.spinner("🔍 Generating 3D visualization..."):
                    chart = create_org_chart(df, chart_title)
                    
                    # Two-column layout
                    col1, col2 = st.columns([3, 1])
                    
                    with col1:
                        st.graphviz_chart(chart, use_container_width=True)
                    
                    with col2:
                        st.markdown("### 📤 Export Options")
                        
                        # PDF Export
                        pdf_bytes = chart.pipe(format='pdf')
                        st.download_button(
                            label="📄 Download PDF",
                            data=pdf_bytes,
                            file_name=f"{chart_title.replace(' ', '_')}.pdf",
                            mime="application/pdf",
                        )
                        
                        # PNG Export
                        png_bytes = chart.pipe(format='png')
                        st.download_button(
                            label="🖼️ Download PNG",
                            data=png_bytes,
                            file_name=f"{chart_title.replace(' ', '_')}.png",
                            mime="image/png",
                        )
            else:
                st.error("❌ CSV must contain 'Parent' and 'Child' columns")
        except Exception as e:
            st.error(f"🚨 Error processing file: {str(e)}")

if __name__ == "__main__":
    main()