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:
    def __init__(self):
        ensure_storage_dir()
        self.storage_path = os.path.join(STORAGE_DIR, 'profiles.json')
        self._ensure_storage_exists()
    
    def _ensure_storage_exists(self):
        if not os.path.exists(self.storage_path):   
            with open(self.storage_path, 'w') as f:
                json.dump({"profiles": []}, f)
    
    def save_profile(self, profile_data: Dict) -> bool:
        try:
            with open(self.storage_path, 'r') as f:
                data = json.load(f)
            
            profiles = data.get("profiles", [])
            profile_exists = False
            for i, profile in enumerate(profiles):
                if profile.get("name") == profile_data["name"]:
                    profiles[i] = profile_data
                    profile_exists = True
                    break
            
            if not profile_exists:
                profiles.append(profile_data)
            
            data["profiles"] = profiles
            
            with open(self.storage_path, 'w') as f:
                json.dump(data, f)
            return True
        except Exception as e:
            print(f"Error saving profile: {e}")
            return False
    
    def get_profile(self, name: str) -> Optional[Dict]:
        try:
            with open(self.storage_path, 'r') as f:
                data = json.load(f)
            
            for profile in data.get("profiles", []):
                if profile.get("name") == name:
                    return profile
            return None
        except Exception as e:
            print(f"Error getting profile: {e}")
            return None
            
class EmailTemplate:
    def __init__(self):
        ensure_storage_dir()
        self.templates_path = os.path.join(STORAGE_DIR, 'templates.json')
        self._ensure_templates_exist()
    
    def _ensure_templates_exist(self):
        if not os.path.exists(self.templates_path):
            default_templates = {
                "sales_pitch": {
                    "subject": "Innovative Solutions for {company}",
                    "body": "Dear {name},\n\nI hope this email finds you well..."
                },
                "job_application": {
                    "subject": "Experienced {role} Application",
                    "body": "Dear {hiring_manager},\n\nI am writing to express my interest..."
                }
            }
            with open(self.templates_path, 'w') as f:
                json.dump(default_templates, f)
    
    def get_template(self, template_name: str) -> Dict:
        with open(self.templates_path, 'r') as f:
            templates = json.load(f)
        return templates.get(template_name, {})


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]


def create_interface():
    emailgenie = EmailGenie()
    
    with gr.Blocks(title="EmailGenie") as demo:
        gr.Markdown("# EmailGenie: AI-Powered Email Outreach")
        
        with gr.Tab("Profile Setup"):
            name_input = gr.Textbox(label="Full Name")
            industry_input = gr.Textbox(label="Industry")
            background_input = gr.Textbox(label="Professional Background", lines=3)
            target_input = gr.Textbox(label="Target Audience", lines=2)
            save_profile_btn = gr.Button("Save Profile")
            profile_status = gr.Textbox(label="Status", interactive=False)
        
        with gr.Tab("Generate Email"):
            profile_name = gr.Textbox(label="Your Profile Name")
            load_profile_btn = gr.Button("Load Profile")
            profile_load_status = gr.Textbox(label="Profile Status", interactive=False)
            
            template_dropdown = gr.Dropdown(
                choices=["sales_pitch", "job_application"],
                label="Select Template"
            )
            subject_input = gr.Textbox(label="Email Subject")
            recipient_input = gr.Textbox(label="Recipient Name/Company")
            recipient_info = gr.Textbox(label="Recipient Information", lines=5)
            generate_btn = gr.Button("Generate Email")
            preview_output = gr.Textbox(label="Email Preview", lines=10)
        
        with gr.Tab("Send Email"):
            recipient_email = gr.Textbox(label="Recipient Email")
            send_btn = gr.Button("Send Email")
            status_output = gr.Textbox(label="Status", interactive=False)
        
        def save_profile(name, industry, background, target):
            profile_data = {
                "name": name,
                "industry": industry,
                "background": background,
                "target_audience": target
            }
            success = emailgenie.user_profile.save_profile(profile_data)
            if success:
                emailgenie.set_current_profile(name)
            return "Profile saved successfully!" if success else "Error saving profile"
        
        def load_profile(name):
            success = emailgenie.set_current_profile(name)
            return "Profile loaded successfully!" if success else "Profile not found"
        
        def generate_email_handler(template, subject, recipient, info):
            if not emailgenie.current_profile:
                return "Please load your profile first"
            
            context = {
                "email_subject": subject,
                "recipient": recipient,
                "recipient_info": info
            }
            email_content = emailgenie.generate_email(template, context)
            return emailgenie.preview_email(email_content)
        
        def send_email_handler(email, content):
            success, message = emailgenie.send_email(email, "Your Subject", content)
            return message
        
        # Connect the interface components with their handlers
        save_profile_btn.click(
            save_profile,
            inputs=[name_input, industry_input, background_input, target_input],
            outputs=profile_status
        )
        
        load_profile_btn.click(
            load_profile,
            inputs=[profile_name],
            outputs=profile_load_status
        )
        
        generate_btn.click(
            generate_email_handler,
            inputs=[template_dropdown, subject_input, recipient_input, recipient_info],
            outputs=preview_output
        )
        
        send_btn.click(
            send_email_handler,
            inputs=[recipient_email, preview_output],
            outputs=status_output
        )
    
    return demo

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