import streamlit as st
import openai
from openai import OpenAI
import time
import gspread
from oauth2client.service_account import ServiceAccountCredentials
import PyPDF2
import io
from datetime import datetime
from PIL import Image

# Add some custom CSS to improve the layout
st.markdown("""
    <style>
    .stImage {
        text-align: center;
        display: block;
        margin-left: auto;
        margin-right: auto;
    }
    
    .stTitle {
        text-align: center;
        padding-bottom: 20px;
    }
    
    div[data-testid="stVerticalBlock"] > div:has(div.stButton) {
        text-align: center;
        padding: 10px 0;
    }

    /* เพิ่ม CSS สำหรับ container ของ logo */
    .logo-container {
        display: flex;
        justify-content: center;
        align-items: center;
        padding: 1rem 0;
    }

    /* ปรับ column ให้อยู่ตรงกลาง */
    div[data-testid="column"] {
        display: flex;
        flex-direction: column;
        align-items: center;
    }
    </style>
    """, unsafe_allow_html=True)

def show_documentation():
    with st.expander("📚 How to Use GeneExam | วิธีการใช้งาน GeneExam"):
        st.markdown("""
        # Program Usage Guide | คู่มือการใช้งานโปรแกรม
        
        ## English Instructions
        1. **Input Your Content**
           - Use the text input box to paste your teaching material
           - Or upload a PDF file (maximum 10,000 words)
        
        2. **Select Question Type**
           - Multiple Choice: Generate questions with 3-5 options
           - Fill in the Blank: Create completion questions
           - True/False: Generate true/false statements
           - Open-ended: Create essay-type questions with scoring criteria
        
        3. **Choose Cognitive Level**
           - Select the appropriate cognitive level (see explanation below)
           - This determines the complexity and depth of questions
        
        4. **Additional Options**
           - Case-based: Toggle for medical case scenarios
           - Extra instructions: Add specific requirements
        
        5. **Generate and Download**
           - Click "Generate Questions" to create your exam
           - Use the download button to save your questions
        
        ## คำแนะนำภาษาไทย
        1. **การใส่เนื้อหา**
           - พิมพ์หรือวางเนื้อหาในช่องข้อความ
           - หรืออัพโหลดไฟล์ PDF (ไม่เกิน 10,000 คำ)
        
        2. **เลือกประเภทคำถาม**
           - ปรนัย: สร้างคำถามพร้อมตัวเลือก 3-5 ข้อ
           - เติมคำ: สร้างคำถามแบบเติมคำในช่องว่าง
           - ถูก/ผิด: สร้างคำถามแบบถูกผิด
           - อัตนัย: สร้างคำถามแบบบรรยายพร้อมเกณฑ์การให้คะแนน
        
        3. **เลือกระดับความคิด**
           - เลือกระดับความคิด (ดูคำอธิบายด้านล่าง)
           - กำหนดความซับซ้อนและความลึกของคำถาม
        
        4. **ตัวเลือกเพิ่มเติม**
           - คำถามเชิงกรณีศึกษา: เลือกสำหรับโจทย์ทางการแพทย์
           - คำแนะนำเพิ่มเติม: ใส่ความต้องการเฉพาะ
        
        5. **สร้างและดาวน์โหลด**
           - คลิก "Generate Questions" เพื่อสร้างข้อสอบ
           - ใช้ปุ่มดาวน์โหลดเพื่อบันทึกข้อสอบ
        """)

    with st.expander("🧠 Cognitive Levels Explanation | คำอธิบายระดับความคิด"):
        st.markdown("""
        # Cognitive Levels Guide | คู่มือระดับความคิด
        
        ## English Explanation

        ### 1. Recall (Knowledge)
        - **Definition**: Basic recall of information
        - **Key Words**: Define, List, Name, Identify
        - **Example**: What is the capital of Thailand?

        ### 2. Understanding (Comprehension)
        - **Definition**: Understanding and explaining ideas
        - **Key Words**: Explain, Describe, Summarize
        - **Example**: Explain how photosynthesis works.

        ### 3. Application
        - **Definition**: Using information in new situations
        - **Key Words**: Apply, Use, Solve, Demonstrate
        - **Example**: Calculate the dosage for a 70kg patient.

        ### 4. Analysis
        - **Definition**: Breaking information into parts
        - **Key Words**: Analyze, Compare, Contrast, Examine
        - **Example**: Compare and contrast viral and bacterial infections.

        ### 5. Synthesis
        - **Definition**: Creating new ideas or solutions
        - **Key Words**: Create, Design, Develop, Plan
        - **Example**: Design a treatment plan for a diabetic patient.

        ### 6. Evaluation
        - **Definition**: Making judgments based on criteria
        - **Key Words**: Evaluate, Judge, Assess, Recommend
        - **Example**: Evaluate the effectiveness of this treatment approach.

        ## คำอธิบายภาษาไทย

        ### 1. การจำ (Recall)
        - **ความหมาย**: การระลึกข้อมูลพื้นฐาน
        - **คำสำคัญ**: บอก, ระบุ, จำแนก, ระลึก
        - **ตัวอย่าง**: เมืองหลวงของประเทศไทยคืออะไร?

        ### 2. ความเข้าใจ (Understanding)
        - **ความหมาย**: เข้าใจและอธิบายแนวคิด
        - **คำสำคัญ**: อธิบาย, บรรยาย, สรุป
        - **ตัวอย่าง**: อธิบายกระบวนการสังเคราะห์แสง

        ### 3. การประยุกต์ใช้ (Application)
        - **ความหมาย**: ใช้ข้อมูลในสถานการณ์ใหม่
        - **คำสำคัญ**: ใช้, แก้ปัญหา, สาธิต
        - **ตัวอย่าง**: คำนวณขนาดยาสำหรับผู้ป่วยน้ำหนัก 70 กิโลกรัม

        ### 4. การวิเคราะห์ (Analysis)
        - **ความหมาย**: แยกแยะข้อมูลเป็นส่วนๆ
        - **คำสำคัญ**: วิเคราะห์, เปรียบเทียบ, จำแนก
        - **ตัวอย่าง**: เปรียบเทียบการติดเชื้อไวรัสและแบคทีเรีย

        ### 5. การสังเคราะห์ (Synthesis)
        - **ความหมาย**: สร้างแนวคิดหรือวิธีการใหม่
        - **คำสำคัญ**: สร้าง, ออกแบบ, พัฒนา, วางแผน
        - **ตัวอย่าง**: ออกแบบแผนการรักษาสำหรับผู้ป่วยเบาหวาน

        ### 6. การประเมินผล (Evaluation)
        - **ความหมาย**: ตัดสินใจบนพื้นฐานของเกณฑ์
        - **คำสำคัญ**: ประเมิน, ตัดสิน, วัดผล, แนะนำ
        - **ตัวอย่าง**: ประเมินประสิทธิภาพของวิธีการรักษานี้
        """)


