import os
import requests
from pypdf import PdfReader
import gradio as gr

# -------------------------------------------------------------------
# CONFIG
# -------------------------------------------------------------------
MODAL_API_URL = "https://chilukalagayathri--llama-resume-outreach-llamaresume-api.modal.run"

# -------------------------------------------------------------------
# HELPERS
# -------------------------------------------------------------------
def extract_resume_text(resume_path: str) -> str:
    """
    Extract text from a PDF or TXT resume.
    """
    if resume_path is None:
        return ""

    path_lower = resume_path.lower()
    text = ""

    try:
        if path_lower.endswith(".pdf"):
            reader = PdfReader(resume_path)
            for page in reader.pages:
                page_text = page.extract_text() or ""
                text += page_text + "\n"
        else:
            # Assume text file
            with open(resume_path, "r", encoding="utf-8", errors="ignore") as f:
                text = f.read()
    except Exception as e:
        return f"[ERROR READING RESUME: {e}]"

    # Truncate to avoid huge prompt
    return text.strip()[:6000]


def build_user_instruction(
    target_type: str,
    channel_type: str,
    job_description: str,
    resume_text: str,
) -> str:
    """
    Build the user prompt text based on:
    - Recruiter vs Referral
    - LinkedIn DM vs Email
    """
    base_context = f"""
Candidate resume:
{resume_text}

Job posting description:
{job_description}
""".strip()

    if target_type == "Recruiter":
        if channel_type == "LinkedIn DM":
            task = """
Write a concise LinkedIn message to a recruiter about this role.

Requirements:
- Tone: professional, friendly, and confident.
- 60–100 words.
- Clearly state why the candidate is a strong fit using 1–2 points from the resume.
- Refer to the role and company from the job description.
- End with a polite question about next steps or a short call.
"""
        else:  # Email
            task = """
Write a short email to a recruiter about this role.

Requirements:
- Tone: professional, warm, and clear.
- 120–180 words.
- Introduce the candidate briefly using information from the resume.
- Highlight 2–3 relevant skills or experiences that match the job posting.
- End with a clear call to action asking for a short call or information on next steps.
"""
    else:  # Referral
        if channel_type == "LinkedIn DM":
            task = """
Write a concise LinkedIn message asking a connection for a referral for this role.

Requirements:
- Tone: respectful, appreciative, and concise.
- 60–100 words.
- Mention why the candidate is a strong fit using 1–2 points from the resume.
- Clearly ask if the person would be comfortable referring them for the role.
"""
        else:  # Email
            task = """
Write a short email asking a connection for a referral for this role.

Requirements:
- Tone: respectful, appreciative, and clear.
- 120–180 words.
- Explain briefly why the candidate is a strong fit based on the resume.
- Reference the role and company.
- Politely ask if the person would be comfortable referring the candidate and offer to share additional details if needed.
"""

    user_content = base_context + "\n\n" + task
    return user_content.strip()


def build_prompt(
    target_type: str,
    channel_type: str,
    job_description: str,
    resume_text: str,
) -> str:
    """
    Build a Llama 3.1 chat-style prompt.
    """
    system_msg = (
        "You are an expert at writing personalized, professional outreach messages "
        "for job applications. Always return only the final message text, with no explanations."
    )

    user_msg = build_user_instruction(
        target_type=target_type,
        channel_type=channel_type,
        job_description=job_description,
        resume_text=resume_text,
    )

    # Llama 3.1 chat template
    prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>

{system_msg}<|eot_id|><|start_header_id|>user<|end_header_id|>

{user_msg}<|eot_id|><|start_header_id|>assistant<|end_header_id|>

