import os
import gradio as gr
from groq import Groq
import pandas as pd
from datetime import datetime
from sendgrid import SendGridAPIClient
from sendgrid.helpers.mail import Mail
import json
from pathlib import Path
from typing import Dict, List, Optional
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_groq import ChatGroq
from langchain_core.messages import HumanMessage, SystemMessage

# Define the storage directory in the Hugging Face Space
STORAGE_DIR = os.environ.get('STORAGE_DIR', 'data')

def ensure_storage_dir():
    """Create storage directory if it doesn't exist"""
    os.makedirs(STORAGE_DIR, exist_ok=True)

def initialize_email_client():
    sendgrid_key = os.environ.get("SENDGRID_API_KEY")
    sender_email = os.environ.get("SENDER_EMAIL")
    
    if not sendgrid_key or not sender_email:
        print("Warning: SendGrid configuration missing. Email sending will be disabled.")
        return None
    
    try:
        return SendGridAPIClient(api_key=sendgrid_key)
    except Exception as e:
        print(f"Error initializing SendGrid client: {e}")
        return None

class UserProfile:
    # [Previous UserProfile code remains the same]
    # ... [UserProfile implementation remains unchanged]
    pass

class EmailTemplate:
    # [Previous EmailTemplate code remains the same]
    # ... [EmailTemplate implementation remains unchanged]
    pass

class EmailGenie:
    def __init__(self):
        groq_api_key = os.environ.get("GROQ_API_KEY")
        self.client = Groq(api_key=groq_api_key)
        self.sendgrid_client = initialize_email_client()
        self.sender_email = os.environ.get("SENDER_EMAIL")
        self.user_profile = UserProfile()
        self.email_template = EmailTemplate()
        self.current_profile = None
        
        # Initialize LangChain components with updated imports
        self.model = ChatGroq(
            model="llama3-8b-8192"
        )
        
        # Create structured prompts for information processing
        self.structure_info_prompt = ChatPromptTemplate.from_messages([
            ("system", "Structure and summarize the following information about a person or company:"),
            ("user", "{info}")
        ])
        
        # Create structured prompts for email generation
        self.email_generation_prompt = ChatPromptTemplate.from_messages([
            ("system", """Create a professional email using the following information:
            Sender: {sender_name} ({sender_industry})
            Sender Background: {sender_background}
            Template Type: {template}"""),
            ("user", """Recipient: {recipient}
            Subject: {subject}
            Recipient Information: {structured_info}""")
        ])
        
        # Create output parser
        self.parser = StrOutputParser()
        
        # Create the chains using the new LCEL syntax
        self.structure_info_chain = self.structure_info_prompt | self.model | self.parser
        self.email_generation_chain = self.email_generation_prompt | self.model | self.parser
    
    def set_current_profile(self, name: str):
        self.current_profile = self.user_profile.get_profile(name)
        return self.current_profile is not None
    
    def structure_info(self, info: str) -> str:
        return self.structure_info_chain.invoke({"info": info})

    def generate_email(self, template_name: str, context: Dict) -> str:
        if not self.current_profile:
            return "Please set up your profile first."
        
        template = self.email_template.get_template(template_name)
        structured_info = self.structure_info(context.get('recipient_info', ''))
        
        email_content = self.email_generation_chain.invoke({
            "template": template,
            "sender_name": self.current_profile['name'],
            "sender_industry": self.current_profile['industry'],
            "sender_background": self.current_profile['background'],
            "recipient": context['recipient'],
            "subject": context['email_subject'],
            "structured_info": structured_info
        })
        
        return email_content

    def preview_email(self, email_content: str) -> str:
        if not self.current_profile:
            return "Please set up your profile first."
        return f"=== Email Preview ===\nFrom: {self.current_profile['name']}\n\n{email_content}"
    
    def send_email(self, to_email: str, subject: str, content: str) -> tuple[bool, str]:
        if not self.current_profile:
            return False, "Please set up your profile first."
        
        if not self.sendgrid_client or not self.sender_email:
            return False, "Email sending is not configured"
            
        message = Mail(
            from_email=self.sender_email,
            to_emails=to_email,
            subject=subject,
            plain_text_content=content
        )
        try:
            self.sendgrid_client.send(message)
            return True, "Email sent successfully"
        except Exception as e:
            return False, f"Error sending email: {e}"

# [The rest of the code (create_interface and main) remains the same]
# ... [Implementation remains unchanged]

if __name__ == "__main__":
    demo = create_interface()
    demo.launch()