import streamlit as st
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import ChatPromptTemplate
import os
import PyPDF2  # For parsing PDFs
from streamlit_extras.add_vertical_space import add_vertical_space
from streamlit_extras.colored_header import colored_header
from pymongo import MongoClient  # For MongoDB connection
import requests

# Streamlit page configuration
st.set_page_config(page_title="Email Generator", page_icon="📧", layout="wide")

# Database connection setup
MONGO_URI = os.getenv("MONGO_URI")  # MongoDB connection string

# Function to connect to MongoDB
def connect_to_mongo():
    """Connect to MongoDB database and return the collection for user details."""
    client = MongoClient(MONGO_URI)
    db = client["email_generator"]  # Use your actual database name
    collection = db["user_details"]  # Use your actual collection name
    return collection

# Function to save data in MongoDB
def save_data_mongo(data):
    """Save user data in MongoDB."""
    collection = connect_to_mongo()
    collection.insert_one(data)

# Function to retrieve user data from MongoDB
def retrieve_user_data():
    """Retrieve user data from MongoDB."""
    collection = connect_to_mongo()
    user_data = collection.find_one()  # Fetch one user data entry
    return user_data

# Sidebar for settings
with st.sidebar:
    st.title("⚙️ Settings")
    temperature = st.slider("Temperature", 0.0, 1.0, 0.5, 0.1)
    add_vertical_space(2)

# Personal Details Screen
if "user_data" not in st.session_state:
    st.title("👤 Personal Details")
    user_name = st.text_input("Name")
    user_email = st.text_input("Email")
    user_profession = st.text_input("Profession")
    user_company = st.text_input("Company")

    if st.button("Save Details"):
        user_data = {
            "name": user_name,
            "email": user_email,
            "profession": user_profession,
            "company": user_company,
        }

        # Save to MongoDB
        save_data_mongo(user_data)

        st.success("Details saved successfully!")
        st.session_state.user_data = user_data
        st.experimental_rerun()  # Rerun to go to the main app

# Main header
colored_header(label="📧 Personalized Email Generator", description="Generate tailored emails with ease", color_name="green-70")

gemini_api_key = os.getenv("GEMINI_API_KEY")
if gemini_api_key is None:
    st.error("❌ Gemini API key not found. Please set it as a secret in Hugging Face Spaces.")
else:
    llm = ChatGoogleGenerativeAI(
        model="gemini-1.5-flash",
        verbose=True,
        temperature=temperature,
        google_api_key=gemini_api_key
    )

prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are a helpful assistant skilled at drafting emails. Create a personalized email for the following context: {email_type} with details: {details} "
            "The mail should be directed to the recipient's name, ( Dear name, with salutations)."
        ),
        ("human", "Generate a personalized email."),
    ]
)

def parse_cv(cv_file):
    """Parse the uploaded CV file to extract relevant information."""
    content = ""
    if cv_file.type == "application/pdf":
        reader = PyPDF2.PdfReader(cv_file)
        for page in reader.pages:
            content += page.extract_text()
    else:
        content = cv_file.read().decode("utf-8")
    return content

def email_generator(email_type, details):
    """Generate a personalized email using the provided details and email type."""
    chain = prompt | llm
    return chain.invoke(
        {
            "email_type": email_type,
            "details": details
        }
    ).content

# Retrieve user data if not already done
if "user_data" not in st.session_state:
    st.session_state.user_data = retrieve_user_data()

# Email Type Selection
email_type = st.selectbox("📧 Type of Email", ["Job Application", "Networking Opportunity", "Business Opportunity"])

details = st.session_state.user_data.copy()  # Use stored user data
if email_type == "Job Application":
    st.subheader("📝 Job Application Details")
    details["company_name"] = st.text_input("🏢 Company Name")
    details["company_details"] = st.text_area("🏢 Company Details")
    cv_file = st.file_uploader("📄 Upload Your CV (Optional)", type=["pdf", "docx", "txt"])
    details["reason_for_application"] = st.text_area("✍️ Reason for Application")

    if cv_file:
        with st.spinner("Parsing CV..."):
            details["cv_content"] = parse_cv(cv_file)
            st.success("CV parsed successfully!")
    else:
        details["cv_content"] = ""

elif email_type == "Networking Opportunity":
    st.subheader("🔗 Networking Opportunity Details")
    details["reason_for_networking"] = st.text_area("✍️ Reason for Networking")

elif email_type == "Business Opportunity":
    st.subheader("💼 Business Opportunity Details")
    details["product"] = st.text_input("🛠 Product or Service")
    details["your_contribution"] = st.text_area("💡 Your Contribution or Value Proposition")
    details["receiving_company"] = st.text_input("🏢 Receiving Company Name")
    details["company_details"] = st.text_area("🏢 Details about the Receiving Company")

# Generate Email Button
if st.button("✉️ Generate Email", key="generate_button"):
    if any(value for value in details.values() if isinstance(value, str)):
        with st.spinner("Generating your personalized email..."):
            generated_email = email_generator(email_type, details)
        st.success("Email generated successfully!")
        st.markdown("### 📝 Generated Email")
        st.write(generated_email)
    else:
        st.error("Please fill in the required fields to generate the email.")

add_vertical_space(2)
st.markdown("---")
st.markdown("Made with Streamlit | Data provided by Gemini API")