import streamlit as st
from langchain_groq import ChatGroq
from langchain_community.document_loaders import WebBaseLoader
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import JsonOutputParser
from PyPDF2 import PdfReader



@st.cache_resource(show_spinner=False)
def load_llm():
    llm = ChatGroq(
        temperature=0,
        groq_api_key='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',
        model_name="llama-3.3-70b-versatile"
    )
    parser = JsonOutputParser()
    return llm, parser

llm, json_parser = load_llm()

if "llm_loaded" not in st.session_state:
    st.toast("✅ LLM Loaded Successfully", icon="🤖")
    st.session_state.llm_loaded = True


def scrape_job_data(url):
    try:
        loader = WebBaseLoader(url)
        data = loader.load().pop().page_content
        return data.strip()
    except Exception as e:
        return f"Error fetching data: {e}"

def read_resume(uploaded_file):
    if uploaded_file.name.endswith(".pdf"):
        reader = PdfReader(uploaded_file)
        text = "\n".join([page.extract_text() for page in reader.pages if page.extract_text()])
    elif uploaded_file.name.endswith(".txt"):
        text = uploaded_file.read().decode("utf-8")
    else:
        text = ""
    return text.strip()


def get_job_desc_json(data):
    template = PromptTemplate.from_template(
        """
        ### SCRAPED DATA FROM WEBSITE:
        {data}

        ### INSTRUCTIONS:
        The scraped text is from the careers page of a website.
        Extract job postings and return valid JSON containing:
        'role', 'experience', 'skills', and 'description'.

        ### VALID JSON (NO PREAMBLE):
        """
    )
    chain = template | llm
    response = chain.invoke({"data": data})
    return json_parser.parse(response.content)

def get_resume_json(data):
    template = PromptTemplate.from_template(
        """
        ### SCRAPED DATA FROM RESUME:
        {resume_data}

        ### INSTRUCTIONS:
        Extract and return valid JSON containing:
        'name', 'about', 'role', 'experience', 'skills', 'projects', 'courses', and 'certificate'.

        ### VALID JSON (NO PREAMBLE):
        """
    )
    chain = template | llm
    response = chain.invoke({"resume_data": data})
    return json_parser.parse(response.content)

def get_cold_email(job, resume):
    template = PromptTemplate.from_template(
        """
        ### JOB DESCRIPTION:
        {job}

        ### RESUME:
        {resume}

        ### INSTRUCTION:
        You are the candidate described in the resume. 
        Write a concise, personalized cold email to the hiring manager expressing interest in the position.
        - Highlight relevant skills and experiences.
        - Maintain a professional yet approachable tone.
        - Avoid generic phrases; make it sound natural and tailored.
        - Keep it under 300-400 words.

        ### EMAIL (NO PREAMBLE):
        """
    )
    chain = template | llm
    response = chain.invoke({"job": job, "resume": resume})
    return response.content.strip()


# ---------------------------------------------------------
#  STREAMLIT APP UI
# ---------------------------------------------------------

st.markdown("""
# **⚡ NeoPitch** ✉️
#### "Turn job descriptions and your resume into perfect cold emails in seconds!"
""")

st.markdown("---")  


st.subheader("1️⃣ Job Description")
job_input_method = st.radio("How would you like to provide the job info?", ["Paste Text", "Enter URL"])

if "job_description" not in st.session_state:
    st.session_state.job_description = ""

if job_input_method == "Paste Text":
    st.session_state.job_description = st.text_area(
        "Paste Job Description Here 📄",
        value=st.session_state.job_description,
        height=150,
    )
else:
    url = st.text_input("Enter Job URL 🌐")
    if st.button("Fetch Job Description"):
        with st.spinner("Scraping job data... 🔍"):
            st.session_state.job_description = scrape_job_data(url)
    st.text_area(
        "Fetched Job Description 📝",
        st.session_state.job_description,
        height=150,
    )

st.markdown("---")  


st.subheader("2️⃣ Resume Input")
resume_input_method = st.radio("How would you like to provide your resume?", ["Upload File", "Paste Text"])

if "resume_data" not in st.session_state:
    st.session_state.resume_data = ""

if resume_input_method == "Upload File":
    uploaded_file = st.file_uploader("Upload Resume (PDF or TXT) 📎", type=["pdf", "txt"])
    if uploaded_file:
        st.session_state.resume_data = read_resume(uploaded_file)
else:
    st.session_state.resume_data = st.text_area(
        "Paste Resume Content Here 📝",
        value=st.session_state.resume_data,
        height=150,
    )

st.markdown("---")  


if st.button("🚀 Generate Email"):
    if not st.session_state.job_description.strip() or not st.session_state.resume_data.strip():
        st.warning("⚠️ Please provide both Job Description and Resume details.")
    else:
        with st.spinner("Analyzing and generating email... 🤖"):
            job = get_job_desc_json(st.session_state.job_description)
            resume = get_resume_json(st.session_state.resume_data)
            email = get_cold_email(job, resume)

        st.success("✅ Email Generated Successfully!")
        st.code(email, language="text")
        st.download_button("📋 Copy Email", email, file_name="cold_email.txt")