import streamlit as st
from PyPDF2 import PdfReader
from transformers import pipeline
import re
import random
from concurrent.futures import ThreadPoolExecutor
import time
import torch
import spacy
from reportlab.lib import colors
from reportlab.lib.pagesizes import letter
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image, Table, TableStyle
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.pdfbase import pdfmetrics
from reportlab.lib.enums import TA_CENTER
from io import BytesIO

# Question Templates
QUESTION_TEMPLATES = {
    "Experience": [
        "Could you describe your experience with {skill} and how you've applied it in your work?",
        "What specific challenges have you overcome while working with {technology}?",
        "Tell me about a successful project where you utilized {skill}. What was your role and the outcome?",
        "How have you implemented {skill} in your previous roles?",
        "What's the most complex problem you've solved using {technology}?"
    ],
    "Technical": [
        "Could you explain the key concepts of {technical_concept} and how you've implemented them?",
        "What approach would you take to implement {technical_process} in a production environment?",
        "How have you leveraged {technical_tool} to solve complex problems in your previous work?",
        "Explain how you would design a system using {technical_concept}.",
        "What are the best practices you follow when working with {technical_tool}?"
    ],
    "Behavioral": [
        "Describe a situation where you had to handle {situation}. What was your approach and the outcome?",
        "Can you share a specific example that demonstrates your {soft_skill} abilities?",
        "Tell me about a time when you faced {work_scenario}. How did you handle it?",
        "How do you typically approach {situation} in your work?",
        "What strategies do you use to demonstrate {soft_skill} in challenging situations?"
    ],
    "Role-specific": [
        "What strategies would you employ to effectively manage {job_task} in this role?",
        "How would you optimize {job_responsibility} to improve team productivity?",
        "What's your perspective on {industry_trend} and its impact on our industry?",
        "How would you handle {job_task} if resources were limited?",
        "What innovations would you bring to {job_responsibility}?"
    ]
}

