import streamlit as st
from openai import OpenAI
import os
from dotenv import load_dotenv
from datetime import datetime
import pytz
from reportlab.lib.pagesizes import letter
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, PageBreak
from reportlab.lib.enums import TA_LEFT, TA_JUSTIFY
import io

# Load environment variables
load_dotenv()

# Page configuration
st.set_page_config(page_title="AI Resume Assistant", layout="wide")
st.title("🤖 AI Resume Assistant")

# Load API keys from environment variables
openrouter_api_key = os.getenv("OPENROUTER_API_KEY")
openai_api_key = os.getenv("OPENAI_API_KEY")

# Check if API keys are available
if not openrouter_api_key or not openai_api_key:
    st.error("❌ API keys not found. Please set OPENROUTER_API_KEY and OPENAI_API_KEY in your environment variables (.env file).")
    st.stop()

def get_est_timestamp():
    """Get current timestamp in EST timezone with format dd-mm-yyyy-HH-MM"""
    est = pytz.timezone('US/Eastern')
    now = datetime.now(est)
    return now.strftime("%d-%m-%Y-%H-%M")


def generate_pdf(content, filename):
    """Generate PDF from content and return as bytes"""
    try:
        pdf_buffer = io.BytesIO()
        doc = SimpleDocTemplate(
            pdf_buffer,
            pagesize=letter,
            rightMargin=0.75*inch,
            leftMargin=0.75*inch,
            topMargin=0.75*inch,
            bottomMargin=0.75*inch
        )
        
        story = []
        styles = getSampleStyleSheet()
        
        # Custom style for body text
        body_style = ParagraphStyle(
            'CustomBody',
            parent=styles['Normal'],
            fontSize=11,
            leading=14,
            alignment=TA_JUSTIFY,
            spaceAfter=12
        )
        
        # Add content only (no preamble)
        # Split content into paragraphs for better formatting
        paragraphs = content.split('\n\n')
        for para in paragraphs:
            if para.strip():
                # Replace line breaks with spaces within paragraphs
                clean_para = para.replace('\n', ' ').strip()
                story.append(Paragraph(clean_para, body_style))
        
        # Build PDF
        doc.build(story)
        pdf_buffer.seek(0)
        return pdf_buffer.getvalue()
    
    except Exception as e:
        st.error(f"Error generating PDF: {str(e)}")
        return None


def categorize_input(resume_finder, cover_letter, select_resume, entry_query):
    """
    Categorize input into one of 4 groups:
    - resume_finder: T, F, No Select
    - cover_letter: F, T, not No Select
    - general_query: F, F, not No Select
    - retry: any other combination
    """
    
    if resume_finder and not cover_letter and select_resume == "No Select":
        return "resume_finder", None
    
    elif not resume_finder and cover_letter and select_resume != "No Select":
        return "cover_letter", None
    
    elif not resume_finder and not cover_letter and select_resume != "No Select":
        if not entry_query.strip():
            return "retry", "Please enter a query for General Query mode."
        return "general_query", None
    
    else:
        return "retry", "Please check your entries and try again"


def load_portfolio(file_path):
    """Load portfolio markdown file"""
    try:
        full_path = os.path.join(os.path.dirname(__file__), file_path)
        with open(full_path, 'r', encoding='utf-8') as f:
            return f.read()
    except FileNotFoundError:
        st.error(f"Portfolio file {file_path} not found!")
        return None


def handle_resume_finder(job_description, ai_portfolio, ds_portfolio, api_key):
    """Handle Resume Finder category using OpenRouter"""
    
    prompt = f"""You are an expert resume matcher. Analyze the following job description and two portfolios to determine which is the best match.

IMPORTANT MAPPING:
- If AI_portfolio is most relevant → Resume = Resume_P
- If DS_portfolio is most relevant → Resume = Resume_Dss

Job Description:
{job_description}

AI_portfolio (Maps to: Resume_P):
{ai_portfolio}

DS_portfolio (Maps to: Resume_Dss):
{ds_portfolio}

Respond ONLY with:
Resume: [Resume_P or Resume_Dss]
Reasoning: [25-30 words explaining the match]

NO PREAMBLE."""
    
    try:
        client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=api_key,
        )
        
        completion = client.chat.completions.create(
            model="openai/gpt-oss-safeguard-20b",
            messages=[
                {
                    "role": "user",
                    "content": prompt
                }
            ]
        )
        
        response = completion.choices[0].message.content
        if response:
            return response
        else:
            st.error("❌ No response received from OpenRouter API")
            return None
    
    except Exception as e:
        st.error(f"❌ Error calling OpenRouter API: {str(e)}")
        return None


