import os
from langchain_text_splitters import CharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.schema import HumanMessage
from langchain.document_loaders import UnstructuredFileLoader
from langchain_community.vectorstores import Chroma
from langchain_groq import ChatGroq
import gradio as gr

# Initialize ChromaDB and Groq API
DB_DIR = "chroma_db"
COLLECTION_NAME = "resume_collection"  # Changed to reflect purpose
embedding_function = HuggingFaceEmbeddings()

GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
llm = ChatGroq(api_key=GROQ_API_KEY, model_name="llama-3.1-8b-instant")

# Store resume content globally
stored_resume = None

def load_and_split_document(file_path):
    """Loads a document and splits it into chunks."""
    loader = UnstructuredFileLoader(file_path)
    documents = loader.load()
    
    text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=50)  # Increased chunk size
    chunks = text_splitter.split_documents(documents)
    
    return chunks

def upload_resume(file):
    """Processes and stores the resume."""
    try:
        global stored_resume
        uploaded_file_path = file.name
        
        # Load and store resume
        chunks = load_and_split_document(uploaded_file_path)
        stored_resume = "\n".join([chunk.page_content for chunk in chunks])
        
        return "Resume successfully uploaded and stored!"
    except Exception as e:
        return f"Error processing resume: {str(e)}"

def generate_cold_email(job_description, hiring_manager_name, company_name):
    """Generates a cold email based on the job description and stored resume."""
    try:
        if not stored_resume:
            return "Please upload your resume first!"

        # Craft prompt for skill extraction and email generation
        prompt = f"""
        Task: Generate a personalized cold email for a job application.

        Resume:
        {stored_resume}

        Job Description:
        {job_description}

        First, analyze the job description and resume to identify matching skills and relevant experiences.
        Then, create a compelling cold email to {hiring_manager_name} at {company_name} that:
        1. Shows enthusiasm for the role and company
        2. Highlights 2-3 most relevant skills/experiences that match the job requirements
        3. Demonstrates knowledge of the company
        4. Includes a clear call to action
        5. Maintains a professional yet engaging tone
        6. Keeps the email concise (max 200 words)

        Format the email with proper salutation and signature.
        """

        messages = [HumanMessage(content=prompt)]
        response = llm.invoke(messages)
        
        return response.content
    except Exception as e:
        return f"Error generating email: {str(e)}"

# Define the Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 📧 AI Cold Email Generator")
    
    with gr.Tab("Upload Resume"):
        resume_input = gr.File(label="Upload Your Resume (PDF)")
        resume_upload_button = gr.Button("Upload Resume")
        resume_status = gr.Textbox(label="Upload Status", interactive=False)
    
    with gr.Tab("Generate Email"):
        job_description = gr.Textbox(
            label="Job Description",
            placeholder="Paste the job description here...",
            lines=5
        )
        with gr.Row():
            hiring_manager = gr.Textbox(
                label="Hiring Manager's Name",
                placeholder="e.g., John Smith"
            )
            company = gr.Textbox(
                label="Company Name",
                placeholder="e.g., Tech Corp"
            )
        
        generate_button = gr.Button("Generate Cold Email")
        email_output = gr.Textbox(
            label="Generated Email",
            interactive=False,
            lines=10
        )
    
    # Event handlers
    resume_upload_button.click(
        upload_resume,
        inputs=[resume_input],
        outputs=[resume_status]
    )
    
    generate_button.click(
        generate_cold_email,
        inputs=[job_description, hiring_manager, company],
        outputs=[email_output]
    )

# Launch the app
demo.launch()