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 research, role problem analysis, and achievement matching 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, "role_problem": str, "achievement_section": str}
    """
    
    prompt = f"""You are an expert career strategist researching a company and role to craft authentic, researched-backed cover letter context.

Your task: Use web search to find SPECIFIC, RECENT company intelligence, then match it to the candidate's achievements.

REQUIRED WEB SEARCHES:
1. Recent company moves (funding rounds, product launches, acquisitions, market expansion, hiring momentum)
2. Current company challenges (what problem are they actively solving?)
3. Company tech stack / tools they use
4. Why they're hiring NOW (growth? new product? team expansion?)
5. Company market position and strategy

Job Description:
{job_description}

Candidate's Portfolio:
{portfolio}

Generate a JSON response with this format (no additional text):
{{
  "company_motivation": "2-3 sentences showing specific, researched interest. Reference recent company moves (funding, product launches, market position) OR specific challenge. Format: '[Company name] recently [specific move/challenge], and your focus on [specific need] aligns with my experience building [domain].' Avoid forced connections—if authenticity is low, keep motivation minimal.",
  "role_problem": "1 sentence defining CORE PROBLEM this role solves for company. Example: 'Improving demand forecasting accuracy for franchisee decision-making' OR 'Building production vision models under real-time latency constraints.'",
  "achievement_section": "ONE specific achievement from portfolio solving role_problem (not just relevant). Format: 'Built [X] to solve [problem/constraint], achieving [metric] across [scale].' Example: 'Built self-serve ML agents (FastAPI+LangChain) to reduce business team dependency on Data Engineering by 60% across 150k+ samples.' This must map directly to role_problem."
}}

REQUIREMENTS FOR AUTHENTICITY:
- company_motivation: Must reference verifiable findings from web search (recent news, funding, product launch, specific challenge)
- role_problem: Explicitly state the core problem extracted from job description + company research
- achievement_section: Must map directly to role_problem with clear cause-effect (not just "relevant to job")
- NO FORCED CONNECTIONS: If no genuine connection exists between candidate achievement and role problem, return empty string rather than forcing weak match
- AUTHENTICITY PRIORITY: A short, genuine motivation beats a longer forced one. Minimize if needed to avoid template-feel.

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", ""),
            "role_problem": result.get("role_problem", ""),
            "achievement_section": result.get("achievement_section", "")
        }
    except json.JSONDecodeError:
        # Fallback if JSON parsing fails
        return {
            "company_motivation": "",
            "role_problem": "",
            "achievement_section": ""
        }


def handle_cover_letter(job_description, portfolio, api_key, company_motivation="", role_problem="", 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: Researched company interest with recent moves/challenges (auto-generated if empty)
        role_problem: The core problem this role solves for the company (auto-generated if empty)
        specific_achievement: One concrete achievement that solves role_problem (auto-generated if empty)
    """
    
    # Build context sections if provided
    motivation_section = ""
    if company_motivation.strip():
        motivation_section = f"\nCompany Research (Recent Moves/Challenges):\n{company_motivation}"
    
    problem_section = ""
    if role_problem.strip():
        problem_section = f"\nRole's Core Problem:\n{role_problem}"
    
    achievement_section = ""
    if specific_achievement.strip():
        achievement_section = f"\nAchievement That Solves This Problem:\n{specific_achievement}"
    
    prompt = f"""You are an expert career coach writing authentic, researched cover letters that prove specific company knowledge and solve real problems.

Your goal: Write a letter showing you researched THIS company (not a template) and authentically connect your achievements to THEIR specific problem.

CRITICAL FOUNDATION:
You have three inputs: company research (recent moves/challenges), role problem (what they're hiring to solve), and one matching achievement.
Construct narrative: "Because you [company context] need to [role problem], my experience with [achievement] makes me valuable."

Cover Letter Structure:
1. Opening (2-3 sentences): Hook with SPECIFIC company research (recent move, funding, product, market challenge)
   - NOT: "I'm interested in your company"
   - YES: "Your recent expansion to [X markets] and focus on [tech] align with my experience"
   
2. Middle (4-5 sentences):
   - State role's core problem (what you understand they're hiring to solve)
   - Connect achievement DIRECTLY to that problem (show cause-effect)
   - Reference job description specifics your achievement addresses
   - Show understanding of their constraint/challenge
   
3. Closing (1-2 sentences): Express genuine enthusiasm about solving THIS specific problem

CRITICAL REQUIREMENTS:
- RESEARCH PROOF: Opening must show specific company knowledge (recent news, not generic mission)
- PROBLEM CLARITY: Explicitly state what problem you're solving for them
- SPECIFIC MAPPING: Achievement → Role Problem → Company Need (clear cause-effect chain)
- NO TEMPLATE: Varied sentence length, conversational tone, human voice
- NO FORCED CONNECTIONS: If something doesn't link cleanly, leave it out
- NO FLUFF: Every sentence serves a purpose (authentic < complete)
- NO SALARY TALK: Omit expectations or negotiations
- NO CORPORATE JARGON: Write like a real human
- NO EM DASHES: Use commas or separate sentences

Formatting:
- Start: "Dear Hiring Manager,"
- End: "Best,\nDhanvanth Voona" (on separate lines)
- Max 250 words
- NO PREAMBLE (start directly)
- Multiple short paragraphs OK

Context for Writing:
Resume:
{portfolio}

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

Response (Max 250 words, researched + authentic 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, salary ranges 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
- For questions related to salary, use the same salary ranges if provided in job description, ONLY if you could not extract salary from 
  job description, use the salary range given in portfolio. 

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("📊 Researching company, analyzing role, and matching achievements (with web search)...")
                context = generate_cover_letter_context(job_description, selected_portfolio, openrouter_api_key)
                company_motivation = context.get("company_motivation", "")
                role_problem = context.get("role_problem", "")
                specific_achievement = context.get("achievement_section", "")
                context_placeholder.success("✅ Company research and achievement matching complete!")
            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 = ""
                role_problem = ""
                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,
                        role_problem=role_problem,
                        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 ❤️"
)