def generate_cover_letter_context(job_description, portfolio, api_key):
    """Generate company motivation and achievement section using web search via Perplexity Sonar
    
    Args:
        job_description: The job posting
        portfolio: Candidate's resume/portfolio
        api_key: OpenRouter API key (used for Perplexity Sonar with web search)
    
    Returns:
        dict: {"company_motivation": str, "achievement_section": str}
    """
    
    prompt = f"""You are an expert career strategist. Your task is to generate two specific, personalized inputs for a cover letter.

Use your web search capability to find relevant company information. Given the job description and candidate's portfolio, you will:
1. Search for relevant company information (mission, values, recent projects, culture)
2. Identify the best achievement from the portfolio that matches the role
3. Generate targeted content for cover letter generation

Job Description:
{job_description}

Candidate's Portfolio:
{portfolio}

Generate a JSON response with exactly this format (no additional text):
{{
  "company_motivation": "1-2 sentences showing specific interest in THIS company based on their mission/values/recent work. Use your web search to find specific details about the company. Should feel genuine and specific, not generic.",
  "achievement_section": "One specific, quantified achievement from the portfolio that directly supports the job requirements. Format: 'achievement that resulted in specific outcome'."
}}

Requirements for company_motivation:
- Must reference something specific from the job description OR from web search about the company (company needs, projects, values, recent news)
- Should show research and genuine interest using real company information
- 1-2 sentences maximum
- Sound natural and authentic

Requirements for achievement_section:
- Must be concrete and specific
- Should include numbers/metrics when possible
- Must be relevant to the job requirements
- Maximum 1 sentence

Return ONLY the JSON object, no other text."""
    
    # Use Perplexity Sonar via OpenRouter (has built-in web search)
    client = OpenAI(
        base_url="https://openrouter.ai/api/v1",
        api_key=api_key,
    )
    
    completion = client.chat.completions.create(
        model="perplexity/sonar",
        messages=[
            {
                "role": "user",
                "content": prompt
            }
        ]
    )
    
    response_text = completion.choices[0].message.content
    
    # Parse JSON response
    import json
    try:
        result = json.loads(response_text)
        return {
            "company_motivation": result.get("company_motivation", ""),
            "achievement_section": result.get("achievement_section", "")
        }
    except json.JSONDecodeError:
        # Fallback if JSON parsing fails
        return {
            "company_motivation": "",
            "achievement_section": ""
        }


def handle_cover_letter(job_description, portfolio, api_key, company_motivation="", specific_achievement=""):
    """Handle Cover Letter category using OpenAI
    
    Args:
        job_description: The job posting
        portfolio: Candidate's resume/portfolio
        api_key: OpenAI API key
        company_motivation: Why candidate is interested in THIS company/role (auto-generated if empty)
        specific_achievement: One concrete achievement to leverage (auto-generated if empty)
    """
    
    # Build context about company motivation if provided
    motivation_section = ""
    if company_motivation.strip():
        motivation_section = f"\nCandidate's Interest in This Role:\n{company_motivation}"
    
    achievement_section = ""
    if specific_achievement.strip():
        achievement_section = f"\nKey Achievement to Reference:\n{specific_achievement}"
    
    prompt = f"""You are an expert career coach writing authentic, human cover letters that stand out—not generic templates.

Your goal: Write a cover letter that feels like it was written by the actual candidate, showing genuine interest and proof of capability.

Cover Letter Structure (follow this order):
1. Opening (2-3 sentences): Hook with specific reason for interest in THIS company, not generic
2. Middle (4-5 sentences): 
   - Show you researched them (reference job description specifics)
   - Connect 1-2 resume achievements directly to their needs
   - Briefly mention the achievement below to prove capability
3. Closing (1-2 sentences): Express enthusiasm and leave door open

Critical Requirements for Authenticity:
- Write like a real person, NOT a template (varied sentence length, conversational where appropriate)
- Show personality through word choice—confident but humble, professional but warm
- Every claim must link to either the job description or the achievement below
- Use specific details from the resume and job posting (shows real attention)
- NO fluff, NO corporate jargon, NO redundancy
- If something doesn't connect, don't force it
- Sound like someone who actually wants this job, not just applying to any job
- Do NOT mention salary expectations or benefits negotiations

Formatting Requirements:
- Start with "Dear Hiring Manager,"
- End with: "Best,\nDhanvanth Voona" (Best on one line, name on next line)
- Maximum 250 words (tight constraint = only include essential points)
- NO PREAMBLE (begin directly with opening)
- STRICTLY NO em dashes (use commas or separate sentences instead)
- Single paragraphs are fine; multiple short paragraphs OK

Context for Authentic Writing:
Resume:
{portfolio}

Job Description:
{job_description}{motivation_section}{achievement_section}

Response (Max 250 words, genuine tone):"""
    
    client = OpenAI(api_key=api_key)
    
    completion = client.chat.completions.create(
        model="gpt-5-mini-2025-08-07",
        messages=[
            {
                "role": "user",
                "content": prompt
            }
        ]
    )
    
    response = completion.choices[0].message.content
    return response