# Streamlit configuration
st.set_page_config(
    page_title="Interview Preparation Coach",
    page_icon="💼",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Apply custom CSS
st.markdown("""
    <style>
        .main {
            padding: 2rem;
            background-color: #f8f9fa;
        }
        .stButton>button {
            width: 100%;
            background-color: #1E88E5;
            color: white;
        }
        .creator-header {
            color: #1E88E5;
            font-size: 2.5em;
            font-weight: bold;
            text-align: center;
            margin-bottom: 1em;
        }
        .creator-subheader {
            color: #424242;
            font-size: 1.2em;
            text-align: center;
            font-style: italic;
            margin-bottom: 2em;
        }
    </style>
""", unsafe_allow_html=True)

def extract_text_from_pdf(pdf_file):
    """Extract text from uploaded PDF file"""
    try:
        pdf = PdfReader(pdf_file)
        text = ""
        for page in pdf.pages:
            text += page.extract_text() + "\n"
        return re.sub(r'\s+', ' ', text).strip()
    except Exception as e:
        st.error(f"Error reading PDF: {str(e)}")
        return ""

@st.cache_resource
def load_nlp():
    try:
        return spacy.load('en_core_web_sm')
    except:
        spacy.cli.download('en_core_web_sm')
        return spacy.load('en_core_web_sm')

@st.cache_resource
def load_model():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    return pipeline(
        "text2text-generation",
        model="google/flan-t5-large",
        device=device,
        model_kwargs={"temperature": 0.7, "top_p": 0.9, "do_sample": True}
    )

def extract_keywords(text, nlp, max_keywords=30):
    """Enhanced keyword extraction using spaCy"""
    doc = nlp(text)
    keywords = []
    for token in doc:
        if token.pos_ in ['NOUN', 'PROPN'] and len(token.text) > 3:
            keywords.append(token.text.lower())
    for chunk in doc.noun_chunks:
        if len(chunk.text.split()) > 1:
            keywords.append(chunk.text.lower())
    return list(set(keywords))[:max_keywords]

def generate_single_qa(args):
    """Generate a single question and answer pair"""
    model, context, keywords, job_specific_terms, number = args
    try:
        category = random.choice(list(QUESTION_TEMPLATES.keys()))
        template = random.choice(QUESTION_TEMPLATES[category])
        
        replacements = {
            'skill': random.choice(job_specific_terms) if job_specific_terms else random.choice(keywords),
            'technology': random.choice(keywords),
            'technical_concept': random.choice(keywords),
            'technical_process': random.choice(job_specific_terms) if job_specific_terms else random.choice(keywords),
            'technical_tool': random.choice(keywords),
            'situation': random.choice([
                'a challenging deadline',
                'team conflicts',
                'unexpected project changes',
                'resource constraints',
                'stakeholder disagreements'
            ]),
            'soft_skill': random.choice([
                'leadership',
                'communication',
                'problem-solving',
                'adaptability',
                'project management',
                'team collaboration'
            ]),
            'work_scenario': random.choice([
                'a challenging project that required innovative solutions',
                'a situation where you had to lead a team through difficulty',
                'a critical decision that impacted project success'
            ]),
            'job_task': random.choice(job_specific_terms) if job_specific_terms else random.choice(keywords),
            'job_responsibility': random.choice(job_specific_terms) if job_specific_terms else random.choice(keywords),
            'industry_trend': random.choice(keywords)
        }
        
        question = template.format(**replacements)
        
        prompt = f"""
        Given the following context and question, provide a detailed, professional answer using the STAR method (Situation, Task, Action, Result) where appropriate. Make sure the response is specific and demonstrates expertise.

        Context: {context[:500]}
        Question: {question}

        Provide a comprehensive answer:
        """
        
        response = model(
            prompt,
            max_length=200,
            num_return_sequences=1,
            temperature=0.7,
            top_p=0.9
        )[0]['generated_text']
        
        response = response.strip()
        if not response.endswith(('.', '!', '?')):
            response += '.'
        
        return question, response
    except Exception as e:
        return f"Question {number}", f"Error generating response: {str(e)}"

def generate_interview_qa(cv_text, job_title, job_description, num_questions=50):
    model = load_model()
    nlp = load_nlp()
    start_time = time.time()
    
    context = f"Job Title: {job_title}\nJob Description: {job_description}\nCV: {cv_text}"
    job_specific_terms = extract_keywords(job_description, nlp, max_keywords=15)
    cv_keywords = extract_keywords(cv_text, nlp, max_keywords=20)
    keywords = list(set(cv_keywords + job_specific_terms))
    
    progress_bar = st.progress(0)
    status_text = st.empty()
    
    qa_pairs = []
    with ThreadPoolExecutor(max_workers=4) as executor:
        args_list = [(model, context, keywords, job_specific_terms, i) for i in range(num_questions)]
        for i, result in enumerate(executor.map(generate_single_qa, args_list)):
            qa_pairs.append(result)
            progress = (i + 1) / num_questions
            progress_bar.progress(progress)
            elapsed_time = time.time() - start_time
            estimated_total = elapsed_time / progress if progress > 0 else 0
            remaining_time = max(0, estimated_total - elapsed_time)
            status_text.text(f"Generated {i+1}/{num_questions} questions. Estimated time remaining: {remaining_time:.1f}s")
    
    return qa_pairs

def generate_pdf(qa_pairs, job_title):
    buffer = BytesIO()
    doc = SimpleDocTemplate(
        buffer,
        pagesize=letter,
        rightMargin=72,
        leftMargin=72,
        topMargin=72,
        bottomMargin=72
    )
    
    styles = getSampleStyleSheet()
    title_style = ParagraphStyle(
        'CustomTitle',
        parent=styles['Heading1'],
        fontSize=24,
        spaceAfter=30,
        alignment=TA_CENTER
    )
    question_style = ParagraphStyle(
        'Question',
        parent=styles['Heading2'],
        fontSize=14,
        textColor=colors.HexColor('#1E88E5'),
        spaceAfter=12
    )
    answer_style = ParagraphStyle(
        'Answer',
        parent=styles['Normal'],
        fontSize=12,
        spaceAfter=20
    )
    
    story = []
    
    story.append(Paragraph(f"Interview Preparation Guide<br/>for {job_title}", title_style))
    story.append(Spacer(1, 30))
    
    watermark_style = ParagraphStyle(
        'Watermark',
        parent=styles['Normal'],
        fontSize=10,
        textColor=colors.gray,
        alignment=TA_CENTER
    )
    story.append(Paragraph("Created by Muhammad Shaheer<br/>Professional Interview Coach", watermark_style))
    story.append(Spacer(1, 30))
    
    quote_style = ParagraphStyle(
        'Quote',
        parent=styles['Italic'],
        fontSize=12,
        textColor=colors.HexColor('#666666'),
        alignment=TA_CENTER,
        spaceAfter=30
    )
    story.append(Paragraph(
        '"Success is not final, failure is not fatal: '
        'it is the courage to continue that counts."<br/>- Winston Churchill',
        quote_style
    ))
    
    for i, (question, answer) in enumerate(qa_pairs, 1):
        story.append(Paragraph(f"Q{i}: {question}", question_style))
        story.append(Paragraph(f"A: {answer}", answer_style))
        story.append(Spacer(1, 10))
    
    footer_style = ParagraphStyle(
        'Footer',
        parent=styles['Normal'],
        fontSize=8,
        textColor=colors.gray,
        alignment=TA_CENTER
    )
    story.append(Spacer(1, 30))
    story.append(Paragraph(
        "This document is generated by Interview Preparation Coach<br/>"
        "Copyright © 2024 Muhammad Shaheer. All rights reserved.",
        footer_style
    ))
    
    doc.build(story)
    buffer.seek(0)
    return buffer

def main():
    st.markdown('<p class="creator-header">Interview Preparation Coach</p>', unsafe_allow_html=True)
    st.markdown('<p class="creator-subheader">Created by Muhammad Shaheer</p>', unsafe_allow_html=True)
    
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.subheader("📄 Upload Your CV")
        cv_file = st.file_uploader("Upload your CV (PDF format only):", type=["pdf"])
        if cv_file:
            cv_text = extract_text_from_pdf(cv_file)
            if cv_text:
                st.success("✅ CV uploaded successfully!")
            else:
                st.error("❌ Error reading PDF. Please try again.")
                return
    
    with col2:
        st.subheader("🎯 Job Details")
        job_title = st.text_input("Enter the job title:", placeholder="e.g., Senior Software Engineer")
        job_description = st.text_area("Paste the job description:", height=150)
    
    with st.expander("⚙️ Advanced Options"):
        num_questions = st.slider("Number of questions to generate", 5, 50, 20)
        categories = st.multiselect(
            "Question categories to include",
            list(QUESTION_TEMPLATES.keys()),
            default=list(QUESTION_TEMPLATES.keys())
        )
    
    if st.button("🚀 Generate Interview Questions & Answers"):
        if cv_file and job_title and job_description:
            with st.spinner("🔄 Generating your personalized interview guide..."):
                qa_pairs = generate_interview_qa(cv_text, job_title, job_description, num_questions)
                
                st.session_state.qa_pairs = qa_pairs
                st.session_state.job_title = job_title
                
                st.subheader("📚 Generated Interview Questions and Answers:")
                for i, (question, answer) in enumerate(qa_pairs, 1):
                    with st.expander(f"Q{i}: {question}"):
                        st.markdown("**Answer:**")
                        st.write(answer)
                
                pdf_buffer = generate_pdf(qa_pairs, job_title)
                st.download_button(
                    label="📄 Download PDF Guide",
                    data=pdf_buffer,
                    file_name=f"Interview_Guide_{job_title.replace(' ', '_')}.pdf",
                    mime="application/pdf"
                )
        else:
            st.error("⚠️ Please provide all required information (CV, job title, and job description).")
    
    st.markdown("---")
    st.markdown("""
        <div style='text-align: center'>
            <h4>💼 Interview Preparation Coach</h4>
            <p>Created by Muhammad Shaheer</p>
            <p><small>Powered by Advanced AI Technology • Professional Interview Preparation Tool</small></p>
        </div>
    """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()