"""
    return prompt


# -------------------------------------------------------------------
# INFERENCE FUNCTION - CALLS MODAL API
# -------------------------------------------------------------------
def generate_outreach_message(
    resume_path,
    target_type,
    channel_type,
    job_description,
    max_new_tokens,
    temperature,
    top_p,
):
    if resume_path is None:
        return "⚠️ Please upload a resume (PDF or TXT)."

    job_description = (job_description or "").strip()
    if not job_description:
        return "⚠️ Please paste the job posting description."

    # Extract resume text
    resume_text = extract_resume_text(resume_path)
    if not resume_text:
        return "⚠️ Could not read resume text. Please check the file."

    # Debug: Print resume text length
    print(f"✅ Resume extracted: {len(resume_text)} characters")

    # Build prompt
    prompt = build_prompt(
        target_type=target_type,
        channel_type=channel_type,
        job_description=job_description,
        resume_text=resume_text,
    )
    
    print(f"✅ Prompt built: {len(prompt)} characters")

    # Call Modal API
    try:
        payload = {
            "prompt": prompt,
            "max_length": int(max_new_tokens),
            "temperature": float(temperature),
        }
        
        print(f"🚀 Calling Modal API...")
        
        response = requests.post(
            MODAL_API_URL,
            json=payload,
            timeout=120  # 2 minute timeout for cold start
        )
        
        response.raise_for_status()
        result = response.json()
        
        print(f"✅ API Response received")
        print(f"Response keys: {result.keys()}")
        
        if "response" in result:
            clean_message = result["response"]
            print(f"✅ Clean message length: {len(clean_message)} characters")
            print(f"First 200 chars: {clean_message[:200]}")
            
            # Return ONLY the clean message
            return clean_message
            
        elif "error" in result:
            return f"❌ API Error: {result['error']}"
        else:
            return f"❌ Unexpected response: {result}"
            
    except requests.exceptions.Timeout:
        return "❌ Request timed out. The API might be cold-starting (takes 30-60s on first request). Please try again."
    except requests.exceptions.RequestException as e:
        return f"❌ Error calling Modal API: {e}"
    except Exception as e:
        return f"❌ Error while generating message: {e}"


# -------------------------------------------------------------------
# GRADIO UI
# -------------------------------------------------------------------
with gr.Blocks(title="Resume Outreach Generator") as demo:
    gr.Markdown(
        """
# 🎯 Resume Outreach Generator
### Powered by Fine-Tuned Llama 3.1 8B via Modal API

Upload your resume + job posting description and generate:
- Recruiter messages (LinkedIn DM or Email), or  
- Referral requests (LinkedIn DM or Email)

_Always review and personalize the message before sending._

---
"""
    )

    with gr.Row():
        with gr.Column():
            resume_file = gr.File(
                label="1️⃣ Upload your resume (PDF or TXT)",
                file_types=[".pdf", ".txt"],
                type="filepath",
            )

            with gr.Row():
                target_type = gr.Radio(
                    ["Recruiter", "Referral"],
                    value="Recruiter",
                    label="2️⃣ Message target",
                )
                channel_type = gr.Radio(
                    ["LinkedIn DM", "Email"],
                    value="LinkedIn DM",
                    label="3️⃣ Message format",
                )

            job_description = gr.Textbox(
                label="4️⃣ Job posting description",
                lines=10,
                placeholder="Paste the job description here...",
            )

            with gr.Accordion("⚙️ Advanced generation settings", open=False):
                max_new_tokens = gr.Slider(
                    64, 512, value=200, step=8, label="Max new tokens"
                )
                temperature = gr.Slider(
                    0.1, 1.5, value=0.7, step=0.05, label="Temperature"
                )
                top_p = gr.Slider(
                    0.1, 1.0, value=0.9, step=0.05, label="Top-p (not used in Modal API)"
                )

            generate_btn = gr.Button("✨ Generate outreach message", variant="primary", size="lg")
            
            gr.Markdown("""
            **⏱️ Response Time:**
            - First request: ~30-60 seconds (cold start)
            - Subsequent requests: ~2-3 seconds
            """)

        with gr.Column():
            output_box = gr.Textbox(
                label="Generated message",
                lines=18,
                placeholder="Your personalized outreach message will appear here...\n\nNote: First generation may take 30-60 seconds due to API cold start.",
            )

    generate_btn.click(
        fn=generate_outreach_message,
        inputs=[
            resume_file,
            target_type,
            channel_type,
            job_description,
            max_new_tokens,
            temperature,
            top_p,
        ],
        outputs=output_box,
    )

    gr.Markdown(
        """
---
### 📚 Resources

- 🤗 **Model**: [llama-3.1-8b-resume-outreach-v1](https://huggingface.co/GChilukala/llama-3.1-8b-resume-outreach-v1)
- 🚀 **API**: Deployed on [Modal](https://modal.com)
- 💾 **Training**: Fine-tuned on 4,091 resume outreach examples

### 🔧 Technical Details
- **Response time**: 2-3 seconds (warm) | 30-60 seconds (cold start)
- **GPU**: A10G
- **Architecture**: Serverless deployment via Modal

⚠️ **Disclaimer:** This tool assists with message drafting. Always review and customize before sending.
"""
    )

if __name__ == "__main__":
    demo.launch(
        share=False,
        server_name="0.0.0.0",
        server_port=7860,
    )