import os
import base64
from PIL import Image
import streamlit as st
from evaluator import (
    IELTSTask1Evaluator,
    IELTSTask1ExerciseGenerator,
    IELTSTask2Evaluator,
    IELTSTask2ExerciseGenerator
)

# Initialize session state for login
if 'logged_in' not in st.session_state:
    st.session_state.logged_in = False

def check_password():
    """Returns `True` if the user had the correct password."""
    
    def password_entered():
        """Checks whether a password entered by the user is correct."""
        if st.session_state["password"] == "hamid1377" and st.session_state["username"] == "afarinesh":
            st.session_state.logged_in = True
            del st.session_state["password"]  # Don't store password
            del st.session_state["username"]  # Don't store username
        else:
            st.session_state.logged_in = False

    if not st.session_state.logged_in:
        st.text_input("Username", key="username")
        st.text_input("Password", type="password", key="password")
        st.button("Login", on_click=password_entered)
        return False
    
    return True

# Configure Streamlit page
st.set_page_config(
    page_title="IELTS Writing Evaluator",
    page_icon="📝",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS
st.markdown("""
<style>
    .stTabs [data-baseweb="tab-list"] {
        gap: 2px;
    }
    .stTabs [data-baseweb="tab"] {
        padding: 10px 20px;
    }
    .main .block-container {
        padding-top: 2rem;
    }
    .stButton button {
        width: 100%;
    }
    .assessment-box {
        padding: 1rem;
        border-radius: 0.5rem;
        border: 1px solid #e0e0e0;
        margin: 1rem 0;
    }
    .info-box {
        background-color: #f8f9fa;
        padding: 1rem;
        border-radius: 0.5rem;
        margin: 1rem 0;
    }
</style>
""", unsafe_allow_html=True)

# Initialize session state
if 'current_task' not in st.session_state:
    st.session_state.current_task = 'Task 1'
if 'evaluation_results' not in st.session_state:
    st.session_state.evaluation_results = None
if 'graph_analysis' not in st.session_state:
    st.session_state.graph_analysis = None

def create_evaluator_instance(task_type):
    """Creates and returns appropriate evaluator instance based on task type"""
    try:
        if task_type == "Task 1":
            return IELTSTask1Evaluator(), IELTSTask1ExerciseGenerator()
        else:
            return IELTSTask2Evaluator(), IELTSTask2ExerciseGenerator()
    except Exception as e:
        st.error(f"Error initializing evaluator: {str(e)}")
        return None, None

def display_task1_interface():
    """Displays Task 1 interface and handles evaluation"""
    st.header("Academic Writing Task 1 Evaluation")
    
    with st.container():
        col1, col2 = st.columns([1, 1])
        
        with col1:
            st.subheader("Upload Graph/Chart")
            uploaded_file = st.file_uploader(
                "Upload image (JPG, PNG)",
                type=['jpg', 'jpeg', 'png'],
                help="Upload the graph or chart you're describing"
            )
            
            # Add question input field
            question = st.text_area(
                "Enter the Task 1 question",
                height=100,
                help="Paste the task question here"
            )
            
            if uploaded_file and question:
                try:
                    image = Image.open(uploaded_file)
                    st.image(image, use_column_width=True)
                    # Save image temporarily
                    temp_path = "temp_graph.jpg"
                    image.save(temp_path)
                    
                    # Analyze graph but don't display analysis
                    evaluator, generator = create_evaluator_instance("Task 1")
                    if evaluator:
                        with st.spinner("Analyzing image..."):
                            graph_analysis = evaluator.analyze_graph(temp_path, question)
                            # Store graph analysis in session state without displaying
                            st.session_state.graph_analysis = graph_analysis
                            st.success("Image analysis complete")
                except Exception as e:
                    st.error(f"Error processing image: {str(e)}")
                finally:
                    if os.path.exists("temp_graph.jpg"):
                        os.remove("temp_graph.jpg")
        
        with col2:
            st.subheader("Your Response")
            written_response = st.text_area(
                "Enter your response here",
                height=300,
                help="Write your Task 1 response here"
            )
            
            if st.button("Evaluate Response", type="primary"):
                if not uploaded_file:
                    st.warning("Please upload a graph/chart first.")
                    return
                if not question:
                    st.warning("Please enter the task question.")
                    return
                if not written_response:
                    st.warning("Please enter your written response.")
                    return
                
                # Pass both question and graph analysis to evaluation
                if st.session_state.graph_analysis is not None:
                    process_evaluation(
                        written_response,
                        "Task 1",
                        {
                            'question': question,
                            'graph_analysis': st.session_state.graph_analysis
                        }
                    )
                else:
                    st.warning("Please wait for image analysis to complete.")
                    
def display_task2_interface():
    """Displays Task 2 interface and handles evaluation"""
    st.header("Academic Writing Task 2 Evaluation")
    
    with st.container():
        st.subheader("Essay Question")
        essay_question = st.text_area(
            "Enter the Task 2 question",
            height=100,
            help="Paste the essay question here"
        )
        
        st.subheader("Your Response")
        written_response = st.text_area(
            "Enter your response here",
            height=300,
            help="Write your Task 2 response here"
        )
        
        if st.button("Evaluate Essay", type="primary"):
            if not essay_question:
                st.warning("Please enter the essay question.")
                return
            if not written_response:
                st.warning("Please enter your written response.")
                return
                
            process_evaluation(written_response, "Task 2", essay_question)

def display_evaluation_results(task_type):
    """Displays evaluation results in organized tabs"""
    if not st.session_state.evaluation_results:
        return
        
    results = st.session_state.evaluation_results
    
    tabs = st.tabs([
        "📊 Overall Score",
        "🔤 Grammar Analysis",
        "📚 Vocabulary Analysis",
        "🎯 Task Achievement",
        "📝 Practice Exercises"
    ])
    
    with tabs[0]:
        st.markdown("### Overall Assessment")
        st.markdown(results['evaluation'])
        
    with tabs[1]:
        st.markdown("### Grammar Analysis")
        st.markdown(results['grammar'])
        
    with tabs[2]:
        st.markdown("### Vocabulary Analysis")
        st.markdown(results['vocabulary'])
        
    with tabs[3]:
        st.markdown("### Task Achievement")
        st.markdown(results['task_achievement'])
        
    with tabs[4]:
        if st.button("Generate Practice Exercises", type="secondary"):
            with st.spinner("Generating exercises..."):
                evaluator, generator = create_evaluator_instance(task_type)
                if generator:
                    grammar_exercises = generator.generate_grammar_exercises(
                        results['grammar'],
                        results['written_response']
                    )
                    vocab_exercises = generator.generate_vocabulary_exercises(
                        results['vocabulary']
                    )
                    
                    st.markdown("### Grammar Exercises")
                    st.markdown(grammar_exercises)
                    
                    st.markdown("### Vocabulary Exercises")
                    st.markdown(vocab_exercises)

def process_evaluation(written_response, task_type, additional_info=None):
    """Processes the evaluation request and stores results"""
    try:
        evaluator, _ = create_evaluator_instance(task_type)
        if not evaluator:
            return
            
        with st.spinner("Evaluating your response..."):
            if task_type == "Task 1":
                question = additional_info['question']
                graph_analysis = additional_info['graph_analysis']
                evaluation = evaluator.evaluate_task1_response(question, graph_analysis, written_response)
                grammar = evaluator.analyze_grammar(written_response)
                vocabulary = evaluator.analyze_vocabulary(written_response)
                task_achievement = evaluator.analyze_task_achievement(question, graph_analysis, written_response)
                cohesion = evaluator.analyze_cohesion_coherence(written_response, graph_analysis)
            else:
                evaluation = evaluator.evaluate_essay(additional_info, written_response)
                grammar = evaluator.analyze_grammar(written_response)
                vocabulary = evaluator.analyze_vocabulary(written_response)
                task_achievement = evaluator.analyze_task_achievement(additional_info, written_response)
                cohesion = evaluator.analyze_cohesion_coherence(written_response)
            
            st.session_state.evaluation_results = {
                'evaluation': evaluation,
                'grammar': grammar,
                'vocabulary': vocabulary,
                'task_achievement': task_achievement,
                'cohesion': cohesion,
                'written_response': written_response
            }
            
            st.success("Evaluation complete! View results in the tabs below.")
    except Exception as e:
        st.error(f"Error during evaluation: {str(e)}")

def main():
    """Main application function"""
    st.title("IELTS Writing Evaluation System")

    if not check_password():
        st.stop()  # Do not continue if check_password is not True
        
    # Task selection in sidebar
    with st.sidebar:
        st.title("Navigation")
        st.session_state.current_task = st.radio(
            "Select Task Type",
            ["Task 1", "Task 2"],
            key="task_selection"
        )
        
        st.markdown("---")
        st.markdown("### About")
        st.markdown("""
        This tool helps evaluate IELTS Writing responses and provides:
        - Detailed scoring
        - Grammar analysis
        - Vocabulary feedback
        - Task achievement assessment
        - Practice exercises
        """)
    
    # Display appropriate interface based on task selection
    if st.session_state.current_task == "Task 1":
        display_task1_interface()
    else:
        display_task2_interface()
    
    # Display evaluation results if available
    if st.session_state.evaluation_results:
        display_evaluation_results(st.session_state.current_task)

if __name__ == "__main__":
    main()