def handle_general_query(job_description, portfolio, query, length, api_key):
    """Handle General Query category using OpenAI"""
    
    word_count_map = {
        "short": "40-60",
        "medium": "80-100",
        "long": "120-150"
    }
    
    word_count = word_count_map.get(length, "40-60")
    
    prompt = f"""You are an expert career consultant helping a candidate answer application questions with authentic, tailored responses.

Your task: Answer the query authentically using ONLY genuine connections between the candidate's experience and the job context.

Word Count Strategy (Important - Read Carefully):
- Target: {word_count} words MAXIMUM
- Adaptive: Use fewer words if the query can be answered completely and convincingly with fewer words
- Examples: "What is your greatest strength?" might need only 45 words. "Why our company?" needs 85-100 words to show genuine research
- NEVER force content to hit word count targets - prioritize authentic connection over word count

Connection Quality Guidelines:
- Extract key company values/needs from job description
- Find 1-2 direct experiences from resume that align with these
- Show cause-and-effect: "Because you need X, my experience with Y makes me valuable"
- If connection is weak or forced, acknowledge limitations honestly
- Avoid generic statements - every sentence should reference either the job, company, or specific experience

Requirements:
- Answer naturally as if written by the candidate
- Start directly with the answer (NO PREAMBLE or "Let me tell you...")
- Response must be directly usable in an application
- Make it engaging and personalized, not templated
- STRICTLY NO EM DASHES
- One authentic connection beats three forced ones

Resume:
{portfolio}

Job Description:
{job_description}

Query:
{query}

Response (Max {word_count} words, use fewer if appropriate):"""
    
    client = OpenAI(api_key=api_key)
    
    completion = client.chat.completions.create(
        model="gpt-5-mini-2025-08-07",
        messages=[
            {
                "role": "user",
                "content": prompt
            }
        ]
    )
    
    response = completion.choices[0].message.content
    return response


# Main input section
st.header("📋 Input Form")

# Create columns for better layout
col1, col2 = st.columns(2)

with col1:
    job_description = st.text_area(
        "Job Description (Required)*",
        placeholder="Paste the job description here...",
        height=150
    )

with col2:
    st.subheader("Options")
    resume_finder = st.checkbox("Resume Finder", value=False)
    cover_letter = st.checkbox("Cover Letter", value=False)

# Length of Resume
length_options = {
    "Short (40-60 words)": "short",
    "Medium (80-100 words)": "medium",
    "Long (120-150 words)": "long"
}
length_of_resume = st.selectbox(
    "Length of Resume",
    options=list(length_options.keys()),
    index=0
)
length_value = length_options[length_of_resume]

# Select Resume dropdown
resume_options = ["No Select", "Resume_P", "Resume_Dss"]
select_resume = st.selectbox(
    "Select Resume",
    options=resume_options,
    index=0
)

# Entry Query
entry_query = st.text_area(
    "Entry Query (Optional)",
    placeholder="Ask any question related to your application...",
    max_chars=5000,
    height=100
)