# Constants
WORD_LIMIT = 11000
DAILY_API_LIMIT = 30  # Set your desired limit per user per day

# Set up OpenAI client
client = OpenAI(api_key=st.secrets["OPENAI_API_KEY"])

# Google Sheets setup
scope = ["https://spreadsheets.google.com/feeds", "https://www.googleapis.com/auth/drive"]
creds = ServiceAccountCredentials.from_json_keyfile_name("genexam-2c8c645ecc0d.json", scope)
client_gs = gspread.authorize(creds)
sheet = client_gs.open("GeneXam user").sheet1

def check_user_in_sheet(username):
    """Check if user exists in sheet"""
    try:
        users_list = sheet.col_values(1)  # UserID column
        if username in users_list:
            return True
        return False
    except Exception as e:
        st.error(f"Error checking user: {str(e)}")
        return False

def get_user_stats(username):
    """Get user's current API usage statistics"""
    try:
        users_list = sheet.col_values(1)  # UserID column
        row_number = users_list.index(username) + 1
        
        daily_count = int(sheet.cell(row_number, 2).value)    # DailyAPICount
        total_count = int(sheet.cell(row_number, 3).value)    # TotalAPICount
        last_used = sheet.cell(row_number, 4).value           # LastUsedDate
        
        return {
            'daily_count': daily_count,
            'total_count': total_count,
            'last_used': last_used
        }
    except Exception as e:
        st.error(f"Error getting user stats: {str(e)}")
        return None

