import streamlit as st
import pandas as pd
from db import get_leads, init_db
import json
from enrich import run_sync, refresh_ce_data
from generate_reply import generate_email_draft

st.set_page_config(page_title="TradingTwins Lead Engine", layout="wide")

st.title("🚀 Lead Engine: TradingTwins")

# Sidebar Actions
st.sidebar.header("Actions")

if st.sidebar.button("1. Ingest Emails (Mock)"):
    from ingest import ingest_mock_leads
    init_db()
    count = ingest_mock_leads()
    st.sidebar.success(f"Ingested {count} new leads.")
    st.rerun()

if st.sidebar.button("2. Ingest Real Emails (Graph API)"):
    try:
        from trading_twins_ingest import process_leads
        with st.spinner("Fetching emails from Microsoft Graph..."):
            count = process_leads()
        if count > 0:
            st.sidebar.success(f"Successfully ingested {count} new leads form inbox!")
        else:
            st.sidebar.info("No new leads found in inbox.")
        st.rerun()
    except Exception as e:
        st.sidebar.error(f"Ingest failed: {e}")

if st.sidebar.button("3. Sync to Company Explorer"):
    with st.spinner("Syncing with Company Explorer API..."):
        # Capture output for debugging
        try:
            # We redirect stdout to capture prints
            import io
            from contextlib import redirect_stdout
            f = io.StringIO()
            with redirect_stdout(f):
                run_sync()
            output = f.getvalue()
            
            st.success("Sync finished!")
            with st.expander("See Process Log", expanded=True):
                st.code(output)
                
        except Exception as e:
            st.error(f"Sync Failed: {e}")

# Main View
leads = get_leads()
df = pd.DataFrame(leads)

if not df.empty:
    col1, col2, col3 = st.columns(3)
    col1.metric("Total Leads", len(df))
    col2.metric("New / Unsynced", len(df[df['status'] == 'new']))
    col3.metric("Synced to CE", len(df[df['status'] == 'synced']))

    st.subheader("Lead Pipeline")
    
    for index, row in df.iterrows():
        with st.expander(f"{row['company_name']} ({row['status']})"):
            c1, c2 = st.columns(2)
            
            # --- Left Column: Lead Data ---
            c1.write(f"**Contact:** {row['contact_name']}")
            c1.write(f"**Email:** {row['email']}")
            
            # Metadata Display
            meta = {}
            if row.get('lead_metadata'):
                try:
                    meta = json.loads(row['lead_metadata'])
                except:
                    pass
            
            if meta:
                c1.write("---")
                c1.write(f"**Area:** {meta.get('area', '-')}")
                c1.write(f"**Purpose:** {meta.get('purpose', '-')}")
                c1.write(f"**Location:** {meta.get('zip', '')} {meta.get('city', '')}")
            
            with c1.expander("Show Original Email Body"):
                st.code(row['raw_body'], language="html")
            
            # --- Right Column: Enrichment & Response ---
            enrichment = json.loads(row['enrichment_data']) if row['enrichment_data'] else {}
            
            if enrichment:
                c2.write("--- Integration Status ---")
                ce_id = enrichment.get('ce_id')
                
                if ce_id:
                    c2.success(f"✅ Linked to Company Explorer (ID: {ce_id})")
                    
                    # CE Data Display
                    ce_data = enrichment.get('ce_data', {})
                    vertical = ce_data.get('vertical')
                    summary = ce_data.get('summary')
                    
                    if vertical:
                        c2.info(f"**Industry:** {vertical}")
                    else:
                        c2.warning("Industry Analysis pending...")
                        
                    if summary:
                        with c2.expander("Show Analysis Summary"):
                            st.write(summary)
                    
                    # Refresh Button
                    if c2.button("🔄 Refresh Analysis Data", key=f"refresh_{row['id']}"):
                        with st.spinner("Fetching latest data from Company Explorer..."):
                            new_data = refresh_ce_data(row['id'], ce_id)
                            st.toast("Data refreshed!")
                            st.rerun()
                            
                    # Generate Reply Button (only if we have data)
                    c2.write("--- Response Generation ---")
                    if c2.button("✨ Generate Expert Reply", key=f"gen_{row['id']}"):
                        with st.spinner("Generating draft with Gemini..."):
                            # Prepare lead data dict from row
                            lead_dict = row.to_dict()
                            draft = generate_email_draft(lead_dict, ce_data)
                            # Store draft in session state to persist
                            st.session_state[f"draft_{row['id']}"] = draft
                    
                    # Display Draft
                    if f"draft_{row['id']}" in st.session_state:
                        c2.text_area("Draft Email", value=st.session_state[f"draft_{row['id']}"], height=300)
                        
                else:
                    c2.warning("⚠️ Not yet synced or failed")
                    c2.info(f"Log: {enrichment.get('message')}")

else:
    st.info("No leads found. Click 'Ingest Emails' in the sidebar.")