# Submit button
if st.button("🚀 Generate", type="primary", use_container_width=True):
    # Validate job description
    if not job_description.strip():
        st.error("❌ Job Description is required!")
        st.stop()
    
    # Categorize input
    category, error_message = categorize_input(
        resume_finder, cover_letter, select_resume, entry_query
    )
    
    if category == "retry":
        st.warning(f"⚠️ {error_message}")
    else:
        st.header("📤 Response")
        
        # Debug info (can be removed later)
        with st.expander("📊 Debug Info"):
            st.write(f"**Category:** {category}")
            st.write(f"**Resume Finder:** {resume_finder}")
            st.write(f"**Cover Letter:** {cover_letter}")
            st.write(f"**Select Resume:** {select_resume}")
            st.write(f"**Has Query:** {bool(entry_query.strip())}")
            st.write(f"**OpenAI API Key Set:** {'✅ Yes' if openai_api_key else '❌ No'}")
            st.write(f"**OpenRouter API Key Set:** {'✅ Yes' if openrouter_api_key else '❌ No'}")
            st.write(f"**OpenAI Key First 10 chars:** {openai_api_key[:10] + '...' if openai_api_key else 'N/A'}")
            st.write(f"**OpenRouter Key First 10 chars:** {openrouter_api_key[:10] + '...' if openrouter_api_key else 'N/A'}")
        
        # Load portfolios
        ai_portfolio = load_portfolio("AI_portfolio.md")
        ds_portfolio = load_portfolio("DS_portfolio.md")
        
        if ai_portfolio is None or ds_portfolio is None:
            st.stop()
        
        response = None
        error_occurred = None
        
        if category == "resume_finder":
            with st.spinner("🔍 Finding the best resume for you..."):
                try:
                    response = handle_resume_finder(
                        job_description, ai_portfolio, ds_portfolio, openrouter_api_key
                    )
                except Exception as e:
                    error_occurred = f"Resume Finder Error: {str(e)}"
        
        elif category == "cover_letter":
            selected_portfolio = ai_portfolio if select_resume == "Resume_P" else ds_portfolio
            
            # Generate company motivation and achievement section
            st.info("🔍 Analyzing company and generating personalized context with web search...")
            context_placeholder = st.empty()
            
            try:
                context_placeholder.info("📊 Generating company motivation and achievement section (with web search)...")
                context = generate_cover_letter_context(job_description, selected_portfolio, openrouter_api_key)
                company_motivation = context.get("company_motivation", "")
                specific_achievement = context.get("achievement_section", "")
                context_placeholder.success("✅ Context generated successfully with web search!")
            except Exception as e:
                error_occurred = f"Context Generation Error: {str(e)}"
                context_placeholder.error(f"❌ Failed to generate context: {str(e)}")
                st.info("💡 Proceeding with cover letter generation without auto-generated context...")
                company_motivation = ""
                specific_achievement = ""
            
            # Now generate the cover letter
            with st.spinner("✍️ Generating your cover letter..."):
                try:
                    response = handle_cover_letter(
                        job_description, selected_portfolio, openai_api_key,
                        company_motivation=company_motivation,
                        specific_achievement=specific_achievement
                    )
                except Exception as e:
                    error_occurred = f"Cover Letter Error: {str(e)}"
        
        elif category == "general_query":
            selected_portfolio = ai_portfolio if select_resume == "Resume_P" else ds_portfolio
            with st.spinner("💭 Crafting your response..."):
                try:
                    response = handle_general_query(
                        job_description, selected_portfolio, entry_query, 
                        length_value, openai_api_key
                    )
                except Exception as e:
                    error_occurred = f"General Query Error: {str(e)}"
        
        # Display error if one occurred
        if error_occurred:
            st.error(f"❌ {error_occurred}")
            st.info("💡 **Troubleshooting Tips:**\n- Check your API keys in the .env file\n- Verify your API key has sufficient credits/permissions\n- Ensure the model name is correct for your API tier")
        
        # Store response in session state only if new response generated
        if response:
            st.session_state.edited_response = response
            st.session_state.editing = False
        elif not error_occurred:
            st.error("❌ Failed to generate response. Please check the error messages above and try again.")

# Display stored response if available (persists across button clicks)
if "edited_response" in st.session_state and st.session_state.edited_response:
    st.header("📤 Response")
    
    # Toggle edit mode
    col_response, col_buttons = st.columns([3, 1])
    
    with col_buttons:
        if st.button("✏️ Edit", key="edit_btn", use_container_width=True):
            st.session_state.editing = not st.session_state.editing
    
    # Display response or edit area
    if st.session_state.editing:
        st.session_state.edited_response = st.text_area(
            "Edit your response:",
            value=st.session_state.edited_response,
            height=250,
            key="response_editor"
        )
        
        col_save, col_cancel = st.columns(2)
        with col_save:
            if st.button("💾 Save Changes", use_container_width=True):
                st.session_state.editing = False
                st.success("✅ Response updated!")
                st.rerun()
        
        with col_cancel:
            if st.button("❌ Cancel", use_container_width=True):
                st.session_state.editing = False
                st.rerun()
    else:
        # Display the response
        st.success(st.session_state.edited_response)
        
        # Download PDF button
        timestamp = get_est_timestamp()
        pdf_filename = f"Dhanvanth_{timestamp}.pdf"
        
        pdf_content = generate_pdf(st.session_state.edited_response, pdf_filename)
        if pdf_content:
            st.download_button(
                label="📥 Download as PDF",
                data=pdf_content,
                file_name=pdf_filename,
                mime="application/pdf",
                use_container_width=True
            )

st.markdown("---")
st.markdown(
    "Say Hi to Griva thalli from her mama ❤️"
)