def update_api_usage(username):
    """Update both daily and total API usage counts"""
    try:
        users_list = sheet.col_values(1)
        row_number = users_list.index(username) + 1
        today = datetime.now().strftime('%Y-%m-%d')
        
        # Get current values
        stats = get_user_stats(username)
        if not stats:
            return False, "Error retrieving user statistics"
        
        # Reset daily count if it's a new day
        daily_count = stats['daily_count']
        if stats['last_used'] != today:
            daily_count = 0
        
        # Check daily limit
        if daily_count >= DAILY_API_LIMIT:
            return False, f"You have reached your daily limit of {DAILY_API_LIMIT} generations. Please try again tomorrow."
        
        # Update counts
        new_daily_count = daily_count + 1
        new_total_count = stats['total_count'] + 1
        
        # Update all values in sheet
        sheet.update_cell(row_number, 2, new_daily_count)  # Update DailyAPICount
        sheet.update_cell(row_number, 3, new_total_count)  # Update TotalAPICount
        sheet.update_cell(row_number, 4, today)            # Update LastUsedDate
        
        return True, None
        
    except Exception as e:
        return False, f"Error updating API usage: {str(e)}"

def extract_text_from_pdf(pdf_file):
    """Simple PDF text extraction with word limit check"""
    try:
        pdf_reader = PyPDF2.PdfReader(io.BytesIO(pdf_file.read()))
        text_content = ""
        for page in pdf_reader.pages:
            text_content += page.extract_text() + "\n"
        
        word_count = len(text_content.split())
        if word_count > WORD_LIMIT:
            return None, f"PDF content exceeds {WORD_LIMIT:,} words (contains {word_count:,} words). Please use a shorter document."
        return text_content, None
    except Exception as e:
        return None, f"Error processing PDF: {str(e)}"

