import streamlit as st
import time
from utils.text_analysis import (
    extract_keywords, generate_summary, generate_mcqs,
    generate_coding_questions, analyze_interview_response,
    get_learning_resources
)
from utils.audio_processor import convert_audio_to_text

def main():
    st.title("AI Interview Preparation Assistant")
    
    # Session state initialization
    if 'current_step' not in st.session_state:
        st.session_state.current_step = 1
    
    # Step 1: Input Collection
    if st.session_state.current_step == 1:
        st.header("Step 1: Upload Resume and Job Description")
        resume = st.text_area("Paste your resume here")
        job_description = st.text_area("Paste the job description here")
        
        if st.button("Analyze") and resume and job_description:
            with st.spinner("Analyzing..."):
                st.session_state.resume_summary = generate_summary(resume)
                st.session_state.jd_summary = generate_summary(job_description)
                st.session_state.keywords = extract_keywords(job_description)
                st.session_state.current_step = 2
    
    # Step 2: Summary and Keywords
    elif st.session_state.current_step == 2:
        st.header("Step 2: Analysis Results")
        st.subheader("Resume Summary")
        st.write(st.session_state.resume_summary)
        
        st.subheader("Job Description Summary")
        st.write(st.session_state.jd_summary)
        
        st.subheader("Key Requirements")
        st.write(", ".join(st.session_state.keywords))
        
        if st.button("Start Assessment"):
            st.session_state.current_step = 3
    
    # Step 3: MCQ Assessment
    elif st.session_state.current_step == 3:
        st.header("Step 3: MCQ Assessment")
        if 'mcqs' not in st.session_state:
            with st.spinner("Generating questions..."):
                st.session_state.mcqs = generate_mcqs(
                    st.session_state.resume_summary,
                    st.session_state.jd_summary
                )
                st.session_state.mcq_answers = {}
                st.session_state.mcq_start_time = time.time()
        
        # Display timer
        elapsed_time = int(time.time() - st.session_state.mcq_start_time)
        st.write(f"Time elapsed: {elapsed_time} seconds")
        
        # Display MCQs
        for i, mcq in enumerate(st.session_state.mcqs):
            answer = st.radio(
                f"Q{i+1}: {mcq['question']}",
                mcq['options'],
                key=f"mcq_{i}"
            )
            st.session_state.mcq_answers[i] = answer
        
        if st.button("Submit MCQs"):
            st.session_state.current_step = 4
    
    # Step 4: Coding Assessment
    elif st.session_state.current_step == 4:
        st.header("Step 4: Coding Assessment")
        if 'coding_questions' not in st.session_state:
            with st.spinner("Generating coding questions..."):
                st.session_state.coding_questions = generate_coding_questions(
                    st.session_state.resume_summary,
                    st.session_state.jd_summary
                )
                st.session_state.coding_start_time = time.time()
        
        # Display timer
        elapsed_time = int(time.time() - st.session_state.coding_start_time)
        st.write(f"Time elapsed: {elapsed_time} seconds")
        
        for i, question in enumerate(st.session_state.coding_questions):
            st.subheader(f"Question {i+1}")
            st.write(question['problem_statement'])
            st.write("Examples:")
            st.write(question['examples'])
            st.write("Constraints:")
            st.write(question['constraints'])
            
            if 'coding_answers' not in st.session_state:
                st.session_state.coding_answers = {}
            
            st.session_state.coding_answers[i] = st.text_area(
                "Your solution:",
                key=f"coding_{i}"
            )
        
        if st.button("Submit Coding Solutions"):
            st.session_state.current_step = 5
    
    # Step 5: Mock Interview
    elif st.session_state.current_step == 5:
        st.header("Step 5: Mock Interview")
        st.write("Click 'Start' when ready to begin the interview.")
        
        if st.button("Start Interview"):
            st.write("Recording... Speak your answer.")
            audio_text = convert_audio_to_text()
            
            with st.spinner("Analyzing your response..."):
                st.session_state.interview_analysis = analyze_interview_response(
                    audio_text,
                    st.session_state.jd_summary
                )
            st.session_state.current_step = 6
    
    # Step 6: Final Analysis
    elif st.session_state.current_step == 6:
        st.header("Final Analysis")
        
        # Calculate overall score
        mcq_score = sum(1 for i, ans in st.session_state.mcq_answers.items() 
                       if ans == st.session_state.mcqs[i]['correct_answer'])
        interview_score = st.session_state.interview_analysis['score']
        
        overall_score = (mcq_score / 30 * 0.4 + interview_score / 100 * 0.6) * 100
        
        st.subheader("Overall Match Score")
        st.progress(overall_score / 100)
        st.write(f"{overall_score:.1f}%")
        
        st.subheader("Strengths")
        st.write(st.session_state.interview_analysis['strengths'])
        
        st.subheader("Areas for Improvement")
        st.write(st.session_state.interview_analysis['improvements'])
        
        # Get learning resources
        resources = get_learning_resources(
            st.session_state.interview_analysis['concepts_to_study']
        )
        
        st.subheader("Recommended Learning Resources")
        for concept, links in resources.items():
            st.write(f"**{concept}:**")
            for link in links:
                st.write(f"- {link}")

if __name__ == "__main__":
    main()