import streamlit as st
import pandas as pd
import time
from typing import List, Dict, Any, Optional
import requests
from dotenv import load_dotenv
import os

# Import the ApolloClient from the existing module
from apollo_contact_finder import ApolloClient

# Load environment variables from .env file
load_dotenv()

# Set page configuration
st.set_page_config(
    page_title="Apollo Contact Finder",
    page_icon="📧",
    layout="wide",
    initial_sidebar_state="expanded",
)

# Add custom CSS
st.markdown(
    """
<style>
    .main {
        padding: 2rem;
    }
    .stButton button {
        width: 100%;
    }
    .stProgress .st-bo {
        background-color: #4CAF50;
    }
    .contact-card {
        border: 1px solid #ddd;
        border-radius: 5px;
        padding: 10px;
        margin-bottom: 10px;
    }
</style>
""",
    unsafe_allow_html=True,
)

# Title and description
st.title("Apollo Contact Finder")
st.markdown("Search for company contacts and retrieve their email addresses.")

# Initialize session state variables if they don't exist
if "organizations" not in st.session_state:
    st.session_state.organizations = []
if "selected_org" not in st.session_state:
    st.session_state.selected_org = None
if "people" not in st.session_state:
    st.session_state.people = []
if "selected_people" not in st.session_state:
    st.session_state.selected_people = []
if "enriched_people" not in st.session_state:
    st.session_state.enriched_people = []
if "search_performed" not in st.session_state:
    st.session_state.search_performed = False
if "people_searched" not in st.session_state:
    st.session_state.people_searched = False
if "contacts_enriched" not in st.session_state:
    st.session_state.contacts_enriched = False


def search_organizations(query: str):
    """Search for organizations using the Apollo API"""
    try:
        client = ApolloClient()
        with st.spinner("Searching for organizations..."):
            org_results = client.search_organizations(query)

        st.session_state.organizations = org_results.get("organizations", [])
        st.session_state.search_performed = True

        if not st.session_state.organizations:
            st.error("No organizations found. Please try a different search.")
    except Exception as e:
        st.error(f"Error: {str(e)}")


def sort_by_title(persons: list[dict]):
    title_map = [
        {"rank": 1, "title": "founder"},
        {"rank": 1, "title": "ceo"},
        {"rank": 2, "title": "partner"},
        {"rank": 2, "title": "managing director"},
        {"rank": 3, "title": "principal"},
        {"rank": 3, "title": "director"},
        {"rank": 4, "title": "vice president"},
        {"rank": 4, "title": "vp"},
        {"rank": 5, "title": "associate"},
        {"rank": 6, "title": "analyst"},
    ]

    def get_rank(person: dict) -> int:
        person_title = person.get("title", "").lower()
        # Check each mapping; if the mapping title is found in the person's title, return its rank.
        for mapping in title_map:
            if mapping["title"] in person_title:
                return mapping["rank"]
        return 7  # Default rank if no match is found

    return sorted(persons, key=get_rank)


def search_people(org_id: str):
    """Search for people at a specific organization"""
    try:
        client = ApolloClient()
        with st.spinner("Searching for people..."):
            people_results = client.search_people(org_id)

        st.session_state.people = people_results.get("people", [])

        ## sort the people by seniority
        st.session_state.people = sort_by_title(st.session_state.people)
        st.session_state.people_searched = True

        if not st.session_state.people:
            st.error("No people found at this organization.")
    except Exception as e:
        st.error(f"Error: {str(e)}")


def enrich_people(person_ids: List[str]):
    """Get detailed information for selected people"""
    try:
        client = ApolloClient()
        enriched_people = []

        with st.spinner("Getting detailed contact information..."):
            # Process each person individually instead of using bulk_enrich
            for person_id in person_ids:
                try:
                    enriched_result = client.enrich_person(person_id)
                    if "person" in enriched_result:
                        enriched_people.append(
                            {"person": enriched_result.get("person", {})}
                        )
                    # Add a small delay to avoid rate limiting
                    # time.sleep(0.5)
                except Exception as e:
                    st.warning(f"Could not retrieve data for one contact: {str(e)}")
                    continue

            st.session_state.enriched_people = enriched_people

        st.session_state.contacts_enriched = True

        if not enriched_people:
            st.error("Could not retrieve contact information. Please try again.")
    except Exception as e:
        st.error(f"Error: {str(e)}")


def reset_state():
    """Reset the application state"""
    st.session_state.organizations = []
    st.session_state.selected_org = None
    st.session_state.people = []
    st.session_state.selected_people = []
    st.session_state.enriched_people = []
    st.session_state.search_performed = False
    st.session_state.people_searched = False
    st.session_state.contacts_enriched = False


# Sidebar for company search
with st.sidebar:
    st.header("Search")
    company_name = st.text_input("Enter company name")

    if st.button("Search Company"):
        if company_name:
            reset_state()
            search_organizations(company_name)
        else:
            st.warning("Please enter a company name")

    if st.session_state.search_performed and st.session_state.organizations:
        st.success(f"Found {len(st.session_state.organizations)} organizations")

    # Add a reset button
    if st.button("Reset"):
        reset_state()