def generate_questions_with_retry(username, knowledge_material, question_type, cognitive_level, extra_instructions, case_based, num_choices=None, max_retries=3):
    """Generate questions and update API usage"""
    
    # Check and update API usage before generating
    can_generate, error_message = update_api_usage(username)
    if not can_generate:
        st.error(error_message)
        return None

    # Adjust number of questions based on type
    if question_type == "Multiple Choice":
        num_questions = 3
        format_instructions = f"""
        For each multiple choice question:
        1. Present the question clearly, ending with '?'
        2. Leave one blank line after the question
        3. Present choices as:
           A) [choice]
           B) [choice]
           C) [choice]
           {f"D) [choice]" if num_choices > 3 else ""}
           {f"E) [choice]" if num_choices > 4 else ""}
        4. After all questions, provide an ANSWER KEY section with:
           - The correct answer letter for each question
           - A brief explanation of why this is the correct answer
        
        Example format:
        1. Your question text here?
    
        A) First choice
        B) Second choice
        C) Third choice
        """
    
    elif question_type == "Fill in the Blank":
        num_questions = 10
        format_instructions = """
        For each fill-in-the-blank question:
        1. Present the question with a clear blank space indicated by _____
        2. After all questions, provide an ANSWER KEY section with:
           - The correct answer for each blank
           - A brief explanation of why this answer is correct
           - Any alternative acceptable answers if applicable
        """
    elif question_type == "True/False":
        num_questions = 5
        format_instructions = """
        For each true/false question:
        1. Present the statement clearly
        2. After all questions, provide an ANSWER KEY section with:
           - Whether the statement is True or False
           - A detailed explanation of why the statement is true or false
           - The specific part of the source material that supports this answer
        """
    else:  # Open-ended
        num_questions = 3
        format_instructions = """
        For each open-ended question:
        1. Present the question clearly
        2. After all questions, provide an ANSWER KEY section with:
           - A structured scoring checklist of key points (minimum 3-5 points per question)
           - Each key point should be worth a specific number of marks
           - Total marks available for each question
           - Sample answer that would receive full marks
           - Common points that students might miss
        """

    # Base prompt
    prompt = f"""Generate {num_questions} {question_type.lower()} exam questions based on {cognitive_level.lower()} level from the following material:

{knowledge_material}

{format_instructions}

{extra_instructions}

Please format the output clearly with:
1. Questions section (numbered 1, 2, 3, etc.)
2. Answer Key section (clearly separated from questions)
3. Each answer should include explanation for better understanding

Make sure all questions and answers are directly related to the provided material."""

    # Modify prompt for case-based medical situations
    if case_based:
        prompt = f"""Generate {num_questions} {question_type.lower()} case-based medical exam questions based on {cognitive_level.lower()} level.

Use this material as the medical knowledge base:
{knowledge_material}

Each question should:
1. Start with a medical case scenario/patient presentation
2. Include relevant clinical details
3. Ask about diagnosis, treatment, or management
4. Be at {cognitive_level.lower()} cognitive level

{format_instructions}

{extra_instructions}

Please format the output with:
1. Cases and Questions (numbered 1, 2, 3, etc.)
2. Detailed Answer Key section including:
   - Correct answers
   - Clinical reasoning
   - Key diagnostic or treatment considerations
   - Common pitfalls to avoid"""

    retries = 0
    while retries < max_retries:
        try:
            response = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": "You are an expert exam question generator with deep knowledge in medical education. Create clear, well-structured questions with detailed answer keys and explanations."},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.7,
                max_tokens=3000  # Increased to accommodate answers and explanations
            )
            return response.choices[0].message.content
        except Exception as e:
            retries += 1
            st.warning(f"Attempt {retries} failed. Retrying... Error: {str(e)}")
            if retries == max_retries:
                st.error(f"Failed to generate questions after {max_retries} attempts. Error: {str(e)}")
                return None
            time.sleep(2)

# Main Streamlit interface
# Initialize session state variables
if 'login_step' not in st.session_state:
    st.session_state.login_step = 'username'
if 'username' not in st.session_state:
    st.session_state.username = None

# Login system
if st.session_state.username is None:
    # Center align the content
    col1, col2, col3 = st.columns([1,2,1])
    
    with col2:
        # Display logo
        st.markdown('<div class="logo-container">', unsafe_allow_html=True)
        st.image("GenExam.png", width=200)
        st.markdown('</div>', unsafe_allow_html=True)
        
        st.title("Login")
        username_input = st.text_input("Enter your username:")

        if st.session_state.login_step == 'username' and st.button("Login", use_container_width=True):
            if username_input:
                if check_user_in_sheet(username_input):
                    stats = get_user_stats(username_input)
                    if stats:
                        st.success(f"Welcome, {username_input}! 👋")
                        st.info(f"""
                        📊 Your API Usage Statistics:
                        - Today's Usage: {stats['daily_count']}/{DAILY_API_LIMIT} generations
                        - Total All-Time Usage: {stats['total_count']} generations
                        """)
                        st.session_state.login_step = 'enter_app'
                else:
                    st.warning("Username not found. Please try again.")
            else:
                st.warning("Please enter a valid username.")

        if st.session_state.login_step == 'enter_app':
            if st.button("🎯 Enter GeneXam Application", use_container_width=True):
                st.session_state.username = username_input
                st.rerun()

        # Show instructions
        if st.session_state.login_step == 'username':
            st.markdown("""
            ### How to Login:
            1. Enter your username and click 'Login' to verify your account
            2. After verification, click 'Enter GeneXam Application' to start using the system
            """)

