import os
from cerebras.cloud.sdk import Cerebras
from markitdown import MarkItDown
from weasyprint import HTML
from docx import Document
from pptx import Presentation
from PyPDF2 import PdfReader
import gradio as gr
from PIL import Image

api_key = os.environ.get("CEREBRAS_API_KEY")
md_converter = MarkItDown()

def extract_file_preview(file_path):
    """
    Extracts a preview of the file based on its format.
    """
    try:
        file_ext = os.path.splitext(file_path)[-1].lower()

        if file_ext in [".jpg", ".jpeg", ".png"]:
            return Image.open(file_path)

        elif file_ext == ".pdf":
            reader = PdfReader(file_path)
            return "\n".join([page.extract_text() for page in reader.pages[:2]])

        elif file_ext in [".docx"]:
            doc = Document(file_path)
            return "\n".join([para.text for para in doc.paragraphs[:20]])

        elif file_ext in [".pptx"]:
            ppt = Presentation(file_path)
            slides_text = []
            for slide in ppt.slides[:5]:
                slide_text = []
                for shape in slide.shapes:
                    if hasattr(shape, "text"):
                        slide_text.append(shape.text)
                slides_text.append("\n".join(slide_text))
            return "\n---\n".join(slides_text)

        else:
            return "File preview not supported for this format."

    except Exception as e:
        return f"Error extracting file preview: {str(e)}"

# Functions for resume optimization
def create_prompt(resume_string: str, jd_string: str) -> str:
    """
    Creates a detailed prompt for AI-powered resume optimization based on a job description.
    """
    return f"""
You are a professional resume optimization expert specializing in tailoring resumes to specific job descriptions. Your goal is to optimize my resume and provide actionable suggestions for improvement to align with the target role.

### Guidelines:
1. **Relevance**:  
   - Prioritize experiences, skills, and achievements **most relevant to the job description**.  
   - Remove or de-emphasize irrelevant details to ensure a **concise** and **targeted** resume.
   - Limit work experience section to 2-3 most relevant roles
   - Limit bullet points under each role to 2-3 most relevant impacts

2. **Action-Driven Results**:  
   - Use **strong action verbs** and **quantifiable results** (e.g., percentages, revenue, efficiency improvements) to highlight impact.  

3. **Keyword Optimization**:  
   - Integrate **keywords** and phrases from the job description naturally to optimize for ATS (Applicant Tracking Systems).  

4. **Additional Suggestions** *(If Gaps Exist)*:  
   - If the resume does not fully align with the job description, suggest:  
     1. **Additional technical or soft skills** that I could add to make my profile stronger.  
     2. **Certifications or courses** I could pursue to bridge the gap.  
     3. **Project ideas or experiences** that would better align with the role.  

5. **Formatting**:  
   - Output the tailored resume in **clean Markdown format**.  
   - Include an **"Additional Suggestions"** section at the end with actionable improvement recommendations.  

---

### Input:
- **My resume**:  
{resume_string}

- **The job description**:  
{jd_string}

---

### Output:  
1. **Tailored Resume**:  
   - A resume in **Markdown format** that emphasizes relevant experience, skills, and achievements.  
   - Incorporates job description **keywords** to optimize for ATS.  
   - Uses strong language and is no longer than **one page**.

2. **Additional Suggestions** *(if applicable)*:  
   - List **skills** that could strengthen alignment with the role.  
   - Recommend **certifications or courses** to pursue.  
   - Suggest **specific projects or experiences** to develop.
"""

def get_resume_response(prompt: str, api_key: str, model: str = "llama-3.3-70b", temperature: float = 0.7) -> str:
    """
    Sends a resume optimization prompt to Cerebras' API and returns the optimized resume response.
    """
    # Initialize the Cerebras client with the API key
    client = Cerebras(api_key=api_key)

    # Make API call using the Llama 3.3 70B model
    stream = client.chat.completions.create(
        messages=[
            {"role": "system", "content": "Expert resume writer"},
            {"role": "user", "content": prompt}
        ],
        model=model,
        stream=True,
        temperature=temperature,
        max_completion_tokens=1024,
        top_p=1
    )

    # Collect the response chunks from the stream
    response_string = ""
    for chunk in stream:
        response_string += chunk.choices[0].delta.content or ""

    return response_string

def process_resume(file, jd_string):
    try:
        file_path = file.name
        original_preview = extract_file_preview(file_path)

        result = md_converter.convert(file_path)
        resume_string = result.text_content

        prompt = f"Optimize resume based on job description: {jd_string}"
        optimized_resume = "Optimized resume placeholder"  # Simulate response for now.

        # Save the files for download
        original_file_path = file_path
        optimized_file_path = "resumes/optimized_resume.md"
        with open(optimized_file_path, "w", encoding="utf-8") as f:
            f.write(optimized_resume)

        return original_preview, resume_string, optimized_resume, original_file_path, optimized_file_path
    except Exception as e:
        return f"Error processing file: {str(e)}", "", "", "", ""

def export_as_pdf(resume_md):
    try:
        pdf_path = "resumes/optimized_resume.pdf"
        HTML(string=resume_md).write_pdf(pdf_path)
        return pdf_path
    except Exception as e:
        return f"Failed to export PDF: {str(e)}"

def export_as_word(resume_md):
    try:
        doc_path = "resumes/optimized_resume.docx"
        doc = Document()
        for line in resume_md.split("\n"):
            doc.add_paragraph(line)
        doc.save(doc_path)
        return doc_path
    except Exception as e:
        return f"Failed to export Word: {str(e)}"

# Gradio UI
with gr.Blocks() as app:
    gr.Markdown("# Resume Optimizer 📄")
    gr.Markdown("Upload your resume, paste the job description, and get actionable insights!")

    with gr.Row():
        resume_input = gr.File(label="Upload Your Resume")
        jd_input = gr.Textbox(label="Paste Job Description", lines=5)

    run_button = gr.Button("Optimize Resume")
    
    with gr.Row():
        before_preview = gr.Markdown(label="Original File Preview")
        before_md = gr.Markdown(label="Original (Markdown)")
        after_md = gr.Markdown(label="Optimized (Markdown)")

    with gr.Row():
        download_before = gr.File(label="Download Original")
        #download_after_md = gr.File(label="Download Optimized (Markdown)")
        download_after_pdf = gr.File(label="Download Optimized (PDF)")
        download_after_word = gr.File(label="Download Optimized (Word)")

    run_button.click(
        process_resume,
        inputs=[resume_input, jd_input],
        outputs=[before_preview, before_md, after_md, download_before, download_after_md]
    )
    
    gr.Button("Export as PDF").click(
        export_as_pdf,
        inputs=[after_md],
        outputs=[download_after_pdf]
    )
    
    gr.Button("Export as Word").click(
        export_as_word,
        inputs=[after_md],
        outputs=[download_after_word]
    )

app.launch()