# Main content area
if st.session_state.search_performed:
    # Display organizations
    if st.session_state.organizations and not st.session_state.selected_org:
        st.header("Select an Organization")

        # Create a DataFrame for better display
        org_data = []
        for org in st.session_state.organizations:
            org_data.append(
                {
                    "Name": org.get("name", "N/A"),
                    "Website": org.get("website_url", "N/A"),
                    "Industry": org.get("industry", "N/A"),
                    "Size": org.get("employee_count", "N/A"),
                    "ID": org.get("id", ""),
                }
            )

        df_orgs = pd.DataFrame(org_data)

        # Display the organizations as a table with selection
        selected_indices = st.dataframe(
            df_orgs[["Name", "Website", "Industry", "Size"]],
            use_container_width=True,
            column_config={
                "Name": st.column_config.TextColumn("Name", width="medium"),
                "Website": st.column_config.TextColumn("Website", width="medium"),
                "Industry": st.column_config.TextColumn("Industry", width="medium"),
                "Size": st.column_config.TextColumn("Size", width="small"),
            },
            height=400,
        )

        # Organization selection
        org_options = [org["name"] for org in st.session_state.organizations]
        selected_org_name = st.selectbox("Select an organization", org_options)

        if st.button("Confirm Selection"):
            # Find the selected organization
            for org in st.session_state.organizations:
                if org["name"] == selected_org_name:
                    st.session_state.selected_org = org
                    search_people(org["id"])
                    break

    # Display people if an organization is selected
    if st.session_state.selected_org and st.session_state.people_searched:
        st.header(f"People at {st.session_state.selected_org['name']}")

        if not st.session_state.people:
            st.warning("No people found at this organization.")
        else:
            # Create a DataFrame for better display
            people_data = []
            for person in st.session_state.people:
                people_data.append(
                    {
                        "Name": f"{person.get('first_name', '')} {person.get('last_name', '')}",
                        "Title": person.get("title", "N/A"),
                        "Seniority": person.get("seniority", "N/A"),
                        "Department": person.get("department", "N/A"),
                        "ID": person.get("id", ""),
                    }
                )

            df_people = pd.DataFrame(people_data)

            # Display the people as a table
            st.dataframe(
                df_people[["Name", "Title", "Seniority", "Department"]],
                use_container_width=True,
                column_config={
                    "Name": st.column_config.TextColumn("Name", width="medium"),
                    "Title": st.column_config.TextColumn("Title", width="medium"),
                    "Seniority": st.column_config.TextColumn(
                        "Seniority", width="small"
                    ),
                    "Department": st.column_config.TextColumn(
                        "Department", width="medium"
                    ),
                },
                height=400,
            )

            # Multi-select for people
            options = [
                f"{p.get('first_name', '')} {p.get('last_name', '')} ({p.get('title', 'N/A')})"
                for p in st.session_state.people
            ]
            selected_options = st.multiselect(
                "Select people to get contact information", options
            )

            # Map selected options back to people objects
            if selected_options:
                st.session_state.selected_people = []
                for option in selected_options:
                    for person in st.session_state.people:
                        person_name = f"{person.get('first_name', '')} {person.get('last_name', '')} ({person.get('title', 'N/A')})"
                        if person_name == option:
                            st.session_state.selected_people.append(person)

            # Button to enrich selected people
            if st.session_state.selected_people:
                if st.button("Get Contact Information"):
                    person_ids = [
                        person["id"] for person in st.session_state.selected_people
                    ]
                    enrich_people(person_ids)

    # Display enriched contact information
    if st.session_state.contacts_enriched and st.session_state.enriched_people:
        st.header("Contact Information")

        # Create a DataFrame for the contact information
        contact_data = []
        for person in st.session_state.enriched_people:
            person_data = person.get("person", {})
            contact_data.append(
                {
                    "Name": f"{person_data.get('first_name', '')} {person_data.get('last_name', '')}",
                    "Title": person_data.get("title", "N/A"),
                    "Email": person_data.get("email", "N/A"),
                    "Phone": person_data.get("phone_number", "N/A"),
                    "LinkedIn": person_data.get("linkedin_url", "N/A"),
                    "Company": person_data.get("organization", {}).get("name", "N/A"),
                    "Department": person_data.get("department", "N/A"),
                    "Seniority": person_data.get("seniority", "N/A"),
                }
            )

        df_contacts = pd.DataFrame(contact_data)

        # Display the contacts as a table
        st.dataframe(df_contacts, use_container_width=True, height=400)

        # Export to CSV option
        if st.button("Export to CSV"):
            csv = df_contacts.to_csv(index=False)
            filename = f"apollo_contacts_{int(time.time())}.csv"

            st.download_button(
                label="Download CSV File",
                data=csv,
                file_name=filename,
                mime="text/csv",
            )

            st.success(f"Data ready for download as {filename}")
else:
    # Initial state - show instructions
    st.info("Enter a company name in the sidebar to start searching for contacts.")

    # Add some example instructions
    st.markdown(
        """
    ### How to use this app:
    
    1. Enter a company name in the sidebar and click "Search Company"
    2. Select an organization from the results
    3. Select people from the organization to get their contact information
    4. Export the contact information to CSV if needed
    
    This app uses the Apollo.io API to find contact information for people at companies.
    """
    )

# Add footer
st.markdown("---")
st.markdown("Apollo Contact Finder - Powered by Apollo.io API")