else:
    # Main application code (ส่วนที่เหลือเหมือนเดิม)
    st.title(f"Welcome to GeneXam, {st.session_state.username}! 🎓")

    # Add Help button in sidebar  # <-- เพิ่มส่วนนี้ตรงนี้
    with st.sidebar:
        st.markdown("### Need Help? | ต้องการความช่วยเหลือ?")
        if st.button("📖 Show Documentation | แสดงคู่มือการใช้งาน"):
            show_documentation()
    
    # Show current usage stats
    stats = get_user_stats(st.session_state.username)
    if stats:
        remaining = DAILY_API_LIMIT - stats['daily_count']
        st.info(f"""
        📊 Usage Statistics:
        - Daily Generations Remaining: {remaining}/{DAILY_API_LIMIT}
        - Total All-Time Generations: {stats['total_count']}
        """)
    
    # Create tabs for input methods
    tab1, tab2 = st.tabs(["Text Input", "PDF Upload"])
    
    with tab1:
        knowledge_material = st.text_area("Enter knowledge material to generate exam questions:")
        word_count = len(knowledge_material.split())
        if word_count > WORD_LIMIT:
            st.error(f"Text exceeds {WORD_LIMIT:,} words. Please shorten your content.")
    
    with tab2:
        st.info(f"Maximum content length: {WORD_LIMIT:,} words")
        uploaded_file = st.file_uploader("Upload a PDF file", type="pdf")
        
        if uploaded_file is not None:
            pdf_content, error = extract_text_from_pdf(uploaded_file)
            if error:
                st.error(error)
            else:
                st.success("PDF processed successfully!")
                knowledge_material = pdf_content

    # Question generation options
    col1, col2 = st.columns(2)
    
    with col1:
        question_type = st.selectbox(
            "Select question type:",
            ["Multiple Choice", "Fill in the Blank", "Open-ended", "True/False"]
        )
        
        if question_type == "Multiple Choice":
            num_choices = st.selectbox("Select number of choices:", [3, 4, 5])
        
        cognitive_level = st.selectbox(
            "Select cognitive level:",
            ["Recall", "Understanding", "Application", "Analysis", "Synthesis", "Evaluation"]
        )
    
    with col2:
        case_based = st.checkbox("Generate case-based medical exam questions")
        extra_instructions = st.text_area("Additional instructions (optional):")

    # Generate questions button
    if st.button("Generate Questions"):
        if 'knowledge_material' in locals() and knowledge_material.strip():
            with st.spinner("Generating questions..."):
                questions = generate_questions_with_retry(
                    st.session_state['username'],
                    knowledge_material,
                    question_type,
                    cognitive_level,
                    extra_instructions,
                    case_based,
                    num_choices if question_type == "Multiple Choice" else None
                )
                
                if questions:
                    st.write("### Generated Exam Questions:")
                    st.write(questions)
                    
                    # Update displayed stats after generation
                    new_stats = get_user_stats(st.session_state['username'])
                    if new_stats:
                        remaining = DAILY_API_LIMIT - new_stats['daily_count']
                        st.info(f"""
                        📊 Updated Usage Statistics:
                        - Daily Generations Remaining: {remaining}/{DAILY_API_LIMIT}
                        - Total All-Time Generations: {new_stats['total_count']}
                        """)
                    
                    # Download button
                    st.download_button(
                        label="Download Questions",
                        data=questions,
                        file_name='generated_questions.txt',
                        mime='text/plain'
                    )
        else:
            st.warning("Please enter knowledge material or upload a PDF file first.")